diff --git a/.github/workflows/prettier.yml b/.github/workflows/prettier.yml new file mode 100644 index 0000000..6d576e3 --- /dev/null +++ b/.github/workflows/prettier.yml @@ -0,0 +1,19 @@ +name: Prettier + +on: + push: + branches: [main] + pull_request: + +permissions: + contents: read + +jobs: + prettier: + name: Check Markdown formatting + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + - uses: oven-sh/setup-bun@0c5077e51419868618aeaa5fe8019c62421857d6 # v2.2.0 + - name: Run Prettier + run: bunx prettier@3.8.3 --check "**/*.md" diff --git a/.prettierrc b/.prettierrc new file mode 100644 index 0000000..64182b1 --- /dev/null +++ b/.prettierrc @@ -0,0 +1,3 @@ +{ + "proseWrap": "preserve" +} diff --git a/CODE_OF_CONDUCT.md b/CODE_OF_CONDUCT.md index 5b627cf..ec98f2b 100644 --- a/CODE_OF_CONDUCT.md +++ b/CODE_OF_CONDUCT.md @@ -1,4 +1,5 @@ ## Code of Conduct + This project has adopted the [Amazon Open Source Code of Conduct](https://aws.github.io/code-of-conduct). For more information see the [Code of Conduct FAQ](https://aws.github.io/code-of-conduct-faq) or contact opensource-codeofconduct@amazon.com with any additional questions or comments. diff --git a/GETTING_STARTED_GUIDE.md b/GETTING_STARTED_GUIDE.md index 683d6fe..ec676c1 100644 --- a/GETTING_STARTED_GUIDE.md +++ b/GETTING_STARTED_GUIDE.md @@ -10,13 +10,13 @@ Common use cases include fused attention kernels, custom normalization operation ## Prerequisites -| # | Requirement | Details | Needed for | -|---|-------------|---------|------------| -| 1 | Trainium/Inferentia instance | trn1, trn2, inf2 EC2 instance (AL2023 DLAMI recommended) | Compiling, profiling, optimization | -| 2 | Neuron SDK | `aws-neuronx-tools` (pre-installed on DLAMI) | All on-device skills | -| 3 | Python venv with Neuron packages | `neuronx-cc`, `torch-neuronx`, `neuron-explorer` | Compilation, profiling, analysis | -| 4 | Kiro or Claude Code | Installed on the Trainium instance | Running agents and skills | -| 5 | Anthropic API key | For Claude model inference | Agent reasoning | +| # | Requirement | Details | Needed for | +| --- | -------------------------------- | -------------------------------------------------------- | ---------------------------------- | +| 1 | Trainium/Inferentia instance | trn1, trn2, inf2 EC2 instance (AL2023 DLAMI recommended) | Compiling, profiling, optimization | +| 2 | Neuron SDK | `aws-neuronx-tools` (pre-installed on DLAMI) | All on-device skills | +| 3 | Python venv with Neuron packages | `neuronx-cc`, `torch-neuronx`, `neuron-explorer` | Compilation, profiling, analysis | +| 4 | Kiro or Claude Code | Installed on the Trainium instance | Running agents and skills | +| 5 | Anthropic API key | For Claude model inference | Agent reasoning | > **Important:** The agent runs on the same machine as the hardware. There is no laptop-to-remote-box file transfer — everything is co-located. Writing and documentation skills work anywhere (no hardware needed), but profiling and debugging require on-instance execution. @@ -44,8 +44,7 @@ source ~/opt/aws_neuronx_venv_pytorch_2_9/bin/activate ## Step 3: Install Neuron Agentic Development - -### Clone from GitHub (for customization or contribution):** +### Clone from GitHub (for customization or contribution):\*\* ```bash git clone https://github.com/aws-neuron/neuron-agentic-development.git @@ -79,24 +78,24 @@ The `neuron-nki-agent` is the unified entry point. It automatically selects the ### Example Prompts -| What you want to do | What to say | Hardware needed? | -|---------------------|-------------|------------------| -| Write a new kernel | "Write a fused softmax kernel for bf16 inputs" | No | -| Debug a compilation error | "Fix this kernel" (with error output) | Yes | -| Profile a kernel | "Profile my kernel and show me the metrics" | Yes | -| Analyze a profile | "What's the bottleneck in this kernel?" | Yes (neuron-explorer) | +| What you want to do | What to say | Hardware needed? | +| ------------------------- | ---------------------------------------------- | --------------------- | +| Write a new kernel | "Write a fused softmax kernel for bf16 inputs" | No | +| Debug a compilation error | "Fix this kernel" (with error output) | Yes | +| Profile a kernel | "Profile my kernel and show me the metrics" | Yes | +| Analyze a profile | "What's the bottleneck in this kernel?" | Yes (neuron-explorer) | ## Skills The package provides five specialized skills that follow the natural kernel development pipeline: **write → debug → profile → analyze**. -| # | Skill | Category | Use when | -|---|-------|----------|----------| -| 1 | `neuron-nki-writing` | Authoring | Writing new kernels or modifying existing ones | -| 2 | `neuron-nki-debugging` | Debugging | Resolving compiler errors or numerical mismatches | -| 3 | `neuron-nki-docs` | Documentation | Looking up API signatures, tutorials, error codes | -| 4 | `neuron-nki-profiling` | Profiling | Capturing execution traces on hardware | -| 5 | `neuron-nki-profile-querying` | Analysis | SQL-based performance bounds and bottleneck analysis | +| # | Skill | Category | Use when | +| --- | ----------------------------- | ------------- | ---------------------------------------------------- | +| 1 | `neuron-nki-writing` | Authoring | Writing new kernels or modifying existing ones | +| 2 | `neuron-nki-debugging` | Debugging | Resolving compiler errors or numerical mismatches | +| 3 | `neuron-nki-docs` | Documentation | Looking up API signatures, tutorials, error codes | +| 4 | `neuron-nki-profiling` | Profiling | Capturing execution traces on hardware | +| 5 | `neuron-nki-profile-querying` | Analysis | SQL-based performance bounds and bottleneck analysis | ### Kernel Authoring (`neuron-nki-writing`) @@ -120,12 +119,12 @@ Used across all stages of development. Provides API signatures and tutorials dur ## Agents -| # | Agent | Focus | What it does | -|---|-------|-------|--------------| -| 1 | `neuron-nki-agent` | Full lifecycle | Top-level entry point. Auto-selects the right workflow based on your request and orchestrates the appropriate skills. | -| 2 | `neuron-nki-writer-agent` | Authoring | Translates PyTorch, NumPy, or natural language descriptions into NKI code. Handles modifications to existing kernels. | -| 3 | `neuron-nki-debugger-agent` | Debugging | Autonomously analyzes compiler errors, searches documentation for fixes, and applies corrections. Tracks iterations (up to 10) and progressively simplifies when stuck. | -| 4 | `neuron-nki-profile-analysis-agent` | Profiling + Analysis | Captures execution profiles on hardware, then runs SQL queries against profile data to compute performance bounds, identify bottleneck engines, and localize inefficiencies. | +| # | Agent | Focus | What it does | +| --- | ----------------------------------- | -------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| 1 | `neuron-nki-agent` | Full lifecycle | Top-level entry point. Auto-selects the right workflow based on your request and orchestrates the appropriate skills. | +| 2 | `neuron-nki-writer-agent` | Authoring | Translates PyTorch, NumPy, or natural language descriptions into NKI code. Handles modifications to existing kernels. | +| 3 | `neuron-nki-debugger-agent` | Debugging | Autonomously analyzes compiler errors, searches documentation for fixes, and applies corrections. Tracks iterations (up to 10) and progressively simplifies when stuck. | +| 4 | `neuron-nki-profile-analysis-agent` | Profiling + Analysis | Captures execution profiles on hardware, then runs SQL queries against profile data to compute performance bounds, identify bottleneck engines, and localize inefficiencies. | ## Architecture @@ -161,7 +160,7 @@ Used across all stages of development. Provides API signatures and tutorials dur Here's a typical end-to-end workflow: -1. **Write** — Ask the agent: *"Write an NKI kernel that computes scaled softmax: softmax(x * scale) along the last dimension, for input shape [batch, seq_len, hidden_dim] in bfloat16."* The agent produces a complete kernel with proper tiling, hardware-accelerated exp, float32 accumulation, and bfloat16 output cast. +1. **Write** — Ask the agent: _"Write an NKI kernel that computes scaled softmax: softmax(x _ scale) along the last dimension, for input shape [batch, seq_len, hidden_dim] in bfloat16."\* The agent produces a complete kernel with proper tiling, hardware-accelerated exp, float32 accumulation, and bfloat16 output cast. 2. **Debug** — Ask the agent to run the kernel and verify numerical parity against a PyTorch reference. If compilation errors occur, the agent diagnoses and fixes them autonomously. diff --git a/README.md b/README.md index 81854ed..bba3782 100644 --- a/README.md +++ b/README.md @@ -1,4 +1,5 @@ # Neuron Agentic Development + This repository contains AI agents and skills for developing on [AWS Neuron](https://awsdocs-neuron.readthedocs-hosted.com/) (Trainium/Inferentia) hardware, including NKI kernel development, profiling and debugging. For an overview of Neuron Agentic Development and the tools it offers for agent-enabled workflows with Neuron, see [the overview of Neuron Agentic Development in the the public Neuron docs](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/about-neuron/agentic-development-overview.html). ## Installation @@ -27,30 +28,29 @@ deploy-neuron-agentic-development-to-claude ## Agents -| Agent | Description | -|-------|-------------| -| [neuron-nki-agent](agents/neuron-nki-agent.md) | Unified NKI kernel development agent. Full lifecycle: writing kernels from PyTorch/NumPy/natural language, debugging compilation errors, profiling performance, optimizing bottlenecks, migrating between API versions, analyzing Perfetto traces, and NKI documentation lookup. | -| [neuron-nki-writer-agent](agents/neuron-nki-writer-agent.md) | NKI kernel authoring and modification. Translates from PyTorch/NumPy/natural language, adds shape/dtype support, refactors tiling strategies, and implements new features following Beta 3 API patterns. | -| [neuron-nki-debugger-agent](agents/neuron-nki-debugger-agent.md) | Autonomous NKI kernel compilation error debugging. Analyzes compiler errors, searches documentation and code examples for fixes, applies corrections following simplicity over performance, and validates fixes. | -| [neuron-nki-profile-analysis-agent](agents/neuron-nki-profile-analysis-agent.md) | Profile and analyze NKI kernels on Neuron hardware. Captures execution traces, computes performance bounds, identifies bottleneck engines, and runs investigations to localize inefficiencies to NKI source lines. | -| [neuron-framework-autoport-agent](agents/neuron-framework-autoport-agent.md) | A model porting agent to port GPU-compatible models to functionally accurate implementation on Neuron. Executes the full porting workflow including architecture analysis, implementation, compilation, inference testing, and validation. | +| Agent | Description | +| -------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| [neuron-nki-agent](agents/neuron-nki-agent.md) | Unified NKI kernel development agent. Full lifecycle: writing kernels from PyTorch/NumPy/natural language, debugging compilation errors, profiling performance, optimizing bottlenecks, migrating between API versions, analyzing Perfetto traces, and NKI documentation lookup. | +| [neuron-nki-writer-agent](agents/neuron-nki-writer-agent.md) | NKI kernel authoring and modification. Translates from PyTorch/NumPy/natural language, adds shape/dtype support, refactors tiling strategies, and implements new features following Beta 3 API patterns. | +| [neuron-nki-debugger-agent](agents/neuron-nki-debugger-agent.md) | Autonomous NKI kernel compilation error debugging. Analyzes compiler errors, searches documentation and code examples for fixes, applies corrections following simplicity over performance, and validates fixes. | +| [neuron-nki-profile-analysis-agent](agents/neuron-nki-profile-analysis-agent.md) | Profile and analyze NKI kernels on Neuron hardware. Captures execution traces, computes performance bounds, identifies bottleneck engines, and runs investigations to localize inefficiencies to NKI source lines. | +| [neuron-framework-autoport-agent](agents/neuron-framework-autoport-agent.md) | A model porting agent to port GPU-compatible models to functionally accurate implementation on Neuron. Executes the full porting workflow including architecture analysis, implementation, compilation, inference testing, and validation. | ## Skills -| Skill | Description | -|-------|-------------| -| [neuron-nki-writing](skills/neuron-nki-writing/SKILL.md) | Write and modify NKI kernels. Covers new kernel creation from PyTorch/NumPy/natural language, editing existing kernels, adding shape/dtype support, refactoring tiling strategies, and implementing new features. | -| [neuron-nki-debugging](skills/neuron-nki-debugging/SKILL.md) | Debug NKI compilation errors on Neuron hardware. | -| [neuron-nki-docs](skills/neuron-nki-docs/SKILL.md) | Research NKI documentation for API lookups, tutorials, error codes, and architecture details. | -| [neuron-nki-profiling](skills/neuron-nki-profiling/SKILL.md) | Profile NKI kernels to analyze performance on Neuron hardware. | -| [neuron-nki-profile-querying](skills/neuron-nki-profile-querying/SKILL.md) | Query and analyze NKI kernel profile data from neuron-explorer parquet files via SQL and Python. | -| [neuron-framework-autoport](skills/neuron-framework-autoport/SKILL.md) | Port a GPU compatible model to NeuronX Distributed Inference for AWS Trainium/Inferentia. Handles the full workflow including architecture analysis, NeuronX implementation, compilation, inference testing, and accuracy validation. | -| [neuron-framework-equivalence](skills/neuron-framework-equivalence/SKILL.md) | Verifies functional equivalence between two implementations of the same model using a hierarchical algorithm. | +| Skill | Description | +| ---------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| [neuron-nki-writing](skills/neuron-nki-writing/SKILL.md) | Write and modify NKI kernels. Covers new kernel creation from PyTorch/NumPy/natural language, editing existing kernels, adding shape/dtype support, refactoring tiling strategies, and implementing new features. | +| [neuron-nki-debugging](skills/neuron-nki-debugging/SKILL.md) | Debug NKI compilation errors on Neuron hardware. | +| [neuron-nki-docs](skills/neuron-nki-docs/SKILL.md) | Research NKI documentation for API lookups, tutorials, error codes, and architecture details. | +| [neuron-nki-profiling](skills/neuron-nki-profiling/SKILL.md) | Profile NKI kernels to analyze performance on Neuron hardware. | +| [neuron-nki-profile-querying](skills/neuron-nki-profile-querying/SKILL.md) | Query and analyze NKI kernel profile data from neuron-explorer parquet files via SQL and Python. | +| [neuron-framework-autoport](skills/neuron-framework-autoport/SKILL.md) | Port a GPU compatible model to NeuronX Distributed Inference for AWS Trainium/Inferentia. Handles the full workflow including architecture analysis, NeuronX implementation, compilation, inference testing, and accuracy validation. | +| [neuron-framework-equivalence](skills/neuron-framework-equivalence/SKILL.md) | Verifies functional equivalence between two implementations of the same model using a hierarchical algorithm. | ## Contributing -We are evaluating the external contribution process. All capabilities undergo internal verification to ensure technical accuracy, security, and architectural alignment. In the interim, we welcome feedback and feature requests via [Issues](https://github.com/aws-neuron/neuron-agentic-development/issues) - +We are evaluating the external contribution process. All capabilities undergo internal verification to ensure technical accuracy, security, and architectural alignment. In the interim, we welcome feedback and feature requests via [Issues](https://github.com/aws-neuron/neuron-agentic-development/issues) ## License diff --git a/agents/neuron-framework-autoport-agent.md b/agents/neuron-framework-autoport-agent.md index 71df053..545b669 100644 --- a/agents/neuron-framework-autoport-agent.md +++ b/agents/neuron-framework-autoport-agent.md @@ -20,7 +20,19 @@ description: | model: opus color: blue -tools: ["Read", "Write", "Edit", "Grep", "Glob", "Bash", "Task", "TodoWrite", "Skill","Agent"] +tools: + [ + "Read", + "Write", + "Edit", + "Grep", + "Glob", + "Bash", + "Task", + "TodoWrite", + "Skill", + "Agent", + ] skills: - neuron-framework-autoport --- @@ -31,19 +43,22 @@ You are an autonomous model porting agent. You accept HuggingFace model paramete ## Workflow Routing -| Request Type | Skill | -|---|---| +| Request Type | Skill | +| ----------------------------------- | ---------------------------- | | Port a HuggingFace model to NeuronX | `/neuron-framework-autoport` | ## Prerequisites Before starting any porting workflow, verify NeuronCores are available: + ```bash neuron-ls ``` + If 0 cores are detected and the user did not specify dry-run mode, tell the user to allocate a compute node with Neuron hardware and STOP. Also clear any stale compile cache: + ```bash rm -rf /var/tmp/neuron-compile-cache ``` @@ -58,27 +73,33 @@ rm -rf /var/tmp/neuron-compile-cache ## Project Guidelines ### Prohibited Packages + - Do not import, reference, or run any code from `transformers_neuronx`. It is an old API library. ### PYTHONPATH Handling + - If you run into issues with imports and PYTHONPATH, do not make changes to the script — change PYTHONPATH instead. When you test, do the same. At the end of the port, include a complete PYTHONPATH in your documentation. ### Error Handling + - Do not generate any `try/except` statements. - Let errors surface directly without catching them. - This allows for cleaner debugging and more transparent error reporting. ### File Organization + - `agent_artifacts/tmp/` — All temporary files (compile scripts, test scripts, intermediate artifacts) - `neuron_port/` — All ported model files (modeling and configuration files) - `agent_artifacts/traces/` — Checkpoint prompts, completions, and tool use for every major step - `agent_artifacts/data/` — All weights, checkpoints, and downloaded artifacts. Do not store weights anywhere else. ### Hardware Context + - You are typically running on a trn1.32xlarge with 32 cores and 16GB per core. - If you question the hardware, use `neuron-ls` to validate. Never assume you are not running on trn1.32xlarge with 32 cores. ### Debugging Tips + - If you get a JSON parse error (`[NLA001]`) or `FileNotFoundError` on neff_output paths, delete `/var/tmp/neuron-compile-cache` and retry. - Compiler logs are in `agent_artifacts/data/neff_output/context_encoding_model/` — look for `log-neuron-cc.txt`. Use bash to read them. - Ignore this warning, it is not important: `WARNING:Neuron:TP degree (XX) and KV heads (YY) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!` diff --git a/agents/neuron-framework-autoport-vllm-agent.md b/agents/neuron-framework-autoport-vllm-agent.md index 3a69ad4..16091b0 100644 --- a/agents/neuron-framework-autoport-vllm-agent.md +++ b/agents/neuron-framework-autoport-vllm-agent.md @@ -19,7 +19,18 @@ description: | model: opus color: blue -tools: ["Read", "Write", "Edit", "Grep", "Glob", "Bash", "Task", "TodoWrite", "Skill"] +tools: + [ + "Read", + "Write", + "Edit", + "Grep", + "Glob", + "Bash", + "Task", + "TodoWrite", + "Skill", + ] skills: - neuron-framework-autoport-vllm-neuron - neuron-framework-equivalence @@ -33,10 +44,10 @@ IMPORTANT: Do NOT validate or check if tools are available. Just use them direct ## Workflow Routing -| Request Type | Skill | -|---|---| -| Port a HuggingFace model to vLLM-Neuron | `/neuron-framework-autoport-vllm-neuron` | -| Deep equivalence validation of a completed port | `/neuron-framework-equivalence` | +| Request Type | Skill | +| ----------------------------------------------- | ---------------------------------------- | +| Port a HuggingFace model to vLLM-Neuron | `/neuron-framework-autoport-vllm-neuron` | +| Deep equivalence validation of a completed port | `/neuron-framework-equivalence` | ## Final Validation: Equivalence @@ -47,16 +58,21 @@ Every port ends with deep equivalence validation (Step 11 of the autoport skill) Before starting any porting workflow, verify the environment: 1. Check for virtual environment: + ```bash echo $NXDI_VENV_PATH ``` + If set, activate it before running any Python commands: + ```bash source $NXDI_VENV_PATH/bin/activate ``` + If not set, check for a local config at `.kiro/local.md` or `.claude/local.md` with `nxdi_venv_path` in YAML frontmatter. If neither is found, report: "NXDI_VENV_PATH not configured" as a warning and continue without a venv. 2. Verify required packages. If anything fails, report what's missing and STOP — do not proceed with the port. + ```python import sys @@ -69,9 +85,11 @@ print("\nPackage check complete.") ``` 3. Verify NeuronCores are available: + ```bash neuron-ls ``` + If 0 cores are detected, tell the user to allocate a compute node with Neuron hardware and STOP. > **Note:** Do NOT clear `/var/tmp/neuron-compile-cache` as a pre-flight step — it is a shared system directory and other processes or users may depend on it. Only clear it reactively if you hit a `[NLA001]` JSON parse error or `FileNotFoundError` on neff_output paths (see Debugging Tips below). @@ -88,17 +106,21 @@ If 0 cores are detected, tell the user to allocate a compute node with Neuron ha ## Project Guidelines ### Prohibited Packages + - Do not import, reference, or run any code from `transformers_neuronx`. It is an old API library. ### PYTHONPATH Handling + - If you run into issues with imports and PYTHONPATH, do not make changes to the script — change PYTHONPATH instead. When you test, do the same. At the end of the port, include a complete PYTHONPATH in your documentation. ### Error Handling + - Do not generate any `try/except` statements - Let errors surface directly without catching them - This allows for cleaner debugging and more transparent error reporting ### File Organization + - Model code: `vllm_neuron/model/MODEL_NAME/` - Examples: `examples/MODEL_NAME/` - Registry: `vllm_neuron/model/registry.py` @@ -106,10 +128,12 @@ If 0 cores are detected, tell the user to allocate a compute node with Neuron ha - `agent_artifacts/traces/` — Checkpoint prompts, completions, and tool use for every major step ### Hardware Context + - You are typically running on a trn2 instance. Use `neuron-ls` to verify available NeuronCores. - Set `NEURON_SKIP_EFA_AFFINITY=1` for trn2 instances where PCI topology doesn't match hardcoded BDF-to-EFA mapping. ### Debugging Tips + - If you get a JSON parse error (`[NLA001]`) or `FileNotFoundError` on neff_output paths, delete `/var/tmp/neuron-compile-cache` and retry. - Compiler logs are in `agent_artifacts/data/neff_output/context_encoding_model/` — look for `log-neuron-cc.txt`. Use bash to read them. - Ignore this warning, it is not important: `WARNING:Neuron:TP degree (XX) and KV heads (YY) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!` diff --git a/agents/neuron-framework-equivalence-agent.md b/agents/neuron-framework-equivalence-agent.md index f385cc1..049ffff 100644 --- a/agents/neuron-framework-equivalence-agent.md +++ b/agents/neuron-framework-equivalence-agent.md @@ -27,7 +27,19 @@ description: | model: opus -tools: ["Read", "Write", "Edit", "Grep", "Glob", "Bash", "Task", "TodoWrite", "Skill", "Agent"] +tools: + [ + "Read", + "Write", + "Edit", + "Grep", + "Glob", + "Bash", + "Task", + "TodoWrite", + "Skill", + "Agent", + ] skills: - neuron-framework-equivalence --- @@ -51,31 +63,32 @@ Run the `/neuron-framework-equivalence` skill's 8-stage pipeline and produce an Route to the correct entry point based on user intent: -| User intent | Entry point | Prerequisites | -|---|---|---| -| Fresh validation (no prior work) | Stage 0 | None — collect Required Inputs first | -| Component tests already written, need to run | Stage 2 | Stage 0 complete, tests exist | -| Known component failures, need debugging | Stage 4 | Stage 2 results exist | -| CPU passes, device fails | Stage 5 + device-e2e-debugging | Stages 0–4 complete | -| All stages done, need report | Step 6 | All stages have concrete results | +| User intent | Entry point | Prerequisites | +| -------------------------------------------- | ------------------------------ | ------------------------------------ | +| Fresh validation (no prior work) | Stage 0 | None — collect Required Inputs first | +| Component tests already written, need to run | Stage 2 | Stage 0 complete, tests exist | +| Known component failures, need debugging | Stage 4 | Stage 2 results exist | +| CPU passes, device fails | Stage 5 + device-e2e-debugging | Stages 0–4 complete | +| All stages done, need report | Step 6 | All stages have concrete results | If the user's intent is ambiguous, ask which stage they want to enter. ## Tolerance Guidelines -| Precision | Threshold | Notes | -|---|---|---| -| FP32 strict | rtol=1e-5 | TP=1 FP32 baseline must match within this | -| BF16 R-ratio | < 1.2 | Component and E2E three-tensor comparison | -| Token match | Exact | Greedy-decoded tokens must match between source and target | -| KL divergence | < 0.01 | Per-position distributional equivalence | -| Cosine similarity | > 0.95 | Per-position semantic consistency | +| Precision | Threshold | Notes | +| ----------------- | --------- | ---------------------------------------------------------- | +| FP32 strict | rtol=1e-5 | TP=1 FP32 baseline must match within this | +| BF16 R-ratio | < 1.2 | Component and E2E three-tensor comparison | +| Token match | Exact | Greedy-decoded tokens must match between source and target | +| KL divergence | < 0.01 | Per-position distributional equivalence | +| Cosine similarity | > 0.95 | Per-position semantic consistency | ## Behavioral Modes ### Validation mode (Stages 0–3, 5–7) You are a **test runner**. Run scripts, record results, continue on failure. Do NOT: + - Investigate why a test failed - Read source code to understand root causes - Write patches or fixes @@ -87,6 +100,7 @@ You are a **test runner**. Run scripts, record results, continue on failure. Do ### Debugging mode (Stage 4 only) You are a **debugger**. Read source code, diagnose root causes, write monkey patches. Follow the escalation workflow in `references/debug-orchestration.md`: + 1. CPU components first (bottom-up, simplest to most complex) 2. Device components second 3. CPU E2E third diff --git a/agents/neuron-nki-agent.md b/agents/neuron-nki-agent.md index 52361ee..c4f5702 100644 --- a/agents/neuron-nki-agent.md +++ b/agents/neuron-nki-agent.md @@ -31,15 +31,24 @@ description: | model: opus color: green -tools: ["Read", "Write", "Edit", "Grep", "Glob", "Bash", "Task", "TodoWrite", "Skill"] +tools: + [ + "Read", + "Write", + "Edit", + "Grep", + "Glob", + "Bash", + "Task", + "TodoWrite", + "Skill", + ] skills: - neuron-nki-writing - neuron-nki-debugging - neuron-nki-docs - neuron-nki-profiling - neuron-nki-profile-querying - - --- # NKI Agent @@ -54,12 +63,12 @@ CRITICAL: Before writing any NKI code, you MUST read `skills/neuron-nki-writing/ Determine which workflow to use based on the request: -| Request Type | Workflow | Key Skills | -|-------------|----------|------------| -| Write new kernel or modify existing | [Write](#write) | `/neuron-nki-writing`, `/neuron-nki-docs` | -| Fix compilation errors | [Debug](#debug) | `/neuron-nki-debugging`, `/neuron-nki-docs` | -| Query profile data with SQL | [Query Profile](#query-profile) | `/neuron-nki-profile-querying` | -| Look up API/error docs | [Explore Docs](#explore-docs) | `/neuron-nki-docs` | +| Request Type | Workflow | Key Skills | +| ----------------------------------- | ------------------------------- | ------------------------------------------- | +| Write new kernel or modify existing | [Write](#write) | `/neuron-nki-writing`, `/neuron-nki-docs` | +| Fix compilation errors | [Debug](#debug) | `/neuron-nki-debugging`, `/neuron-nki-docs` | +| Query profile data with SQL | [Query Profile](#query-profile) | `/neuron-nki-profile-querying` | +| Look up API/error docs | [Explore Docs](#explore-docs) | `/neuron-nki-docs` | ## Write @@ -83,14 +92,15 @@ When fixing compilation errors, follow these principles in order: Common fixes: -| Error Pattern | Fix | -|---------------|-----| -| "missing `dst` parameter" | Add `dst=result` to ISA function | -| "PSUM buffer required" | Change `buffer=nl.sbuf` to `buffer=nl.psum` | -| "exceeds SBUF limit" | Reduce tile size in free dimension | -| "deprecated API" | Consult `nki-language-constraint.md` for correct patterns | +| Error Pattern | Fix | +| ------------------------- | --------------------------------------------------------- | +| "missing `dst` parameter" | Add `dst=result` to ISA function | +| "PSUM buffer required" | Change `buffer=nl.sbuf` to `buffer=nl.psum` | +| "exceeds SBUF limit" | Reduce tile size in free dimension | +| "deprecated API" | Consult `nki-language-constraint.md` for correct patterns | **Simplification hierarchy** (apply in order when stuck): + 1. Reduce tile sizes → 2. Simplify tiling strategy → 3. Break apart fused operations → 4. Use simpler data types → 5. Reduce parallelism Max 10 iterations. Save backup before starting: `cp {kernel_file} {kernel_file}.pre-debug` @@ -118,12 +128,12 @@ For API lookups, error codes, tutorials: ## Hardware Constraints Reference -| Constraint | Limit | Buffer | -|------------|-------|--------| -| Partition dimension (P) | ≤ 128 | SBUF/PSUM | -| PSUM free dimension | ≤ 512 (gen2/3) / ≤ 4096 (gen4) | PSUM | -| SBUF free dimension | ≤ 32767 | SBUF | -| MatMul K dimension | ≤ 2048 | N/A | +| Constraint | Limit | Buffer | +| ----------------------- | ------------------------------ | --------- | +| Partition dimension (P) | ≤ 128 | SBUF/PSUM | +| PSUM free dimension | ≤ 512 (gen2/3) / ≤ 4096 (gen4) | PSUM | +| SBUF free dimension | ≤ 32767 | SBUF | +| MatMul K dimension | ≤ 2048 | N/A | ## Neuron Core Isolation @@ -138,10 +148,10 @@ os.environ['NEURON_RT_INSPECT_OUTPUT_DIR'] = f'./output/nki-{os.getpid()}' ## Skill Invocations -| Situation | Skill | -|-----------|-------| -| Write/modify kernel | `/neuron-nki-writing` | -| Debug compilation error | `/neuron-nki-debugging` | -| Look up API/error code | `/neuron-nki-docs {topic}` | -| Profile kernel | `/neuron-nki-profiling {kernel_file}` | -| Query profile with SQL | `/neuron-nki-profile-querying` | +| Situation | Skill | +| ----------------------- | ------------------------------------- | +| Write/modify kernel | `/neuron-nki-writing` | +| Debug compilation error | `/neuron-nki-debugging` | +| Look up API/error code | `/neuron-nki-docs {topic}` | +| Profile kernel | `/neuron-nki-profiling {kernel_file}` | +| Query profile with SQL | `/neuron-nki-profile-querying` | diff --git a/agents/neuron-nki-debugger-agent.md b/agents/neuron-nki-debugger-agent.md index b623805..470fab9 100644 --- a/agents/neuron-nki-debugger-agent.md +++ b/agents/neuron-nki-debugger-agent.md @@ -32,7 +32,18 @@ description: | model: opus color: orange -tools: ["Read", "Write", "Edit", "Grep", "Glob", "Bash", "Task", "TodoWrite", "Skill"] +tools: + [ + "Read", + "Write", + "Edit", + "Grep", + "Glob", + "Bash", + "Task", + "TodoWrite", + "Skill", + ] skills: - neuron-nki-docs - neuron-nki-writing @@ -49,7 +60,6 @@ CRITICAL: All NKI code you generate MUST follow the language constraints defined **Read `/neuron-nki-writing` reference `nki-language-constraint.md` for the full constraint table and reference kernel. If you cannot load the skill, follow the reference kernel in the description examples above.** - ## Debugging Philosophy Follow these core principles in order: @@ -68,6 +78,7 @@ Execute these phases in order. Track iterations to prevent infinite loops (max 1 ### Phase 1: Analyze Error Message 1. **Run compilation** to capture the full error output: + ```bash source $NKI_VENV_PATH/bin/activate python test_{kernel_name}.py @@ -80,6 +91,7 @@ python test_{kernel_name}.py - Any suggestions in the error message 3. **Create error analysis** in your report: + ```markdown ## Error Analysis @@ -98,14 +110,14 @@ python test_{kernel_name}.py 2. **Common obvious fixes:** -| Error Pattern | Fix | -|---------------|-----| -| "missing `dst` parameter" | Add `dst=result` to ISA function | -| "PSUM buffer required" | Change `buffer=nl.sbuf` to `buffer=nl.psum` | -| "exceeds SBUF limit" | Reduce tile size in free dimension | -| "exceeds PSUM limit" | Reduce MatMul result tile size | -| "dimension must be <= 128" | Set partition dimension to 128 or less | -| "deprecated API" | Use Beta 2 API (e.g., `nisa.dma_copy` not `nl.load`) | +| Error Pattern | Fix | +| -------------------------- | ---------------------------------------------------- | +| "missing `dst` parameter" | Add `dst=result` to ISA function | +| "PSUM buffer required" | Change `buffer=nl.sbuf` to `buffer=nl.psum` | +| "exceeds SBUF limit" | Reduce tile size in free dimension | +| "exceeds PSUM limit" | Reduce MatMul result tile size | +| "dimension must be <= 128" | Set partition dimension to 128 or less | +| "deprecated API" | Use Beta 2 API (e.g., `nisa.dma_copy` not `nl.load`) | 3. **If fix is obvious:** - Apply the fix using Edit tool @@ -136,6 +148,7 @@ When the fix is not obvious from the error message, search for reference impleme - Look for simpler implementations of the same operation **Example search strategy:** + ```python # If error is in tensor_reduce operation: /neuron-nki-docs tensor_reduce # Get API documentation @@ -171,8 +184,10 @@ If error persists after trying obvious fixes and documented patterns, progressiv - Process data in smaller batches **Document trade-offs:** + ```markdown **Performance Trade-off:** + - Original: Fused matmul + softmax in single pass - Simplified: Separated into two passes with intermediate HBM write - Impact: ~2x increase in memory bandwidth, ~30% slower execution @@ -182,6 +197,7 @@ If error persists after trying obvious fixes and documented patterns, progressiv ### Phase 5: Test and Validate 1. **Compile the fixed kernel:** + ```bash source $NKI_VENV_PATH/bin/activate python test_{kernel_name}.py @@ -211,7 +227,7 @@ python test_{kernel_name}.py Every debugging session produces a structured report: -```markdown +````markdown # Debugging Report: {kernel_name} **Status:** {RESOLVED | BLOCKED | IN_PROGRESS} @@ -224,16 +240,19 @@ Every debugging session produces a structured report: ### Iteration 1: {error_code} **Error Analysis:** + - Line: {line_number} - Issue: {description} **Fix Applied:** + - Type: {obvious_fix | documented_pattern | simplification} - Changes: {description} **Performance Trade-off:** {if applicable} **Code Changes:** + ```python # Before {old_code} @@ -241,6 +260,7 @@ Every debugging session produces a structured report: # After {new_code} ``` +```` **Result:** {COMPILATION_SUCCESS | COMPILATION_FAILED | NEW_ERROR} @@ -252,15 +272,16 @@ Every debugging session produces a structured report: ## Artifacts -| Type | Path | -|------|------| +| Type | Path | +| --------------- | ----------------------- | | Original kernel | {kernel_file}.pre-debug | -| Fixed kernel | {kernel_file} | -| Test script | test_{kernel_name}.py | +| Fixed kernel | {kernel_file} | +| Test script | test\_{kernel_name}.py | ## Recommendations {Any suggestions for further improvements, performance recovery, or alternative approaches} + ``` ## Skill Invocations @@ -291,11 +312,13 @@ Use the Skill tool to invoke these skills as part of the workflow: **Recovery Strategy:** ``` + Iterations 1-3: Apply obvious fixes and documented solutions Iterations 4-6: Search examples and apply reference patterns Iterations 7-9: Simplify aggressively, sacrifice performance Iteration 10: Report blocked state, request user guidance -``` + +```` ## Hardware Constraints Reference @@ -316,7 +339,7 @@ When running concurrently with other agents (e.g., optimizer profiling on anothe import os os.environ["NEURON_RT_VISIBLE_CORES"] = "0" # Pin to core 0 os.environ["NEURON_CC_FLAGS"] = "--target trn2 --lnc 1" -``` +```` Also use a session-unique output directory for NEFF artifacts: @@ -329,14 +352,17 @@ See `references/neuron-core-isolation.md` for core detection and allocation patt ## Before You Begin 1. **Save backup:** + ```bash cp {kernel_file} {kernel_file}.pre-debug ``` 2. **Verify environment:** + - `$NKI_VENV_PATH` is set (from `.claude/nki-dev-suite.local.md` or environment) - Kernel test file exists or create minimal test 3. **Initialize tracking:** + - Create debugging report structure - Set iteration counter to 0 diff --git a/agents/neuron-nki-profile-analysis-agent.md b/agents/neuron-nki-profile-analysis-agent.md index b41c599..28b501d 100644 --- a/agents/neuron-nki-profile-analysis-agent.md +++ b/agents/neuron-nki-profile-analysis-agent.md @@ -84,11 +84,13 @@ When comparing before/after an optimization step: ## Error Handling If profiling fails: + 1. Check venv is activated with `neuronxcc` and `nki` packages 2. Verify `neuron-explorer` is on PATH 3. Check NEFF file exists and is valid If analysis is incomplete: + 1. Check DmaPacket table has data (re-profile with `NEURON_RT_ENABLE_DGE_NOTIFICATIONS=1` if empty) 2. Check `bir_debug_info_source_location` is populated (re-compile with debug info if NULL) 3. Verify parquet files exist at the expected data-path diff --git a/agents/neuron-nki-writer-agent.md b/agents/neuron-nki-writer-agent.md index 4d61634..1ef6185 100644 --- a/agents/neuron-nki-writer-agent.md +++ b/agents/neuron-nki-writer-agent.md @@ -46,7 +46,7 @@ skills: # NKI Writer Agent -You are an expert NKI kernel author. Your role is to write new NKI kernels and modify existing ones — whether translating from PyTorch/NumPy/natural language, adding shape/dtype support, refactoring tiling, or implementing new features. All output follows the latest NKI version API pattern. +You are an expert NKI kernel author. Your role is to write new NKI kernels and modify existing ones — whether translating from PyTorch/NumPy/natural language, adding shape/dtype support, refactoring tiling, or implementing new features. All output follows the latest NKI version API pattern. ## NKI Language Constraints (MANDATORY) @@ -54,7 +54,6 @@ CRITICAL: All NKI code you generate MUST follow the language constraints defined **Read `/neuron-nki-writing` reference `nki-language-constraint.md` for the full constraint table and reference kernel. If you cannot load the skill, follow the reference kernel in the description examples above.** - ## Workflow: New Kernel When creating a kernel from a PyTorch/NumPy/natural language specification: @@ -65,6 +64,7 @@ When creating a kernel from a PyTorch/NumPy/natural language specification: 4. **Validate** — build a test harness comparing against a CPU reference (never XLA device — each on-device graph generates a separate NEFF). For complex kernels, validate incrementally stage-by-stage per the skill's validation guidance **Capabilities worth reaching for (look up details via `/neuron-nki-docs`):** + - **Native `NkiTensor` view methods** — call zero-copy views directly on a tensor (`t.slice`, `t.select`, `t.permute`, `t.broadcast`, `t.expand_dim`, `t.squeeze_dim`, `t.reshape_dim`, `t.flatten_dims`, `t.rearrange`, `t.reshape`, `t.view`, `t.vector_select`, plus `t.is_contiguous()` / `t.is_indirect()`) instead of hand-coding `.ap()` for reshapes/slices. See `api-nki-tensor.md`. - **Tensor indirection on compute ops (`.indirect()`)** — on NeuronCore-v4+, do on-chip gather/scatter by passing a `.indirect(index)` view as `dst`/`data` to compute ops (`nc_matmul`, `nc_matmul_mx`, `tensor_tensor`, `tensor_scalar`, `tensor_reduce`, `tensor_copy`, `tensor_copy_predicated`, `tensor_scalar_reduce`, `tensor_scalar_cumulative`, `activation`, `activation_reduce`, `activate2`, `exponential`), subject to quadrant/partition-alignment rules. Extends the DMA-only `vector_offset` indirection to compute. See `NkiTensor.indirect`. diff --git a/skills/neuron-explorer-profile-schema/SKILL.md b/skills/neuron-explorer-profile-schema/SKILL.md index ab392e0..6658fd9 100644 --- a/skills/neuron-explorer-profile-schema/SKILL.md +++ b/skills/neuron-explorer-profile-schema/SKILL.md @@ -99,18 +99,18 @@ curl -s -X POST http://localhost:3002/api/v1/db/${PROFILE_NAME}/_search \ Every table is one of the following modalities. Knowing the modality tells you how to read the rows. -| Modality | What a row is | Examples | -|---|---|---| -| Timeline of events | One row per discrete event with `start_ts`/`end_ts` (or single `timestamp`) | `Instruction`, `DmaPacket`, `DmaPacketAggregated`, `ActiveTime`, `SemaphoreUpdate`, `Throttle`, `Error`, `CcOp`, `CoreBarriers`, `SystemProfileEvents`, `SbufAllocation` | -| Time-series samples | One row per sampled tick on a time axis | `DmaUsage`, `HbmUsage`, `PsumUsage`, `SbufUsage`, `PendingDma`, `CpuUsage`, `HostMemUsage`, `SystemProfileHbmUsage` | -| Dependency graph edges | One row per directed edge between rows in timeline tables (e.g. an instruction → the DMA it triggered) | `Flow` | -| Hierarchical aggregation | One row per node in a compiler IR hierarchy (Framework → HLO → Penguin → BIR → Instruction) with rolled-up statistics | `FrameworkInstruction`, `HloInstruction`, `PenguinInstruction`, `BirInstruction`, `FrameworkNode` | -| Aggregated summary | Computed roll-up across the whole profile (or by a key) | `Summary`, `OpcodeSummary`, `ThrottleSummary`, `HbmUsageSummaryByType` | -| Reference / lookup | Static dimension table referenced by other rows via foreign keys | `TensorInfo`, `DmaQueuesInfo`, `CcStream`, `StackFrame`, `StackFrameFileLocation`, `StackFrameFunctionName`, `StackFrameFileName`, `KernelStackFrames`, `KernelIterationVariables`, `KernelInstructions`, `AssemblyInstruction`, `DeviceProfileList` | -| Profile-level metadata | Single-row table describing the profile as a whole | `Metadata`, `NeffHeader`, `SystemProfileMetadata`, `ExecutionInfo` | -| Diagnostic messages | One row per warning emitted by `neuron-explorer` during ingestion. Always check this table first because rows can indicate a data quality problem. | `Warning` | -| Transient API response | Computed at query time and returned via the HTTP API; **not** written to parquet | `MemoryBandwidthPoint`, `MemoryBandwidthSeries`, `MemoryBandwidthResponse` | -| Enum | String enum referenced by other tables; not a standalone parquet table | `DmaQueueType`, `ErrorType`, `PerformanceMode`, `MemoryBandwidthDirection` | +| Modality | What a row is | Examples | +| ------------------------ | -------------------------------------------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| Timeline of events | One row per discrete event with `start_ts`/`end_ts` (or single `timestamp`) | `Instruction`, `DmaPacket`, `DmaPacketAggregated`, `ActiveTime`, `SemaphoreUpdate`, `Throttle`, `Error`, `CcOp`, `CoreBarriers`, `SystemProfileEvents`, `SbufAllocation` | +| Time-series samples | One row per sampled tick on a time axis | `DmaUsage`, `HbmUsage`, `PsumUsage`, `SbufUsage`, `PendingDma`, `CpuUsage`, `HostMemUsage`, `SystemProfileHbmUsage` | +| Dependency graph edges | One row per directed edge between rows in timeline tables (e.g. an instruction → the DMA it triggered) | `Flow` | +| Hierarchical aggregation | One row per node in a compiler IR hierarchy (Framework → HLO → Penguin → BIR → Instruction) with rolled-up statistics | `FrameworkInstruction`, `HloInstruction`, `PenguinInstruction`, `BirInstruction`, `FrameworkNode` | +| Aggregated summary | Computed roll-up across the whole profile (or by a key) | `Summary`, `OpcodeSummary`, `ThrottleSummary`, `HbmUsageSummaryByType` | +| Reference / lookup | Static dimension table referenced by other rows via foreign keys | `TensorInfo`, `DmaQueuesInfo`, `CcStream`, `StackFrame`, `StackFrameFileLocation`, `StackFrameFunctionName`, `StackFrameFileName`, `KernelStackFrames`, `KernelIterationVariables`, `KernelInstructions`, `AssemblyInstruction`, `DeviceProfileList` | +| Profile-level metadata | Single-row table describing the profile as a whole | `Metadata`, `NeffHeader`, `SystemProfileMetadata`, `ExecutionInfo` | +| Diagnostic messages | One row per warning emitted by `neuron-explorer` during ingestion. Always check this table first because rows can indicate a data quality problem. | `Warning` | +| Transient API response | Computed at query time and returned via the HTTP API; **not** written to parquet | `MemoryBandwidthPoint`, `MemoryBandwidthSeries`, `MemoryBandwidthResponse` | +| Enum | String enum referenced by other tables; not a standalone parquet table | `DmaQueueType`, `ErrorType`, `PerformanceMode`, `MemoryBandwidthDirection` | ## Data flow at a glance @@ -161,15 +161,15 @@ flowchart LR Producer → input artifact → `neuron-explorer` -> output data table: -| Producer | Input artifact | Output Data Tables | -|---|---|---| -| Neuron Hardware | `*.ntff` (binary device trace) | `Instruction`, `DmaPacket`, `SemaphoreUpdate`, `Throttle`, `Error`, `CoreBarriers`, etc. | -| Neuron Runtime | `ntrace.pb` + `cpu_util.pb` + `host_mem.pb` + `trace_info.pb` (host protobuf) | `SystemProfileEvents` (`trace_event_source = neuron_rt` / `neuron_hw`), `CpuUsage`, `HostMemUsage`, `SystemProfileMetadata`, etc. | -| Frameworks (PyTorch, JAX, vLLM) | `*/plugins/*/trace.json.gz` (Chrome Trace JSON) | `SystemProfileEvents` (`trace_event_source = framework`) | -| Neuron Compiler — model | `*.neff` archive (header, tensor + queue manifest) | `NeffHeader`, `TensorInfo`, `DmaQueuesInfo`, `Metadata`, etc. | -| Neuron Compiler — debug info | `/debug_info/{framework,hlo,penguin,backend}.dbg` + `stack_frame_index.dbg` | IR hierarchy `FrameworkInstruction` → `HloInstruction` → `PenguinInstruction` → `BirInstruction`, plus `FrameworkNode` and the four `StackFrame*` tables | -| NKI front-end | `/kernel_debug_info.json` + per-kernel JSON | `KernelInstructions`, `KernelStackFrames`, `KernelIterationVariables`, `Instruction.nki_source_location` | -| User upload | `source_folder.tar.gz` (gzipped tar of `.py` files) | Not in any table — served by the `/fs/*` API for UI source rendering | +| Producer | Input artifact | Output Data Tables | +| ------------------------------- | --------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------- | +| Neuron Hardware | `*.ntff` (binary device trace) | `Instruction`, `DmaPacket`, `SemaphoreUpdate`, `Throttle`, `Error`, `CoreBarriers`, etc. | +| Neuron Runtime | `ntrace.pb` + `cpu_util.pb` + `host_mem.pb` + `trace_info.pb` (host protobuf) | `SystemProfileEvents` (`trace_event_source = neuron_rt` / `neuron_hw`), `CpuUsage`, `HostMemUsage`, `SystemProfileMetadata`, etc. | +| Frameworks (PyTorch, JAX, vLLM) | `*/plugins/*/trace.json.gz` (Chrome Trace JSON) | `SystemProfileEvents` (`trace_event_source = framework`) | +| Neuron Compiler — model | `*.neff` archive (header, tensor + queue manifest) | `NeffHeader`, `TensorInfo`, `DmaQueuesInfo`, `Metadata`, etc. | +| Neuron Compiler — debug info | `/debug_info/{framework,hlo,penguin,backend}.dbg` + `stack_frame_index.dbg` | IR hierarchy `FrameworkInstruction` → `HloInstruction` → `PenguinInstruction` → `BirInstruction`, plus `FrameworkNode` and the four `StackFrame*` tables | +| NKI front-end | `/kernel_debug_info.json` + per-kernel JSON | `KernelInstructions`, `KernelStackFrames`, `KernelIterationVariables`, `Instruction.nki_source_location` | +| User upload | `source_folder.tar.gz` (gzipped tar of `.py` files) | Not in any table — served by the `/fs/*` API for UI source rendering | ## Source code linking @@ -186,10 +186,10 @@ names from the model root down to the op (e.g. Captured at compile time by the framework's tracer; populated only for compiled PyTorch flows. -| Schema location | Field | Example | -|---|---|---| -| Per-instruction string | `Instruction.layer` | `LlamaDecoderLayer[0]_dot.4` | -| Top level of hierarchy | `FrameworkInstruction.framework_name` | `LlamaDecoderLayer[1]/function[2]/aten.add` | +| Schema location | Field | Example | +| ---------------------- | -------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------- | +| Per-instruction string | `Instruction.layer` | `LlamaDecoderLayer[0]_dot.4` | +| Top level of hierarchy | `FrameworkInstruction.framework_name` | `LlamaDecoderLayer[1]/function[2]/aten.add` | | Path decomposed by `/` | `FrameworkNode.{node_name, parent_name, children_names}` | `node_name=LlamaAttention[attn][0]`, `parent_name=LlamaDecoderLayer[0]`, `children_names=[Linear[q_proj][0], Linear[k_proj][0]]` | ### 2. Python source location + stack frame index (compiled flows) @@ -198,12 +198,12 @@ The Python file, line number, and function name for each instruction, plus the full caller chain. Captured from PyTorch frame info at compile time and embedded in the NEFF debug info; empty for eager-mode profiles. -| Schema location | Field | Example | -|---|---|---| -| Per-instruction list of frame ids | `Instruction.stack_frame_ids` | `[101, 102, 103]` | -| Frame, with parent pointer | `StackFrame.{id, parent_frame_id, file_location_id}` | `id=103`, `parent_frame_id=102`, `file_location_id=42` | -| Resolved location | `StackFrameFileLocation.{file_name_id, function_name_id, line_number}` | `file_name_id=7`, `function_name_id=15`, `line_number=58` | -| Interned strings | `StackFrameFileName.name`, `StackFrameFunctionName.name` | `StackFrameFileName.name=model.py`, `StackFrameFunctionName.name=forward` | +| Schema location | Field | Example | +| --------------------------------- | ---------------------------------------------------------------------- | ------------------------------------------------------------------------- | +| Per-instruction list of frame ids | `Instruction.stack_frame_ids` | `[101, 102, 103]` | +| Frame, with parent pointer | `StackFrame.{id, parent_frame_id, file_location_id}` | `id=103`, `parent_frame_id=102`, `file_location_id=42` | +| Resolved location | `StackFrameFileLocation.{file_name_id, function_name_id, line_number}` | `file_name_id=7`, `function_name_id=15`, `line_number=58` | +| Interned strings | `StackFrameFileName.name`, `StackFrameFunctionName.name` | `StackFrameFileName.name=model.py`, `StackFrameFunctionName.name=forward` | A single instruction may carry multiple `stack_frame_ids` because compiler fusions collapse multiple source locations onto one hardware instruction. @@ -215,13 +215,13 @@ for each kernel instruction, plus the kernel call stack and the surrounding loop-nest iteration variables. Captured by the NKI front-end and bundled into the NEFF; empty for non-NKI workloads. -| Schema location | Field | Example | -|---|---|---| -| Direct `:` for the NKI op | `Instruction.nki_source_location` | `/home/ubuntu/decoder.py:139` | -| BIR-recorded source location | `Instruction.bir_debug_info_source_location` | `/home/ubuntu/decoder.py:139` | -| Kernel-instruction lookup | `KernelInstructions.{file_path, line_number, stack_frame_id, iteration_variables_id}` | `file_path=/home/ubuntu/kernel.py`, `line_number=76`, `stack_frame_id=42`, `iteration_variables_id=156` | -| Kernel call stack (linked list of frames) | `KernelStackFrames.{fully_qualified_function_name, file_path, line_number, parent_stack_frame_id}` | `fully_qualified_function_name=nki.kernels.matmul`, `file_path=/home/ubuntu/kernel.py`, `line_number=17`, `parent_stack_frame_id=41` | -| Loop nest (linked list of iter vars) | `KernelIterationVariables.{variable_name, variable_value, file_path, line_number, parent_iteration_variable_id}` | `variable_name=k`, `variable_value=3`, `file_path=/home/ubuntu/kernel.py`, `line_number=72`, `parent_iteration_variable_id=101` | +| Schema location | Field | Example | +| ----------------------------------------- | ---------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------ | +| Direct `:` for the NKI op | `Instruction.nki_source_location` | `/home/ubuntu/decoder.py:139` | +| BIR-recorded source location | `Instruction.bir_debug_info_source_location` | `/home/ubuntu/decoder.py:139` | +| Kernel-instruction lookup | `KernelInstructions.{file_path, line_number, stack_frame_id, iteration_variables_id}` | `file_path=/home/ubuntu/kernel.py`, `line_number=76`, `stack_frame_id=42`, `iteration_variables_id=156` | +| Kernel call stack (linked list of frames) | `KernelStackFrames.{fully_qualified_function_name, file_path, line_number, parent_stack_frame_id}` | `fully_qualified_function_name=nki.kernels.matmul`, `file_path=/home/ubuntu/kernel.py`, `line_number=17`, `parent_stack_frame_id=41` | +| Loop nest (linked list of iter vars) | `KernelIterationVariables.{variable_name, variable_value, file_path, line_number, parent_iteration_variable_id}` | `variable_name=k`, `variable_value=3`, `file_path=/home/ubuntu/kernel.py`, `line_number=72`, `parent_iteration_variable_id=101` | ### 4. Framework call stack (system profile) @@ -231,10 +231,10 @@ PyTorch's profiler at runtime when execution is wrapped in `torch.profiler.profile`; works with both eager and compiled flows. With `with_stack=True`, event names include the source file and line. -| Schema location | Field | Example | -|---|---|---| -| Rows with `trace_event_source = framework` | `SystemProfileEvents.name` | `_prepare_inputs` | -| Rows with `trace_event_source = framework` captured with PyTorch `with_stack=True` | `SystemProfileEvents.name` | `torch_neuronx/neuron_dynamo_backend/executor.py(83): _prepare_inputs` | +| Schema location | Field | Example | +| ---------------------------------------------------------------------------------- | -------------------------- | ---------------------------------------------------------------------- | +| Rows with `trace_event_source = framework` | `SystemProfileEvents.name` | `_prepare_inputs` | +| Rows with `trace_event_source = framework` captured with PyTorch `with_stack=True` | `SystemProfileEvents.name` | `torch_neuronx/neuron_dynamo_backend/executor.py(83): _prepare_inputs` | ### Accessing full source code files @@ -257,13 +257,13 @@ curl -s "http://localhost:3002/api/v1/profiles/namespace/global/profile_name/${P ## Bundled scripts -| Script | Purpose | -|---|---| +| Script | Purpose | +| -------------------------------------------------------- | ---------------------------------------------------------------------------------- | | `scripts/write_profile_schema_to_separate_yaml_files.py` | Split `neuron-explorer --show-profile-schema` output into one YAML file per table. | ## Related skills -| Skill | Purpose | -|---|---| +| Skill | Purpose | +| ------------------------------ | -------------------------------------------------------- | | `/neuron-nki-profile-querying` | Run SQL / Python on parquet against an ingested profile. | -| `/neuron-nki-profiling` | Capture NEFF + NTFF on Trainium/Inferentia hardware. | +| `/neuron-nki-profiling` | Capture NEFF + NTFF on Trainium/Inferentia hardware. | diff --git a/skills/neuron-framework-autoport-vllm-neuron/SKILL.md b/skills/neuron-framework-autoport-vllm-neuron/SKILL.md index faea3b4..3ab6868 100644 --- a/skills/neuron-framework-autoport-vllm-neuron/SKILL.md +++ b/skills/neuron-framework-autoport-vllm-neuron/SKILL.md @@ -15,12 +15,13 @@ Port the HuggingFace model `$1` to the vLLM-Neuron Trainium2 backend as `$0`. - `HF_MODEL_ID` = `$1` (e.g., `01-ai/Yi-6B-Chat`) - `--review` flag (if present in `$ARGUMENTS`): pause at Step 2 for user confirmation - Derive `PascalName` from MODEL_NAME (e.g., `yi` → `Yi`, `gpt_neox` → `GPTNeoX`) -- Model code goes in `vllm_neuron/model` +- Model code goes in `vllm_neuron/model` - Example script goes in `examples/vllm_neuron/models/MODEL_NAME/` ## Dry-run When the user specifies `dry-run`: + - **Skip** all agent-level prerequisites (package import checks, NeuronCore check via `neuron-ls`) - **Activate** the provided venv and resolve source paths by filesystem lookup (do NOT use `import` — dependencies are not installed): ```bash @@ -54,6 +55,7 @@ Ensure the agent-level prerequisites have been completed (venv activation, packa Gather all information needed to port the model. Do NOT write any files yet. **1a. Fetch HF config:** + ```python python3 -c " from transformers import AutoConfig @@ -64,6 +66,7 @@ print(json.dumps(c.to_dict(), indent=2, default=str)) ``` Extract and record: + - `architectures` field (exact string for registry, e.g., `"YiForCausalLM"`) - `hidden_size`, `intermediate_size`, `num_hidden_layers` - `num_attention_heads`, `num_key_value_heads`, `head_dim` (compute if missing: hidden_size / num_attention_heads) @@ -77,6 +80,7 @@ Extract and record: **1b. Read HF transformers source:** Find the modeling file for this architecture. Check: + - Exact weight names (for checkpoint key mapping in `load_weights`) - Whether Q, K, V projections are fused or separate in the checkpoint - Bias presence on each projection (Q, K, V, O, gate, up, down, layernorm) @@ -85,6 +89,7 @@ Find the modeling file for this architecture. Check: **CRITICAL — Verify normalization type by reading the forward() implementation, not the class name.** Some models name their norm class "RMSNorm" but actually implement full LayerNorm (with mean subtraction and bias). The difference: + - RMSNorm: `variance = x.pow(2).mean(); x = x * rsqrt(variance + eps)` — no mean subtraction, no bias - LayerNorm: `mean = x.mean(); variance = (x - mean).pow(2).mean(); x = (x - mean) / sqrt(variance + eps)` — has mean subtraction, may have bias @@ -96,26 +101,27 @@ Read the canonical model bringup guide at `doc/vllm_neuron/source/design/framewo **1d. Select best reference model:** Based on the architecture analysis, pick the closest existing vLLM-Neuron model as your copy source: -| If the model has... | Use reference | -|---|---| -| Standard GQA + RoPE (most common) | `llama3/` | -| QKV bias | `llama3/` (add bias params to QKV/O projections) | -| Q/K per-head RMSNorm | `qwen3/` (dense) or `qwen3_moe/` (MoE) | -| ALiBi attention (no RoPE) | `bloom/` | -| Mixture of Experts (fits in TP) | `qwen3_moe/` | -| Mixture of Experts (needs EP) | `gpt_oss/` (see EP section below) | -| Large MoE with MLA/multi-latent attention | `deepseek_v32/` | -| Learned position embeddings | `gpt2/` | -| Parallel residual | `gptj/` | -| Non-gated MLP (no gate_proj) | `starcoder2/` | -| Multi-Query Attention (1 KV head) | `starcoder2/` (set num_key_value_heads=1) | -| Vision-language model | `qwen3_vl/` | -| Full LayerNorm with bias (not RMSNorm) | `starcoder2/` | +| If the model has... | Use reference | +| ----------------------------------------- | ------------------------------------------------ | +| Standard GQA + RoPE (most common) | `llama3/` | +| QKV bias | `llama3/` (add bias params to QKV/O projections) | +| Q/K per-head RMSNorm | `qwen3/` (dense) or `qwen3_moe/` (MoE) | +| ALiBi attention (no RoPE) | `bloom/` | +| Mixture of Experts (fits in TP) | `qwen3_moe/` | +| Mixture of Experts (needs EP) | `gpt_oss/` (see EP section below) | +| Large MoE with MLA/multi-latent attention | `deepseek_v32/` | +| Learned position embeddings | `gpt2/` | +| Parallel residual | `gptj/` | +| Non-gated MLP (no gate_proj) | `starcoder2/` | +| Multi-Query Attention (1 KV head) | `starcoder2/` (set num_key_value_heads=1) | +| Vision-language model | `qwen3_vl/` | +| Full LayerNorm with bias (not RMSNorm) | `starcoder2/` | Read the selected reference model's `model.py`, `config.py`, and `factory.py`. **1e. Compute valid TP sizes:** All 5 rules must be satisfied: + 1. `num_attention_heads % tp_size == 0` 2. `(num_attention_heads / tp_size)` is **even** (NKI decode megakernel constraint) 3. `num_key_value_heads % tp_size == 0` OR `tp_size % num_key_value_heads == 0` (GQA replication) @@ -124,9 +130,10 @@ All 5 rules must be satisfied: **Memory constraint for real hardware:** Each Neuron device has ~24GB HBM (shared between 2 NeuronCores in LNC=2 config). With TP=N, each rank holds ~(model_size_bytes / N) of weights plus KV cache. If per-rank weight memory exceeds ~20GB, increase TP size. MoE models are especially memory-hungry — 16 experts × 3 projections × hidden × intermediate × 2 bytes adds up fast. -Recommend the smallest valid TP size ≥ 2 that also fits the model in memory (rule of thumb: 2x model params in bytes < tp_size * 96GB per device). +Recommend the smallest valid TP size ≥ 2 that also fits the model in memory (rule of thumb: 2x model params in bytes < tp_size \* 96GB per device). **1f. Inspect checkpoint weight keys:** + ```python python3 -c " from safetensors import safe_open @@ -139,6 +146,7 @@ for f in files: print(f'{k}: {sf.get_tensor(k).shape}') " ``` + This reveals which layers have biases, the exact key naming convention, and weight shapes. Essential for building the `load_weights` mapping correctly. **1g. Evaluate Expert Parallelism (EP) need (MoE models only):** @@ -146,6 +154,7 @@ This reveals which layers have biases, the exact key naming convention, and weig If the model is MoE and no single TP size satisfies all 5 rules above while also fitting in memory, the model needs **Expert Parallelism (EP)**. EP uses a two-level parallelism: TP_sub for attention/dense layers, EP for distributing experts across groups. Determine EP need: + - Compute `model_bytes = num_params * 2` (bf16) - If `model_bytes / max_valid_tp > 24GB per NC` (with lnc=2 on trn2.48xlarge, 96GB per device) → EP required - EP config: `world_size = 64` (full trn2.48xlarge), `ep_degree = world_size / tp_sub` @@ -256,9 +265,11 @@ Match the model's rotary embedding variant. Remove scaling if standard. Add/remove bias. Handle sliding window, GQA/MQA head counts. **CRITICAL — Bias shapes for NKI kernels:** The `NF.attention_decode` megakernel requires bias tensors to be 2D `[1, size]`, not 1D `(size,)`. When passing `bias_qkv` or `bias_out`, always unsqueeze: + ```python bias_qkv=self.qkv_proj_bias.unsqueeze(0) if self.qkv_proj_bias is not None else None ``` + This passes on CPU simulator (which accepts 1D) but fails on real Neuron hardware with: "Bias shape must be [1, I], got (768,), expected (1, 768)". For the prefill path, biases can be applied manually after the matmul (1D is fine for `torch.matmul` + bias addition). Only the decode megakernel has the 2D requirement. @@ -266,16 +277,20 @@ For the prefill path, biases can be applied manually after the matmul (1D is fin **Section 4 (MLP)**: Match activation function and gating. SwiGLU (gate+up+down) vs non-gated (up+down). For MoE models: keep router gate weights in float32, not bf16. The softmax over num_experts is extremely sensitive to precision — a tiny difference in bf16 can select a completely different pair of experts. In the model init: + ```python self.gate_weight = nn.Parameter(torch.empty(..., dtype=torch.float32)) ``` + In `load_weights`, preserve float32 for router weights: + ```python if 'gate_weight' in name: rank_sharded[name] = tensor.to(torch.float32) ``` **Section 4b (MoE with EP)**: If EP is required, implement the unified sp_group pattern: + ```python # In every module that does collectives: ep_degree = get_neuron_ep_degree() @@ -286,6 +301,7 @@ else: self.tp_group = get_tp_group() self.sp_group = self.tp_group ``` + - Use `self.tp_group.world_size` for parameter dimensions (num_heads_per_rank, etc.) - Use `self.sp_group` for every all_gather, reduce_scatter, all_reduce call - This applies to ALL modules: Attention, MLP, MoE, Embedding, LM Head, Sampler @@ -307,15 +323,18 @@ Rename ALL classes with the model's PascalCase prefix (e.g., `LlamaRMSNorm` → ### Step 5: Register the Model Edit `vllm_neuron/model/registry.py`: + 1. Add import: `from .MODEL_NAME import PascalNameForCausalLM` 2. Add tuple to `get_models()`: `("HFArchitectureString", PascalNameForCausalLM),` **CRITICAL: The HF architecture string must EXACTLY match the `architectures` field from config.json, including capitalization.** Example: + - HF config: `"architectures": ["PhiMoEForCausalLM"]` - Registry: `("PhiMoEForCausalLM", PhiMoEForCausalLM)` ← correct - Registry: `("PhimoeForCausalLM", PhiMoEForCausalLM)` ← WRONG, model won't load Always verify with: + ```bash python3 -c "from transformers import AutoConfig; c = AutoConfig.from_pretrained('HF_MODEL_ID'); print(c.architectures)" ``` @@ -329,12 +348,13 @@ python3 -c "from transformers import AutoConfig; c = AutoConfig.from_pretrained( ### Step 7: Self-Review Before reporting completion, verify: + - [ ] Every `nn.Parameter` in model.py has a corresponding key in `load_weights` mappings - [ ] Bias shapes are `[1, size]` (2D) for any param passed to `NF.attention_decode` (`bias_qkv`, `bias_out`) - [ ] No `.to(device)` calls in the forward path (tensors are already on device) - [ ] No `.item()` calls in the forward path (breaks torch.compile) - [ ] RoPE inv_freq is computed lazily in `forward()`, not in `__init__()` (avoids meta tensor) -- [ ] Class names are consistent across config.py, factory.py, __init__.py, model.py +- [ ] Class names are consistent across config.py, factory.py, **init**.py, model.py - [ ] `__init__.py` exports the factory's ForCausalLM (not the model's — they have the same name) - [ ] Registry uses the exact HF architecture string from config.json (case-sensitive) - [ ] `from_configs` classmethod properly maps all model-specific config fields @@ -343,6 +363,7 @@ Before reporting completion, verify: - [ ] Normalization type matches HF source (read forward(), don't trust class name) **Additional EP checks (MoE models with Expert Parallelism only):** + - [ ] ALL collectives use a single group (sp_group) — no mixed group sizes in the NEFF - [ ] `tp_group` used only for weight sizing, `sp_group` for all collective ops - [ ] O_proj and dense MLP down_proj weights scaled by `1/ep_degree` after loading @@ -361,6 +382,7 @@ Print summary: files created/modified, line counts, and any warnings. ### Step 8: Smoke Test Run the example script: + ```bash # Standard (non-EP) models: NEURON_SKIP_EFA_AFFINITY=1 python examples/vllm_neuron/models/MODEL_NAME/run.py @@ -372,29 +394,31 @@ NEURON_SKIP_EFA_AFFINITY=1 VLLM_NEURON_SWITCH_CC=1 python examples/vllm_neuron/m Note: `NEURON_SKIP_EFA_AFFINITY=1` is needed on trn2 instances where the PCI topology doesn't match the hardcoded BDF-to-EFA mapping. Safe for single-node TP. Check for: + - Successful weight loading (no missing key errors) - Successful compilation (no NKI shape mismatch errors) - Reasonable generated text (not garbage or degenerate repetition) **Diagnosing common failures:** -| Symptom | Likely Cause | Fix | -|---|---|---| -| `Cannot copy out of meta tensor` | RoPE or other tensor computed in `__init__` | Move computation to `forward()`, compute lazily | -| `Bias shape must be [1, I], got (N,)` | 1D bias passed to NKI decode megakernel | `.unsqueeze(0)` before passing to `NF.attention_decode` | -| `unimplemented _copy_from xla:0 neuron:N` | `.to(device)` call in forward path | Remove `.to(device)`, tensors are already on device | -| `Unsupported Tensor.item()` | `.item()` call in forward path | Use static values or tensor ops | -| `Checkpoint key(s) not found` | Weight mapping mismatch | Check `load_weights` mappings vs checkpoint keys | -| `size mismatch for lm_head.bias` | lm_head bias not mapped correctly | Use separate `nn.Parameter` + explicit mapping | -| `nrt_tensor_allocate status=4` | HBM out of memory | Increase TP size | -| `No EFA device found` | PCI topology mismatch | `NEURON_SKIP_EFA_AFFINITY=1` | -| Degenerate output (`the the the...`) | Decode path bug (bias, KV cache, or norm) | Test prefill only (max_tokens=1), then debug decode | -| Registry `AttributeError: no attribute 'from_configs'` | Architecture string mismatch in registry | Verify exact HF architecture string | +| Symptom | Likely Cause | Fix | +| ------------------------------------------------------ | ------------------------------------------- | ------------------------------------------------------- | +| `Cannot copy out of meta tensor` | RoPE or other tensor computed in `__init__` | Move computation to `forward()`, compute lazily | +| `Bias shape must be [1, I], got (N,)` | 1D bias passed to NKI decode megakernel | `.unsqueeze(0)` before passing to `NF.attention_decode` | +| `unimplemented _copy_from xla:0 neuron:N` | `.to(device)` call in forward path | Remove `.to(device)`, tensors are already on device | +| `Unsupported Tensor.item()` | `.item()` call in forward path | Use static values or tensor ops | +| `Checkpoint key(s) not found` | Weight mapping mismatch | Check `load_weights` mappings vs checkpoint keys | +| `size mismatch for lm_head.bias` | lm_head bias not mapped correctly | Use separate `nn.Parameter` + explicit mapping | +| `nrt_tensor_allocate status=4` | HBM out of memory | Increase TP size | +| `No EFA device found` | PCI topology mismatch | `NEURON_SKIP_EFA_AFFINITY=1` | +| Degenerate output (`the the the...`) | Decode path bug (bias, KV cache, or norm) | Test prefill only (max_tokens=1), then debug decode | +| Registry `AttributeError: no attribute 'from_configs'` | Architecture string mismatch in registry | Verify exact HF architecture string | **EP-specific failures (MoE models with Expert Parallelism):** + - `NEFF Warmup failed with status 1006` / DGE scatter/gather out-of-bound → Mixed collective group sizes in one NEFF. Grep the FX graph dump for different `replica_groups` sizes. Fix: unify ALL collectives to one group (sp_group). - OOM during weight loading → Expert weights not filtered by EP rank. Check that `load_weights` only maps `range(local_expert_start, local_expert_start + num_local_experts)`. -- Garbage output but no crash → Check O_proj/down_proj scaling (missing `div_(ep_degree)`?). Check router gate is replicated (NOT EP-sharded). Check shared expert division by ep_degree after collective. +- Garbage output but no crash → Check O*proj/down_proj scaling (missing `div*(ep_degree)`?). Check router gate is replicated (NOT EP-sharded). Check shared expert division by ep_degree after collective. - Compilation OOM/timeout → With many experts in a loop, the compiler unrolls all iterations. For production, switch to `NF.moe_cte` (prefill) and `NF.moe_block_tkg` (decode) kernels. - Run EP models with: `NEURON_SKIP_EFA_AFFINITY=1 VLLM_NEURON_SWITCH_CC=1 python run.py` @@ -403,6 +427,7 @@ Update `results/results.md` with offline inference results. ### Step 9: Online Serving & APC Start the vLLM API server: + ```bash NEURON_SKIP_EFA_AFFINITY=1 python3 -m vllm.entrypoints.openai.api_server \ --model HF_MODEL_ID \ @@ -411,6 +436,7 @@ NEURON_SKIP_EFA_AFFINITY=1 python3 -m vllm.entrypoints.openai.api_server \ ``` Test battery: + 1. Basic completion: `POST /v1/completions` with "The capital of France is" 2. Batch: 4 prompts in one request 3. Streaming: `stream=true`, verify SSE chunks @@ -429,16 +455,19 @@ Generate the canonical logit validation test from `test/vllm_neuron/model/templa 1. Create the test directory following the current layout convention: `test/vllm_neuron/model/MODEL_NAME/bf16/e2e/` (precision tier, then `e2e/`) — see existing tests like `test/vllm_neuron/model/qwen3/bf16/e2e/test_logits.py` for the pattern. 2. Generate `test/vllm_neuron/model/MODEL_NAME/bf16/e2e/test_logits.py` from the template, filling in all `{{...}}` variables based on the model's architecture (hidden_size, TP sizes, batch sizes, HF model ID, etc.). Reference existing tests in `test/vllm_neuron/model/` for examples of how other models fill in these variables. 3. Run the sanity tests first (they're fast): + ```bash NEURON_SKIP_EFA_AFFINITY=1 pytest test/vllm_neuron/model/MODEL_NAME/bf16/e2e/test_logits.py -m "offline_serving" -v --timeout=3600 ``` 4. If sanity tests pass, run the online serving test at a small seq_len: + ```bash NEURON_SKIP_EFA_AFFINITY=1 pytest test/vllm_neuron/model/MODEL_NAME/bf16/e2e/test_logits.py -m "online_serving and seq256" -v --timeout=7200 ``` **Interpreting results:** The logit validation compares Neuron output logits against a HuggingFace CPU reference at various top-k levels (5, 50, 1000, all). Results are reported as sigma values: + - **< 3 sigma**: Excellent — within normal numerical noise - **3-5 sigma**: Acceptable — minor numerical differences, usually from bf16 quantization - **5-10 sigma**: Investigate — may indicate a real issue but could also be model-specific @@ -468,6 +497,7 @@ TP_SIZE: {recommended TP from Step 1e} ``` The equivalence skill runs the 8-stage pipeline: + - Stage 0: Build model trees, component mapping - Stage 2: Component-level R-ratio tests (per submodule) - Stages 3-4: Fault localization and debugging (if failures found) @@ -477,6 +507,7 @@ The equivalence skill runs the 8-stage pipeline: For vLLM-Neuron-specific behaviors (weight transpositions, forward signature, TP detection, KV cache shapes), the equivalence skill's `references/vllm-neuron-adaptation.md` is the authority — defer to it if anything conflicts. **Interpreting results:** + - All R < 1.2 and E2E passes → port is verified at component level - Component failures found → equivalence report includes patches showing what's wrong - Use the patches as a guide to fix the actual model code in `vllm_neuron/model/MODEL_NAME/model.py` @@ -486,6 +517,7 @@ For vLLM-Neuron-specific behaviors (weight transpositions, forward signature, TP ## Completion Print final status: + ``` ## Port Complete: MODEL_NAME (HF_MODEL_ID) diff --git a/skills/neuron-framework-autoport/SKILL.md b/skills/neuron-framework-autoport/SKILL.md index d8b9221..e7ba713 100644 --- a/skills/neuron-framework-autoport/SKILL.md +++ b/skills/neuron-framework-autoport/SKILL.md @@ -20,6 +20,7 @@ This document provides the agent direct instructions on how to port a model from ## Dry-run When the user specifies `dry-run`: + - **Skip** the "Resolve Dependencies" step - **Run** these commands to activate the venv and resolve source paths: ```bash @@ -44,11 +45,11 @@ Follow `references/setup_flow.md`. It handles venv validation, install consent, After success, retain the 3 resolved paths from the script output for use throughout the workflow: -| Variable | Description | -|---|---| -| `${NXDI_SRC}` | Path to NeuronX Distributed Inference source | -| `${NXD_SRC}` | Path to NeuronX Distributed source | -| `${TRANSFORMERS_SRC}` | Path to HuggingFace Transformers source | +| Variable | Description | +| --------------------- | -------------------------------------------- | +| `${NXDI_SRC}` | Path to NeuronX Distributed Inference source | +| `${NXD_SRC}` | Path to NeuronX Distributed source | +| `${TRANSFORMERS_SRC}` | Path to HuggingFace Transformers source | ### Read Project Guidelines @@ -58,14 +59,14 @@ READ `references/systemPrompts/systemPrompt.md` in this skill directory. It cont Extract these six required parameters from the user's request before starting: -| Parameter | Description | -|---|---| -| `ModelName` | The HuggingFace model class name (e.g., `ArceeForCausalLM`) | +| Parameter | Description | +| ------------------------------------ | --------------------------------------------------------------------------------------- | +| `ModelName` | The HuggingFace model class name (e.g., `ArceeForCausalLM`) | | `pathToModelImplementationDirectory` | Path to the model source directory (e.g., `transformers/src/transformers/models/arcee`) | -| `NameOfImplementationFile` | The modeling file name (e.g., `modeling_arcee.py`) | -| `NameOfConfigurationFile` | The configuration file name (e.g., `configuration_arcee.py`) | -| `huggingFaceModelID` | The HuggingFace model ID (e.g., `arcee-ai/AFM-4.5B-Base`) | -| `pathToModelWeightsDirectory` | Path to store/load model weights (e.g., `agent_artifacts/data`) | +| `NameOfImplementationFile` | The modeling file name (e.g., `modeling_arcee.py`) | +| `NameOfConfigurationFile` | The configuration file name (e.g., `configuration_arcee.py`) | +| `huggingFaceModelID` | The HuggingFace model ID (e.g., `arcee-ai/AFM-4.5B-Base`) | +| `pathToModelWeightsDirectory` | Path to store/load model weights (e.g., `agent_artifacts/data`) | If any required parameter is missing, prompt the user for it before starting the workflow. @@ -92,7 +93,8 @@ Please give me an architectural description of each of the existing supported mo Based on your understanding, including all of the existing available components in the Neuron SDK you have analyzed in directories NeuronxDistributed and NeuronxDistributedInference, please now analyze a CUDA specific implementation of {{ modelName }} in the project root sub-directory {{ pathToModelImplementationDirectory }} it contains the implementation of the model in file {{ NameOfImplementationFile }} and configuration of the model in {{ NameOfConfigurationFile }} and then create a version of this model that works based on the neuronx_distributed_inference framework in this repository. In the same directory will be a configuration file which will provide you the configuration you need. Please pay attention to the configurations, particularly the quantization, torch_dtype. Do a basic implementation of {{ modelName }} model using no sharding be thorough in your explanation, and ensure that the code produced is well documented and refers to the functions and files in the project root sub-directory {{ pathToModelImplementationDirectory }} directory where possible. #### Component and approach instructions -You will need to reuse all of the existing Neuron components and where you cant you need to flag them in comments. Before concluding a component cannot be reused, check if the base class supports override hooks or accepts None for optional parameters. When you port this model from huggingface please keep the names for each component of the model consistent with huggingface. If you can not please include a _u in the name to indicate that it does not have a 1:1 mapping. Implement it a component at a time and test the individual components before proceeding onto the next step, dont skip tests if distributed initialization or other features of the framework are required. For instance if there is a MLP, or attention component, implement the MLP component and test it first by doing a forward pass before going onto the Attention component. Always read the helper function's signature before calling it, and use only the parameters it accepts. Leverage the documents in the references/knowledge_base/ directory as a guide to ensure you avoid all the past mistakes. When you create the new code only leverage the NeuronxDistributed and NeuronxDistributedInference frameworks and pytorch, and do not use anything outside of the NeuronxDistributed and NeuronxDistributedInference. Do not modify any of the existing framework code. Do not pip install any additional packages. When you create the resulting ported implementation please place it in neuron_port sub-directory in the project. + +You will need to reuse all of the existing Neuron components and where you cant you need to flag them in comments. Before concluding a component cannot be reused, check if the base class supports override hooks or accepts None for optional parameters. When you port this model from huggingface please keep the names for each component of the model consistent with huggingface. If you can not please include a \_u in the name to indicate that it does not have a 1:1 mapping. Implement it a component at a time and test the individual components before proceeding onto the next step, dont skip tests if distributed initialization or other features of the framework are required. For instance if there is a MLP, or attention component, implement the MLP component and test it first by doing a forward pass before going onto the Attention component. Always read the helper function's signature before calling it, and use only the parameters it accepts. Leverage the documents in the references/knowledge_base/ directory as a guide to ensure you avoid all the past mistakes. When you create the new code only leverage the NeuronxDistributed and NeuronxDistributedInference frameworks and pytorch, and do not use anything outside of the NeuronxDistributed and NeuronxDistributedInference. Do not modify any of the existing framework code. Do not pip install any additional packages. When you create the resulting ported implementation please place it in neuron_port sub-directory in the project. ### Step 3: Compile @@ -119,6 +121,7 @@ Validate against the HuggingFace golden reference using the **Validation Tool**: **Success criteria: >= 95% greedy token match rate.** Exit code 0 means passed. **Iteration loop if validation fails (<95% match rate):** + 1. Fix code 2. Delete compiled model: `rm -rf agent_artifacts/data/compiled_model && rm -rf /var/tmp/neuron-compile-cache` 3. Re-compile diff --git a/skills/neuron-framework-autoport/references/knowledge_base/COMPREHENSIVE_LLAMA3_NEURONX_GUIDE.md b/skills/neuron-framework-autoport/references/knowledge_base/COMPREHENSIVE_LLAMA3_NEURONX_GUIDE.md index a06f989..ab7ebae 100644 --- a/skills/neuron-framework-autoport/references/knowledge_base/COMPREHENSIVE_LLAMA3_NEURONX_GUIDE.md +++ b/skills/neuron-framework-autoport/references/knowledge_base/COMPREHENSIVE_LLAMA3_NEURONX_GUIDE.md @@ -34,6 +34,7 @@ This comprehensive guide documents the complete implementation of Meta's Llama3 **Primary Goal**: Port Meta's Llama3.2-1B model from the original CUDA implementation to run efficiently on AWS Neuron hardware using the NeuronxDistributed framework. **Key Requirements**: + - Maintain architectural fidelity to the original Llama3 implementation - Support both original Llama3 checkpoint format and HuggingFace format - Follow established NeuronxDistributed framework patterns @@ -43,6 +44,7 @@ This comprehensive guide documents the complete implementation of Meta's Llama3 ### Final Status: PRODUCTION READY ✅ The model successfully: + 1. ✅ **Compiles** for NeuronX hardware without errors 2. ✅ **Loads** in the compiled environment with proper initialization 3. ✅ **Processes** GQA conversion correctly for single-device deployment @@ -56,21 +58,23 @@ The model successfully: ### Original Llama3 Architecture (Llama3.2-1B) #### Core Model Parameters + ```json { - "dim": 2048, // Hidden size - "n_heads": 32, // Query attention heads - "n_kv_heads": 8, // Key-value heads (4:1 GQA ratio) - "n_layers": 16, // Transformer layers - "vocab_size": 128256, // Vocabulary size - "ffn_dim_multiplier": 1.5, // MLP dimension multiplier - "rope_theta": 500000.0, // RoPE base frequency - "use_scaled_rope": true, // Scaled RoPE for long context - "norm_eps": 1e-05 // RMSNorm epsilon + "dim": 2048, // Hidden size + "n_heads": 32, // Query attention heads + "n_kv_heads": 8, // Key-value heads (4:1 GQA ratio) + "n_layers": 16, // Transformer layers + "vocab_size": 128256, // Vocabulary size + "ffn_dim_multiplier": 1.5, // MLP dimension multiplier + "rope_theta": 500000.0, // RoPE base frequency + "use_scaled_rope": true, // Scaled RoPE for long context + "norm_eps": 1e-5 // RMSNorm epsilon } ``` #### Key Architectural Features + 1. **Grouped-Query Attention (GQA)**: 32 query heads share 8 key-value heads (4:1 ratio) 2. **Rotary Position Embeddings (RoPE)**: With θ=500,000 and scaling support 3. **SwiGLU Activation**: `w2(silu(w1(x)) * w3(x))` in MLP layers @@ -78,7 +82,8 @@ The model successfully: 5. **KV Caching**: Explicit key-value caching for autoregressive generation #### Original Implementation Structure -```python + +````python # From source/llama3/llama/model.py class Transformer: - tok_embeddings: VocabParallelEmbedding @@ -97,11 +102,13 @@ onxDistributed Framework Architecture #### Framework Structure The NeuronxDistributed framework follows a specific pattern for model implementations: -``` -src/neuronx_distributed_inference/models/{model_name}/ -├── __init__.py -└── modeling_{model_name}.py -``` +```` + +src/neuronx*distributed_inference/models/{model_name}/ +├── **init**.py +└── modeling*{model_name}.py + +```` #### Base Classes 1. **NeuronApplicationBase**: Root application class requiring `model_path` parameter @@ -167,7 +174,7 @@ class NeuronLlama3Model(NeuronBaseModel): """Initialize model components""" # Initialize embed_tokens, layers, norm, lm_head pass -``` +```` --- @@ -176,15 +183,16 @@ class NeuronLlama3Model(NeuronBaseModel): ### Multi-Format Configuration Support #### Dual Format Loading + ```python @classmethod def from_pretrained(cls, model_path: str, **kwargs): # Expand user home directory if needed model_path = os.path.expanduser(model_path) - + params_path = os.path.join(model_path, "params.json") config_path = os.path.join(model_path, "config.json") - + if os.path.exists(params_path): # Load original Llama3 format with open(params_path, 'r') as f: @@ -200,7 +208,9 @@ def from_pretrained(cls, model_path: str, **kwargs): ``` #### Parameter Name Mapping + **Original Format** → **Framework Format**: + - `dim` → `hidden_size` - `n_layers` → `num_hidden_layers` - `n_heads` → `num_attention_heads` @@ -210,6 +220,7 @@ def from_pretrained(cls, model_path: str, **kwargs): - `rope_theta` → `rope_theta` #### Configuration Conversion + ```python @classmethod def from_original_params(cls, params): @@ -230,23 +241,24 @@ def from_original_params(cls, params): ### Intermediate Size Calculation #### Original Llama3 Logic Replication + ```python def calculate_intermediate_size(params): """Calculate intermediate_size from ffn_dim_multiplier like original Llama3""" hidden_dim = params['dim'] multiple_of = params.get('multiple_of', 256) ffn_dim_multiplier = params.get('ffn_dim_multiplier') - + # Base calculation: 2/3 of 4 * hidden_dim hidden_dim = int(2 * hidden_dim / 3) - + # Apply multiplier if specified if ffn_dim_multiplier is not None: hidden_dim = int(ffn_dim_multiplier * hidden_dim) - + # Round to nearest multiple hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of) - + return hidden_dim ``` @@ -259,11 +271,13 @@ def calculate_intermediate_size(params): #### 1. Base Class Integration Issues **Error**: Missing Required Methods + ``` AttributeError: 'NeuronLlama3Model' object has no attribute 'setup_attr_for_model' ``` **Fix**: Added required methods to the `NeuronLlama3Model` class: + ```python def setup_attr_for_model(self): """Setup attributes for model initialization""" @@ -279,18 +293,20 @@ def init_model(self): #### 2. Constructor Signature Mismatch **Error**: Incompatible Constructor Parameters + ``` TypeError: __init__() got unexpected keyword arguments ``` **Fix**: Modified the constructor to match the base class signature: + ```python def __init__(self, config): if isinstance(config, str): config = Llama3InferenceConfig.from_pretrained(config) elif isinstance(config, dict): config = Llama3InferenceConfig(**config) - + super().__init__(config) self.setup_attr_for_model() ``` @@ -298,6 +314,7 @@ def __init__(self, config): #### 3. Forward Method Signature Conflicts **Error**: Framework Forward Method Conflict + ``` ERROR: You cannot specify both input_ids and inputs_embeds at the same time ``` @@ -305,6 +322,7 @@ ERROR: You cannot specify both input_ids and inputs_embeds at the same time **Root Cause**: The NeuronxDistributedInference framework expects models to NOT implement their own forward method. **Solution**: Remove the custom forward method entirely: + ```python # WRONG - Don't implement forward method class NeuronLlama3Model(NeuronBaseModel): @@ -315,21 +333,23 @@ class NeuronLlama3Model(NeuronBaseModel): class NeuronLlama3Model(NeuronBaseModel): def setup_attr_for_model(self, config): pass - + def init_model(self, config): pass - + # No forward method - base class handles this ✅ ``` #### 4. Layer Return Format Mismatches **Error**: Tuple Unpacking Mismatch + ``` ERROR: too many values to unpack (expected 3) ``` **Solution**: Match the framework's expected return format: + ```python # CORRECT - Unpack 4 values as framework expects hidden_states, present_key_value, cos_cache, sin_cache = self.self_attn( @@ -349,12 +369,14 @@ return outputs #### 5. Checkpoint Conversion Issues **Problem**: Parameter Count Mismatch + ``` Original checkpoint: 147 parameters Converted checkpoint: 164 parameters ``` **Analysis**: The increase was due to: + - Framework splitting combined weight matrices - Adding metadata and configuration parameters - Tensor reshaping for distributed training compatibility @@ -364,11 +386,13 @@ Converted checkpoint: 164 parameters #### 6. Import and Module Structure Issues **Error**: Module Import Problems + ``` ImportError: cannot import name 'NeuronLlama3Model' from 'neuronx_llama3' ``` **Fix**: Created proper `__init__.py` files with correct imports: + ```python # src/neuronx_llama3/__init__.py from .modeling_llama3 import ( @@ -379,7 +403,7 @@ from .modeling_llama3 import ( __all__ = [ "NeuronLlama3Model", - "Llama3InferenceConfig", + "Llama3InferenceConfig", "NeuronLlama3ForCausalLM" ] ``` @@ -407,12 +431,13 @@ neuron_config = NeuronConfig( ### Compilation Success Indicators The successful compilation showed: + ``` ============================================================ COMPILATION COMPLETED SUCCESSFULLY ✅ ============================================================ - Both context encoding and token generation models compiled -- GQA correctly converted to MHA for single-device deployment +- GQA correctly converted to MHA for single-device deployment - All 164 parameters loaded and converted - Compilation time: ~138 seconds - Target hardware: AWS Trn1 (Neuron optimized) @@ -431,9 +456,11 @@ INFO:Neuron:Finished Compilation for all HLOs in 49.70 seconds ### Weight Conversion System #### Multi-Format Support + The conversion system handles three checkpoint formats: 1. **Original Llama3 Format** (`consolidated.00.pth`): + ``` tok_embeddings.weight → model.embed_tokens.weight layers.{i}.attention.wq.weight → model.layers.{i}.self_attn.q_proj.weight @@ -441,6 +468,7 @@ The conversion system handles three checkpoint formats: ``` 2. **HuggingFace SafeTensors Format**: + ``` model.embed_tokens.weight → model.embed_tokens.weight model.layers.{i}.self_attn.q_proj.weight → model.layers.{i}.self_attn.qkv_proj.q_proj.weight @@ -449,24 +477,25 @@ The conversion system handles three checkpoint formats: 3. **Neuron Format**: Final format expected by the NeuronxDistributed model #### Tensor Parallelism Metadata Addition + ```python @staticmethod def convert_hf_to_neuron_state_dict(state_dict: dict, config: InferenceConfig) -> dict: neuron_config = config.neuron_config - + # Add rank utilities for tensor parallel support if neuron_config.vocab_parallel: state_dict["embed_tokens.rank_util.rank"] = torch.arange( 0, neuron_config.local_ranks_size ) - + num_layers = config.num_hidden_layers tp_degree = neuron_config.tp_degree for i in range(num_layers): state_dict[f"layers.{i}.self_attn.rank_util.rank"] = torch.arange( 0, tp_degree, dtype=torch.int32 ) - + return state_dict ``` @@ -477,6 +506,7 @@ def convert_hf_to_neuron_state_dict(state_dict: dict, config: InferenceConfig) - ### Model Loading Architecture #### Checkpoint Loader Override + ```python class NeuronLlama3ForCausalLM(NeuronLlamaForCausalLM): def checkpoint_loader_fn(self, mmap: bool = False): @@ -498,6 +528,7 @@ class NeuronLlama3ForCausalLM(NeuronLlamaForCausalLM): ``` ### Inference Script Features + - **Robust tokenizer loading**: Falls back to original checkpoint directory - **Dummy token testing**: Allows model validation without tokenizer - **Progressive testing**: Tests forward pass before attempting generation @@ -506,6 +537,7 @@ class NeuronLlama3ForCausalLM(NeuronLlamaForCausalLM): ### Working Inference Implementation #### Model Loading and Initialization + ```python # Load model model = NeuronLlama3ForCausalLM(model_path) @@ -522,6 +554,7 @@ generated_ids = input_ids.clone() ``` #### Generation Loop Implementation + ```python # Generation loop with torch.no_grad(): @@ -529,18 +562,18 @@ with torch.no_grad(): # Create position_ids seq_len = generated_ids.shape[1] position_ids = torch.arange(seq_len).unsqueeze(0) - + # Forward pass outputs = model(generated_ids, position_ids=position_ids) logits = outputs.logits if hasattr(outputs, 'logits') else outputs[0] - + # Get next token (greedy decoding) next_token_logits = logits[:, -1, :] next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True) - + # Append to sequence generated_ids = torch.cat([generated_ids, next_token], dim=-1) - + # Check for EOS if tokenizer.eos_token_id and next_token.item() == tokenizer.eos_token_id: break @@ -552,19 +585,22 @@ output_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True) ### Final Status: SUCCESS ✅ #### What Works + 1. **Model Loading**: Successfully loads compiled TorchScript model 2. **Weight Initialization**: Properly initializes with weights from original checkpoint 3. **Forward Pass**: Model forward pass executes successfully 4. **Neuron Integration**: All Neuron-specific optimizations active (GQA conversion, etc.) #### Test Results + ```bash python run_inference.py --model_path ./llama3_compiled --prompt "Hello" --max_new_tokens 3 ``` **Output Indicators of Success**: + - ✅ Weight sharding completed: `INFO:Neuron:Sharding weights on load...` -- ✅ GQA conversion working: All 33 layers processed correctly +- ✅ GQA conversion working: All 33 layers processed correctly - ✅ Weights loaded: `INFO:Neuron:Loading weights from original checkpoint` - ✅ Model warming up: `INFO:Neuron:Warming up the model.` - ✅ Forward pass successful: Model inference working @@ -577,7 +613,8 @@ python run_inference.py --model_path ./llama3_compiled --prompt "Hello" --max_ne #### 1. Model Initialization Errors -**Error**: +**Error**: + ``` This model is not initialized, please call traced_model.nxd_model.initialize(sharded_checkpoint) or traced_model.nxd_model.initialize_with_saved_weights() ``` @@ -589,6 +626,7 @@ This model is not initialized, please call traced_model.nxd_model.initialize(sha #### 2. Tokenizer Loading Issues **Error**: Multiple tokenizer-related errors: + - Missing SentencePiece library for `LlamaTokenizer` - Compiled directory only contained `tokenizer.model` (SentencePiece format) - HuggingFace tokenizer format not available in compiled directory @@ -596,18 +634,23 @@ This model is not initialized, please call traced_model.nxd_model.initialize(sha **Solution**: Modified tokenizer loading to fall back to original checkpoint directory and implemented dummy token testing for model validation. ##### 2.1 Model Type Recognition Error -**Error**: + +**Error**: + ``` The checkpoint you are trying to load has model type `llama3_neuron` but Transformers does not recognize this architecture. ``` **Solution**: Fixed the config.json: + ```python config['model_type'] = 'llama' # Changed from 'llama3_neuron' ``` ##### 2.2 Missing Tokenizer Files + **Solution**: Created minimal tokenizer configuration files: + ```python # tokenizer_config.json { @@ -627,12 +670,15 @@ config['model_type'] = 'llama' # Changed from 'llama3_neuron' ``` #### 3. Forward Pass Parameter Issues -**Error**: + +**Error**: + ``` AssertionError: need to call forward with position_ids if attention_mask is not provided ``` **Solution**: Added proper `position_ids` generation in the inference loop: + ```python seq_len = generated_ids.shape[1] position_ids = torch.arange(seq_len).unsqueeze(0) @@ -641,7 +687,8 @@ outputs = model(generated_ids, position_ids=position_ids) #### 4. Generation Method Issues -**Error**: +**Error**: + ``` 'NeuronLlama3ForCausalLM' object has no attribute 'generate' ``` @@ -653,6 +700,7 @@ outputs = model(generated_ids, position_ids=position_ids) ### Debugging Workflow #### Progressive Testing Approach + 1. **Configuration Check**: Verify config files load correctly 2. **Checkpoint Files Check**: Validate all required files exist 3. **Weight Loading Check**: Ensure parameters load successfully @@ -660,12 +708,14 @@ outputs = model(generated_ids, position_ids=position_ids) 5. **Compiled Model Loading**: Verify model loads in compiled environment #### Debug Utilities Created + - **debug_keys.py**: Parameter structure debugging - **debug_mistral_keys.py**: Mistral model comparison - **debug_qwen_keys.py**: Qwen model comparison - **debug_t5_keys.py**: T5 model comparison ### File Structure Requirements + ``` neuronx_llama3/ ├── llama3_compiled/ # Compiled model with fixed tokenizer files @@ -689,23 +739,27 @@ neuronx_llama3/ All models in the NeuronxDistributed framework share several key architectural components: #### Parallelization Strategies + - **Tensor Parallelism (TP)**: Splits model parameters across multiple devices - **Sequence Parallelism (SP)**: Distributes sequence processing across devices - **Context Parallelism (CP)**: Optimizes attention computation for long sequences - **Data Parallelism (DP)**: Processes different batches on different devices #### Attention Mechanisms + - **Multi-head Attention**: All models use some form of multi-head attention - **Group Query Attention (GQA)**: Optimizes KV cache usage by sharing key-value heads - **Rotary Position Embeddings (RoPE)**: Used for positional encoding in most models #### Normalization + - **RMSNorm**: Most models use RMSNorm instead of LayerNorm for better performance - **CustomRMSNorm**: Optimized implementation for Neuron hardware ### Model-Specific Architectures #### 1. Mistral + - **Attention**: Uses Grouped-Query Attention (GQA) - **Normalization**: RMSNorm for input and post-attention normalization - **Position Encoding**: Rotary Position Embeddings (RoPE) @@ -713,11 +767,13 @@ All models in the NeuronxDistributed framework share several key architectural c - **Special Features**: Sliding window attention mechanism #### 2. Mixtral (MoE) + - **Attention**: Same GQA mechanism as Mistral - **MLP Layer**: Replaced with a Mixture-of-Experts layer - **MoE Architecture**: 8 experts per layer, Top-k routing (k=2) #### 3. Qwen3 + - **Attention**: Uses GQA with Q-K normalization - **Normalization**: RMSNorm for all normalization layers - **Position Encoding**: Rotary Position Embeddings (RoPE) @@ -725,21 +781,25 @@ All models in the NeuronxDistributed framework share several key architectural c - **Special Features**: Q-K normalization applies RMSNorm to query and key vectors #### 4. T5 + - **Architecture**: Encoder-decoder transformer - **Special Features**: Supports text-to-text tasks #### 5. CLIP + - **Architecture**: Dual encoder (vision and text) - **Special Features**: Supports image-text tasks ### Key Architectural Differences #### Attention Mechanisms + 1. **Standard Multi-head Attention**: Used in basic transformer models 2. **Grouped-Query Attention (GQA)**: Used in Mistral, Mixtral, Qwen3 - reduces KV cache size 3. **Q-K Normalization**: Used in Qwen3 - applies normalization to query and key vectors #### MLP Architectures + 1. **Standard MLP**: Used in basic transformer models 2. **SwiGLU Activation**: Used in LLaMA-style models including Mistral 3. **Mixture-of-Experts**: Used in Mixtral and Qwen3-MoE @@ -747,6 +807,7 @@ All models in the NeuronxDistributed framework share several key architectural c ### Llama3 Architecture Implementation #### Attention Implementation + ```python class NeuronLlama3Attention(NeuronAttentionBase): def __init__(self, config): @@ -755,7 +816,7 @@ class NeuronLlama3Attention(NeuronAttentionBase): max_position_embeddings=config.max_position_embeddings, base=config.rope_theta, # 500000.0 ) - + super().__init__( config=config, hidden_size=config.hidden_size, # 2048 @@ -767,6 +828,7 @@ class NeuronLlama3Attention(NeuronAttentionBase): ``` #### MLP Implementation + ```python class NeuronLlama3MLP(nn.Module): def forward(self, x): @@ -783,24 +845,28 @@ class NeuronLlama3MLP(nn.Module): ## Best Practices ### Framework Understanding + 1. **Base Class Requirements**: Understanding constructor parameters is crucial 2. **Distributed Initialization**: Proper sequence is essential for model creation 3. **State Dict Conventions**: Each framework has specific naming expectations 4. **Parallelization Integration**: Models must be designed with distribution in mind ### Implementation Strategy + 1. **Start Simple**: Begin with minimal configurations and build up 2. **Follow Patterns**: Existing models provide excellent templates 3. **Incremental Development**: Add features one at a time 4. **Comprehensive Testing**: Create debug utilities early in the process ### Common Pitfalls + 1. **State Dict Mismatches**: Most common source of loading errors 2. **Missing Initialization**: Distributed environment must be set up first 3. **Parameter Naming**: Inconsistent naming across frameworks causes issues 4. **Dependency Management**: Model interdependencies can cause import failures ### Development Workflow Recommendations + 1. **Start with Configuration**: Get config loading working first 2. **Test Checkpoint Loading**: Verify weight loading before model creation 3. **Implement Components**: Build attention, MLP, and model classes incrementally @@ -828,12 +894,14 @@ When implementing a new model for NeuronxDistributed, ensure: ## Performance Optimization ### Current Performance Baseline + - **Compilation Time**: ~142 seconds total - **Model Size**: 9.4MB compiled artifacts - **Memory Usage**: Optimized for single device deployment - **Sequence Length**: 128 tokens (expandable) ### Optimization Opportunities + 1. **Tensor Parallelism**: Enable multi-device distribution 2. **Longer Sequences**: Increase context length with optimizations 3. **Mixed Precision**: Use bfloat16 for better performance @@ -841,6 +909,7 @@ When implementing a new model for NeuronxDistributed, ensure: 5. **On-Device Sampling**: Reduce host-device communication ### Inference Optimizations + 1. **Flash Attention**: Optimized attention computation for faster inference 2. **KV Cache Management**: Efficient handling of key-value pairs for autoregressive generation 3. **On-device Sampling**: Optimized token generation directly on device @@ -850,6 +919,7 @@ When implementing a new model for NeuronxDistributed, ensure: ### GQA (Grouped Query Attention) Handling The framework automatically handles GQA conversion: + - For TP=1 (single device): Converts GQA to MHA by replicating key-value heads - For TP>1: Maintains GQA structure across tensor parallel ranks @@ -866,9 +936,10 @@ This is standard behavior and provides memory efficiency while maintaining perfo ✅ **Weight Conversion**: Proper parameter name mapping and metadata addition ✅ **Model Compilation**: Successful compilation with GQA handling ✅ **Model Loading**: Compiled artifacts load correctly -✅ **Architecture Integrity**: All Llama3 components properly implemented +✅ **Architecture Integrity**: All Llama3 components properly implemented ### Key Validation Insights + - **Parameter Shapes**: Verified correct GQA dimensions (q_proj: 2048x2048, k_proj: 512x2048) - **Weight Loading**: All 164 parameters loaded with expected tensor parallelism metadata - **Compilation Warnings**: GQA automatically converted to MHA for TP=1 (expected behavior) @@ -876,17 +947,20 @@ This is standard behavior and provides memory efficiency while maintaining perfo ### Testing Results #### Final Test Results + ```bash python simple_inference.py --model_path ./llama3_compiled --prompt "Hello, how are you?" --max_new_tokens 5 ``` **Output**: + ``` Prompt: Hello, how are you? Generated: Hello, how are you? I am I am I ``` #### Performance Indicators + - ✅ Model loads successfully with proper weight sharding - ✅ Tokenizer loads and processes input correctly - ✅ Forward pass executes without errors @@ -907,6 +981,7 @@ model.load_state_dict(mapped_checkpoint) ``` **Key Findings**: + - Both Neuron and CPU versions produce identical output patterns - Parameter name mapping was crucial: Neuron checkpoint uses `layers.X.*` while HuggingFace expects `model.layers.X.*` - The simple generation loop approach works consistently across both implementations @@ -917,15 +992,17 @@ model.load_state_dict(mapped_checkpoint) ## Utility Scripts and Tools ### Complete Toolchain + 1. **convert_checkpoint.py**: Multi-format checkpoint conversion 2. **compile_model.py**: Model compilation for Neuron hardware 3. **run_inference.py**: Inference execution 4. **example_chat.py**: Interactive chat interface 5. **test_model.py**: Model testing without compilation 6. **run_pipeline.sh**: Complete pipeline automation -7. **debug_*.py**: Various debugging utilities +7. **debug\_\*.py**: Various debugging utilities ### Configuration Management + - **Minimal Stable Settings**: Following best practices from MODEL_IMPLEMENTATION_GUIDE.md - **Progressive Optimization**: Start simple, add optimizations incrementally - **Comprehensive Error Handling**: Detailed logging and error reporting @@ -933,12 +1010,14 @@ model.load_state_dict(mapped_checkpoint) ### Usage Instructions #### Basic Inference + ```bash cd neuronx_llama3 python run_inference.py --model_path ./llama3_compiled --prompt "Your prompt here" --max_new_tokens 50 ``` #### Advanced Options + ```bash python run_inference.py \ --model_path ./llama3_compiled \ @@ -957,6 +1036,7 @@ python run_inference.py \ This comprehensive implementation demonstrates the successful porting of Meta's Llama3 model to the AWS NeuronxDistributed framework. The project achieved: ### Key Achievements + - **Complete Architecture Port**: Faithful implementation of all Llama3 components - **Multi-format Support**: Handles both original and HuggingFace checkpoint formats - **Hardware Optimization**: Successfully compiles and optimizes for Neuron hardware @@ -965,20 +1045,24 @@ This comprehensive implementation demonstrates the successful porting of Meta's - **Extensible Design**: Easy to add optimizations and new features ### Technical Achievements + #### Architecture Fidelity: **100%** ✅ + - **GQA**: 32 query heads, 8 key-value heads (4:1 ratio) ✅ -- **RoPE**: θ=500,000 with scaling support ✅ -- **SwiGLU**: w2(silu(w1(x)) * w3(x)) activation ✅ +- **RoPE**: θ=500,000 with scaling support ✅ +- **SwiGLU**: w2(silu(w1(x)) \* w3(x)) activation ✅ - **RMSNorm**: ε=1e-05 layer normalization ✅ - **Parameter Count**: 1B parameters (Llama3.2-1B) ✅ #### Framework Integration: **100%** ✅ + - **Base Classes**: Proper `NeuronBaseModel` and `NeuronBaseForCausalLM` inheritance ✅ - **Method Implementation**: All required framework methods implemented ✅ - **Return Formats**: Consistent with other framework models (Qwen3, Mistral) ✅ - **Configuration**: Dual-format support with automatic parameter mapping ✅ #### Performance Optimization: **100%** ✅ + - **Hardware Target**: AWS Trn1 instances ✅ - **Memory Efficiency**: GQA reduces KV cache requirements ✅ - **Compilation**: Both context encoding and token generation models ✅ @@ -991,6 +1075,7 @@ This implementation demonstrates the feasibility and methodology for porting com The successful compilation and validation prove that the NeuronxDistributed framework can effectively support modern transformer architectures while providing the hardware optimizations necessary for efficient inference on AWS Neuron hardware. ### Next Steps + 1. **Complete Generation Interface**: Implement proper text generation API 2. **Performance Optimization**: Add tensor parallelism and longer sequences 3. **Advanced Features**: On-device sampling, flash attention, quantization @@ -999,11 +1084,13 @@ The successful compilation and validation prove that the NeuronxDistributed fram ### Files Created #### Core Implementation + - `neuronx_llama3/src/neuronx_llama3/modeling_llama3.py`: Complete model implementation - `neuronx_llama3/src/neuronx_llama3/__init__.py`: Module exports - `neuronx_llama3/setup.py`: Package configuration #### Utility Scripts + - `neuronx_llama3/convert_checkpoint.py`: Multi-format checkpoint conversion - `neuronx_llama3/compile_model.py`: Model compilation for Neuron - `neuronx_llama3/run_inference.py`: Inference execution @@ -1012,6 +1099,7 @@ The successful compilation and validation prove that the NeuronxDistributed fram - `neuronx_llama3/run_pipeline.sh`: Complete pipeline automation #### Documentation + - `neuronx_llama3/README.md`: Usage instructions and examples - `neuronx_llama3/IMPLEMENTATION_DETAILS.md`: Detailed implementation guide - `docs/`: Comprehensive documentation and analysis @@ -1024,4 +1112,4 @@ This comprehensive implementation demonstrates the complexity and depth required **Framework**: NeuronxDistributedInference **Target Hardware**: AWS Trn1 **Model**: Llama3.2-1B with GQA -**Status**: **COMPLETE SUCCESS - READY FOR PRODUCTION** 🚀 \ No newline at end of file +**Status**: **COMPLETE SUCCESS - READY FOR PRODUCTION** 🚀 diff --git a/skills/neuron-framework-autoport/references/knowledge_base/Category1_Porting_Config_Compilation_Issues.md b/skills/neuron-framework-autoport/references/knowledge_base/Category1_Porting_Config_Compilation_Issues.md index b2cbaa9..4c1a4bc 100644 --- a/skills/neuron-framework-autoport/references/knowledge_base/Category1_Porting_Config_Compilation_Issues.md +++ b/skills/neuron-framework-autoport/references/knowledge_base/Category1_Porting_Config_Compilation_Issues.md @@ -11,6 +11,7 @@ This document summarizes all issues, solutions, and learnings related to porting ## Document Organization ### Source Documents Analyzed: + - MAIN_comprehensive_generic_moe_port_analysis.md - genericmoe_complete_success_summary.md - genericmoe_compilation_solutions_guide.md @@ -30,6 +31,7 @@ This document summarizes all issues, solutions, and learnings related to porting ### Problem Description **Symptom:** + ``` HloVerifier failed: Expert routing patterns not recognized Compilation blocked at verification stage @@ -37,6 +39,7 @@ Error: Shape mismatches during weight layout optimization ``` **Impact:** + - 0% compilation success rate - Complete blocker for model deployment - Affected all MoE compilation attempts @@ -51,6 +54,7 @@ MoE models use **dynamic expert routing** that creates conditional computation g 4. **Complex Patterns**: All-to-all communication for expert routing flagged as invalid **Why It Happened:** + - HLO verifier designed for static computation graphs - MoE dynamic routing patterns not in verifier's supported pattern list - Compiler toolchain limitation, not model architecture issue @@ -60,11 +64,13 @@ MoE models use **dynamic expert routing** that creates conditional computation g **Approach**: Selective verifier disabling with comprehensive post-compilation validation **Location**: + - File: `NeuronxDistributed/src/neuronx_distributed/trace/model_builder.py` - Line: 2064 - Change: `--verify-hlo=false` (was `--verify-hlo=true`) **Implementation:** + ```python # In compilation script os.environ['NEURON_CC_FLAGS'] = '--disable-hlo-verifier' @@ -79,6 +85,7 @@ compiler_args = [ ``` **Post-Compilation Validation:** + ```python def validate_compiled_model(compiled_model_path): """Comprehensive validation to replace HloVerifier""" @@ -113,6 +120,7 @@ def validate_compiled_model(compiled_model_path): ### Results **Before Fix:** + ``` Compilation Status: FAILED Error: HloVerifier failed on expert routing patterns @@ -121,6 +129,7 @@ Success Rate: 0% ``` **After Fix:** + ``` Compilation Status: SUCCESS Compilation Time: 45 minutes (tiny model), ~5 minutes (full model with cache) @@ -147,14 +156,15 @@ Inference: Fully functional **Options Available:** -| Framework | Models Using | Expert Parallelism | Production Ready | Complexity | -|-----------|--------------|-------------------|------------------|------------| -| MoE v1 | Mixtral, DBRX | Manual | Limited | High | -| MoE v2 | Qwen3, DeepSeek | Automatic | ✅ Yes | Low | +| Framework | Models Using | Expert Parallelism | Production Ready | Complexity | +| --------- | --------------- | ------------------ | ---------------- | ---------- | +| MoE v1 | Mixtral, DBRX | Manual | Limited | High | +| MoE v2 | Qwen3, DeepSeek | Automatic | ✅ Yes | Low | ### Solution: MoE v2 Framework Selection **Rationale:** + 1. ✅ **Built-in expert parallelism support** - automatic process group management 2. ✅ **Advanced optimization kernels** - better performance 3. ✅ **Automatic process group management** - less manual configuration @@ -162,6 +172,7 @@ Inference: Fully functional 5. ✅ **Production-ready** - used in Qwen3 and DeepSeek deployments **Implementation:** + ```python # Use MoE v2 framework from neuronx_distributed_inference.modules.moe_v2 import initialize_moe_module @@ -179,6 +190,7 @@ class GenericMoEDecoderLayer(nn.Module): ``` **Configuration:** + ```python class GenericMoEInferenceConfig(InferenceConfig): """Complete configuration with HuggingFace compatibility""" @@ -214,11 +226,13 @@ class GenericMoEInferenceConfig(InferenceConfig): ### Problem Description **Symptom:** + ```python AttributeError: 'builtin_function_or_method' object has no attribute 'is_initialized' ``` **Impact:** + - Distributed initialization failed - Process groups couldn't be created - Model instantiation blocked @@ -232,6 +246,7 @@ Incomplete `InferenceConfig` inheritance missing **required abstract methods**: 3. `add_derived_config()` - missing 20+ framework-expected attributes **Why This Happened:** + - Base class has abstract methods that must be overridden - Framework expects specific configuration attributes for MoE models - Process group initialization depends on proper config @@ -318,6 +333,7 @@ class GenericMoEInferenceConfig(InferenceConfig): ### Results **Before Fix:** + ``` Error: AttributeError during initialization Distributed initialization: FAILED @@ -326,6 +342,7 @@ Model instantiation: BLOCKED ``` **After Fix:** + ``` Configuration validation: PASSED Distributed initialization: SUCCESS @@ -349,6 +366,7 @@ MoE framework integration: COMPLETE ### Problem Description **Symptom:** + ``` Missing keys: 547 Unexpected keys: 517 @@ -406,6 +424,7 @@ def convert_generic_moe_hf_to_neuron_state_dict(hf_state_dict, config): ``` **Validation:** + ```python # Verify key mappings print("HF keys sample:") @@ -418,6 +437,7 @@ print("\nMapping: ✅ Correct") ### Results **Before Fix:** + ``` Missing keys: 547 Unexpected keys: 517 @@ -426,6 +446,7 @@ Forward pass: FAILED ``` **After Fix:** + ``` Missing keys: 96 (only MoE weights, fixed separately) Unexpected keys: 1 (harmless rank tensor) @@ -440,6 +461,7 @@ Forward pass: SUCCESS (after MoE fix) ### Problem Description **Symptom:** + ```python RuntimeError: shape mismatch: value (1, 2048, 32, 128) vs expected (1, 2048, 128) ``` @@ -515,6 +537,7 @@ class GenericMoEAttention(NeuronAttentionBase): ### Results **Before Migration:** + ``` Shape errors: Frequent Performance: Suboptimal @@ -523,6 +546,7 @@ Optimization: None ``` **After Migration:** + ``` Shape errors: None Performance: Optimized for Neuron @@ -538,6 +562,7 @@ Correctness: 100% ### Problem Description **Original Model Constraints:** + ``` Parameters: 29B total (16 experts × ~1.8B each) Compilation Memory: ~28GB (exceeded 24GB limit) @@ -596,6 +621,7 @@ CONFIGURATIONS = { ``` **Weight Mapping Strategy:** + ```python def create_small_model_from_full(full_model_path, config_size): """Extract weights from full model for smaller variant""" @@ -649,15 +675,16 @@ python recompile_generic_moe_final.py --tp_degree 16 **Memory Reduction:** -| Configuration | Parameters | Memory/Rank | Compilation Time | Success | -|---------------|------------|-------------|------------------|---------| -| Original (unopt) | 29B | 28GB | N/A | ❌ OOM | -| Tiny | 1B | 0.5GB | 10 min | ✅ | -| Small | 3B | 1GB | 20 min | ✅ | -| Medium | 8B | 3GB | 30 min | ✅ | -| Full (optimized) | 29B | <16GB | 5 min | ✅ | +| Configuration | Parameters | Memory/Rank | Compilation Time | Success | +| ---------------- | ---------- | ----------- | ---------------- | ------- | +| Original (unopt) | 29B | 28GB | N/A | ❌ OOM | +| Tiny | 1B | 0.5GB | 10 min | ✅ | +| Small | 3B | 1GB | 20 min | ✅ | +| Medium | 8B | 3GB | 30 min | ✅ | +| Full (optimized) | 29B | <16GB | 5 min | ✅ | **Benefits of Progressive Approach:** + - ✅ Validated compilation pipeline early - ✅ Identified issues on manageable model sizes - ✅ Tested expert sharding systematically @@ -791,6 +818,7 @@ Memory Usage: **Lesson**: Always use MoE v2 framework for new MoE implementations **Why**: + - Built-in expert parallelism support - Automatic process group management - Better optimization kernels @@ -804,12 +832,14 @@ Memory Usage: **Lesson**: HLO verifier can have bugs with complex models; disabling with validation is acceptable **When to Apply**: + - Dynamic computation graphs (MoE routing) - Conditional operations - Complex communication patterns - New or cutting-edge architectures **Best Practice**: + - Disable verifier with explicit flag - Implement comprehensive post-compilation validation - Document the workaround clearly @@ -820,6 +850,7 @@ Memory Usage: **Lesson**: Proper InferenceConfig inheritance is critical for distributed initialization **Required Methods**: + ```python def get_required_attributes(self) -> list: # Return all required config attributes @@ -839,12 +870,14 @@ def add_derived_config(self): **Lesson**: Base classes may modify state dicts before custom conversion **Best Practice**: + - Always check what base class does to keys - Test conversion function standalone - Validate key mappings explicitly - Handle both prefixed and non-prefixed keys **Example**: + ```python # Don't assume key format - handle both if key.startswith('model.'): @@ -857,6 +890,7 @@ if key.startswith('model.'): **Lesson**: NeuronAttentionBase integration provides significant benefits **Advantages**: + - Fused QKV projections (performance) - Hardware-optimized kernels - Native GQA support @@ -870,6 +904,7 @@ if key.startswith('model.'): **Lesson**: Tiny → small → medium → full progression accelerates development **Benefits**: + - Faster iteration cycles - Earlier problem identification - Better expert sharding testing @@ -883,6 +918,7 @@ if key.startswith('model.'): **Lesson**: Monitor memory usage throughout compilation pipeline **Key Points**: + - Compilation memory != Runtime memory - Weight sharding is memory-intensive - Use cached NEFFs when available @@ -893,6 +929,7 @@ if key.startswith('model.'): **Lesson**: Comprehensive testing is essential when bypassing built-in checks **Validation Checklist**: + - ✅ Model loading (no errors) - ✅ Shape validation (output matches expected) - ✅ Expert routing (correct selection) @@ -935,6 +972,7 @@ Integration: ### Next Steps With compilation complete, proceed to: + 1. **Category 2**: Understand sharding and memory distribution 2. **Category 3**: Debug accuracy and achieve HuggingFace parity 3. **Inference Testing**: Validate token generation quality diff --git a/skills/neuron-framework-autoport/references/knowledge_base/Category1_Scripts_Compilation_Config_Summary.md b/skills/neuron-framework-autoport/references/knowledge_base/Category1_Scripts_Compilation_Config_Summary.md index 7b430ec..445f2e0 100644 --- a/skills/neuron-framework-autoport/references/knowledge_base/Category1_Scripts_Compilation_Config_Summary.md +++ b/skills/neuron-framework-autoport/references/knowledge_base/Category1_Scripts_Compilation_Config_Summary.md @@ -20,11 +20,13 @@ This category contains scripts focused on the initial phase of the MoE port: get **Problem**: The Neuron compiler's HLO (High-Level Optimizer) verifier fails during weight layout optimization with shape mismatch errors. **Scripts**: + - `recompile_tp16_disable_hlo_verifier.py` - `restart_compilation_tp8.py` - `force_fresh_compilation.py` **Solution Pattern**: + ```python # Disable HLO verifier to work around compiler bug os.environ['NEURON_CC_FLAGS'] = '--internal-hlo2tensorizer-options=--verify-hlo=false' @@ -41,17 +43,20 @@ os.environ['NEURON_CC_FLAGS'] = '--internal-hlo2tensorizer-options=--verify-hlo= **Problem**: Different TP degrees have different compilation success rates and performance characteristics. **Scripts**: + - `restart_compilation_tp8.py` - Tests TP=8 as workaround - `recompile_tp16_disable_hlo_verifier.py` - Targets TP=16 for full utilization - `compile_rank.py` - Implements rank-based compilation for distributed setup - `recompile_minimal_parallelism.py` - Tests minimal parallelism settings **Exploration Timeline**: + 1. **TP=16**: Initial target (uses 16 of 32 cores) - failed with HLO verifier 2. **TP=8**: Intermediate fallback - still failed without HLO verifier disable 3. **TP=16 with HLO verifier disabled**: Final successful configuration **Configuration Pattern**: + ```python config = CompilationConfig( model_class=NeuronGenericMoEForCausalLM, @@ -67,6 +72,7 @@ config = CompilationConfig( ``` **Key Insights**: + - TP=16 provides optimal core utilization (50% of 32 cores) - Expert parallelism (EP) initially avoided due to compilation complexity - Batch size = 1 for inference workloads @@ -79,21 +85,23 @@ config = CompilationConfig( **Problem**: NeuronX has two MoE framework implementations with different characteristics. **Scripts**: + - `examine_setup_all_experts_and_test_flag.py` - `fix_config_attributes.py` - `apply_configuration_fix_and_validate.py` **Framework Comparison**: -| Feature | MoE v1 | MoE v2 | -|---------|--------|--------| -| Expert Parallelism | Limited | Full support | -| Routing Configuration | Simple | Advanced (early_expert_affinity_modulation flag) | -| Compilation Stability | More stable | Requires careful configuration | -| Performance | Good | Better | -| Used in Port | No | Yes | +| Feature | MoE v1 | MoE v2 | +| --------------------- | ----------- | ------------------------------------------------ | +| Expert Parallelism | Limited | Full support | +| Routing Configuration | Simple | Advanced (early_expert_affinity_modulation flag) | +| Compilation Stability | More stable | Requires careful configuration | +| Performance | Good | Better | +| Used in Port | No | Yes | **Key Configuration Flags**: + ```python neuron_config = MoENeuronConfig( tp_degree=16, @@ -109,6 +117,7 @@ neuron_config = MoENeuronConfig( ``` **Critical Discovery**: The `early_expert_affinity_modulation` flag in MoE v2 controls routing weight application: + - **True** (default): Binary expert masking - loses routing weight precision - **False** (correct): Weighted routing - preserves precision and matches HuggingFace behavior @@ -119,12 +128,14 @@ neuron_config = MoENeuronConfig( **Problem**: GenericMoE models require specialized inference configuration that differs from standard transformer models. **Scripts**: + - `debug_neuron_config_dtype.py` - `fix_config_attributes.py` - `debug_model_initialization.py` - `fix_model_wrapper_initialization.py` **Configuration Class Hierarchy**: + ``` GenericMoeInferenceConfig ├── Extends PretrainedConfig (HuggingFace) @@ -133,6 +144,7 @@ GenericMoeInferenceConfig ``` **Key Parameters**: + ```python model_config = GenericMoeInferenceConfig.from_pretrained( model_path, @@ -160,6 +172,7 @@ model_config = GenericMoeInferenceConfig.from_pretrained( ``` **Common Configuration Errors**: + 1. **torch_dtype mismatch**: String "bfloat16" vs torch.bfloat16 object 2. **Missing neuron_config**: Requires explicit MoENeuronConfig initialization 3. **Incorrect GQA heads**: num_key_value_heads must be 8, not 32 @@ -172,6 +185,7 @@ model_config = GenericMoeInferenceConfig.from_pretrained( **Problem**: NeuronX uses a ModelWrapper pattern that requires careful initialization. **Scripts**: + - `fix_model_wrapper_initialization.py` - `debug_model_initialization.py` @@ -188,6 +202,7 @@ model.context_encoding_model.model is not None # ✅ ``` **Workaround Pattern**: + ```python # Create model model = NeuronGenericMoEForCausalLM(model_path, model_config) @@ -285,6 +300,7 @@ if os.path.exists(compiled_output_path): ``` **When to Use**: + - After modifying source code (modeling files) - After changing configuration flags - When debugging mysterious compilation failures @@ -330,6 +346,7 @@ def compile_rank(rank: int): ``` **Launch Pattern**: + ```bash # Compile each rank in parallel for rank in {0..7}; do @@ -425,12 +442,14 @@ def verify_compiled_model(compiled_path): ### 4.1 HLO Verifier Shape Mismatch **Error**: + ``` neuronx-cc error: HLO verifier failed with shape mismatch Exit code: 70 ``` **Fix**: + ```python os.environ['NEURON_CC_FLAGS'] = '--internal-hlo2tensorizer-options=--verify-hlo=false' ``` @@ -442,11 +461,13 @@ os.environ['NEURON_CC_FLAGS'] = '--internal-hlo2tensorizer-options=--verify-hlo= ### 4.2 Missing Configuration Attributes **Error**: + ``` AttributeError: 'GenericMoeInferenceConfig' object has no attribute 'early_expert_affinity_modulation' ``` **Fix**: + ```python # Option 1: Use MoE v2 framework explicitly from neuronx_distributed.modules.moe.expert_mlps_v2 import ExpertMLPsV2 @@ -462,11 +483,13 @@ config.early_expert_affinity_modulation = False ### 4.3 Dtype Mismatches **Error**: + ``` TypeError: expected torch.dtype but got str ``` **Fix**: + ```python # Wrong torch_dtype="bfloat16" @@ -485,12 +508,14 @@ torch_dtype="torch.bfloat16" # String representation for JSON ### 4.4 Model Wrapper Not Initialized **Error**: + ``` AttributeError: 'NoneType' object has no attribute 'forward' # Because model.context_encoding_model.model is None ``` **Fix**: + ```python # Initialize wrappers before use model.context_encoding_model.load_state_dict({}, strict=False) @@ -532,18 +557,21 @@ model.token_generation_model.load_state_dict({}, strict=False) **Systematic Approach**: 1. **Check environment variables**: + ```bash echo $NEURON_CC_FLAGS echo $NEURONX_CACHE ``` 2. **Verify source code**: + ```bash grep -n "early_expert_affinity_modulation" modeling_genericmoe_neuronx.py grep -n "pad=True" modeling_genericmoe_neuronx.py ``` 3. **Clear all caches**: + ```bash rm -rf /tmp/neuron-compile-cache* rm -rf ~/.cache/neuron @@ -566,21 +594,21 @@ model.token_generation_model.load_state_dict({}, strict=False) ### 6.1 Compilation Times -| Configuration | Compilation Time | Memory Usage | -|--------------|------------------|--------------| -| TP=1, 32 layers | 10-15 minutes | 20-30 GB | -| TP=8, 32 layers | 25-35 minutes | 40-50 GB | -| TP=16, 32 layers | 40-60 minutes | 60-80 GB | -| TP=16, EP=8 | 50-70 minutes | 70-90 GB | +| Configuration | Compilation Time | Memory Usage | +| ---------------- | ---------------- | ------------ | +| TP=1, 32 layers | 10-15 minutes | 20-30 GB | +| TP=8, 32 layers | 25-35 minutes | 40-50 GB | +| TP=16, 32 layers | 40-60 minutes | 60-80 GB | +| TP=16, EP=8 | 50-70 minutes | 70-90 GB | ### 6.2 Compilation Artifacts Size -| Component | Size | -|-----------|------| -| Compiled NEFF files | 5-10 GB | +| Component | Size | +| --------------------------- | -------- | +| Compiled NEFF files | 5-10 GB | | Weight files (.safetensors) | 15-20 GB | -| neuron_config.json | < 1 KB | -| Total | 20-30 GB | +| neuron_config.json | < 1 KB | +| Total | 20-30 GB | --- @@ -686,21 +714,25 @@ for prompt in test_prompts: ## 9. Timeline of Compilation Evolution ### Phase 1: Initial Compilation Attempts (Days 1-2) + - Tried TP=16 → Failed with HLO verifier error - Dropped to TP=8 → Still failed - Discovered HLO verifier workaround ### Phase 2: Configuration Refinement (Days 3-4) + - Implemented InferenceConfig properly - Fixed dtype mismatches - Addressed ModelWrapper initialization ### Phase 3: Framework Selection (Days 5-6) + - Explored MoE v1 vs v2 - Discovered early_expert_affinity_modulation flag - Validated MoE v2 as correct choice ### Phase 4: Successful Compilation (Days 7-8) + - TP=16 compilation succeeded with HLO verifier disabled - Generated working compiled artifacts - Validated basic inference functionality @@ -710,6 +742,7 @@ for prompt in test_prompts: ## 10. Reusable Compilation Patterns ### Pattern 1: Quick Compile for Testing + ```python # Minimal configuration for rapid iteration config = CompilationConfig( @@ -722,6 +755,7 @@ config = CompilationConfig( ``` ### Pattern 2: Production Compile + ```python # Full configuration for deployment os.environ['NEURON_CC_FLAGS'] = '--internal-hlo2tensorizer-options=--verify-hlo=false' @@ -737,6 +771,7 @@ config = CompilationConfig( ``` ### Pattern 3: Debug Compile + ```python # Maximum verbosity for troubleshooting os.environ['NEURON_CC_FLAGS'] = '--verbose' @@ -768,6 +803,7 @@ shutil.rmtree(compiled_output_path) The compilation phase of the GenericMoE port required systematic exploration of NeuronX compiler behavior, configuration options, and framework choices. The 22 scripts in this category document a journey from initial compilation failures to reliable, reproducible compilation with TP=16. **Key Takeaway**: Successful MoE compilation on Neuron requires: + 1. Disabling HLO verifier 2. Using MoE v2 framework 3. Proper InferenceConfig implementation diff --git a/skills/neuron-framework-autoport/references/knowledge_base/Category2_Scripts_Sharding_Memory_Weights_Summary.md b/skills/neuron-framework-autoport/references/knowledge_base/Category2_Scripts_Sharding_Memory_Weights_Summary.md index 84af179..77e6ef3 100644 --- a/skills/neuron-framework-autoport/references/knowledge_base/Category2_Scripts_Sharding_Memory_Weights_Summary.md +++ b/skills/neuron-framework-autoport/references/knowledge_base/Category2_Scripts_Sharding_Memory_Weights_Summary.md @@ -42,6 +42,7 @@ Stage 4: Inference Format ``` **Scripts Documenting Pipeline**: + - `transform_spmd_weights.py` - Stage 3 → Stage 4 transformation - `fix_compiled_weights.py` - Complete HF → Inference transformation - `analyze_real_weight_loading.py` - Pipeline analysis and verification @@ -55,6 +56,7 @@ Stage 4: Inference Format **Script**: `fix_compiled_weights.py`, `fix_neuronx_model_weight_loading.py` **Transformation Logic**: + ```python def convert_hf_to_compilation_format(hf_state_dict, config): """Convert HF expert weights to compilation format""" @@ -101,6 +103,7 @@ def convert_hf_to_compilation_format(hf_state_dict, config): ``` **Critical Details**: + 1. **Weight naming**: HF uses `w1/w2/w3`, NeuronX uses `gate_proj/down_proj/up_proj` 2. **Transposition**: All weights must be transposed (.T) 3. **Concatenation**: gate and up projections are concatenated along last dimension @@ -113,6 +116,7 @@ def convert_hf_to_compilation_format(hf_state_dict, config): **Purpose**: Neuron compiler automatically shards weights across TP ranks **What happens during compilation**: + ``` Input (Stage 2): layers.0.block_sparse_moe.expert_mlps.mlp_op.gate_up_proj.weight @@ -130,11 +134,13 @@ After Compilation (Stage 3): ``` **Scripts Analyzing SPMD**: + - `investigate_compiled_model_artifacts.py` - `check_tp_ep_weights.py` - `investigate_expert_routing.py` **Key Discovery**: The SPMD format keys are **automatically generated** by the compiler and depend on: + - TP degree (tensor_model_parallel_size) - EP degree (expert_model_parallel_size) - Total number of experts (16) @@ -150,6 +156,7 @@ After Compilation (Stage 3): **Problem**: Inference framework expects `mlp_op.gate_up_proj.weight` but SPMD creates `spmd_rank.rank_X` **Solution**: Post-compilation weight transformation + ```python def transform_spmd_to_inference(compiled_path): """Transform SPMD weights back to mlp_op format""" @@ -199,21 +206,23 @@ def transform_spmd_to_inference(compiled_path): **Problem**: HuggingFace and NeuronX use different weight key conventions **Scripts**: + - `debug_weight_key_mapping.py` - `fix_weight_key_mismatch.py` - `investigate_missing_weights.py` **Key Mapping Table**: -| HuggingFace Key | NeuronX Key | Notes | -|----------------|-------------|-------| -| `model.embed_tokens.weight` | `embed_tokens.weight` | Remove "model." prefix | -| `lm_head.weight` | `lm_head.weight` | Same (no prefix in HF) | -| `model.layers.X.self_attn.q_proj.weight` | `layers.X.self_attn.qkv_proj.q_proj.weight` | QKV combined in NeuronX | +| HuggingFace Key | NeuronX Key | Notes | +| ----------------------------------------------------- | ------------------------------------------------------------------ | ----------------------- | +| `model.embed_tokens.weight` | `embed_tokens.weight` | Remove "model." prefix | +| `lm_head.weight` | `lm_head.weight` | Same (no prefix in HF) | +| `model.layers.X.self_attn.q_proj.weight` | `layers.X.self_attn.qkv_proj.q_proj.weight` | QKV combined in NeuronX | | `model.layers.X.block_sparse_moe.experts.E.w1.weight` | `layers.X.block_sparse_moe.expert_mlps.mlp_op.gate_up_proj.weight` | Expert weights combined | -| `model.layers.X.block_sparse_moe.gate.weight` | `layers.X.block_sparse_moe.router.linear_router.weight` | Router renamed | +| `model.layers.X.block_sparse_moe.gate.weight` | `layers.X.block_sparse_moe.router.linear_router.weight` | Router renamed | **Conversion Function**: + ```python def convert_key_hf_to_neuronx(hf_key): """Convert HuggingFace key to NeuronX format""" @@ -244,7 +253,9 @@ def convert_key_hf_to_neuronx(hf_key): **Script**: `investigate_missing_keys.py`, `investigate_missing_weights.py` **Common Missing Keys**: + 1. **rank_util.rank tensors**: Need to be added manually + ```python state_dict["rank_util.rank"] = torch.arange(0, tp_degree, dtype=torch.int32) ``` @@ -268,12 +279,14 @@ def convert_key_hf_to_neuronx(hf_key): **Problem**: GenericMoE has ~20GB of weights, and loading/converting them can use 50GB+ RAM **Scripts**: + - `memory_optimized_weight_conversion.py` - `memory_efficient_float32_test.py` - `test_memory_solution_demo.py` - `test_cpu_tp2_memory_safe.py` **Memory Usage Breakdown**: + ``` HuggingFace model in memory: ~20 GB (bfloat16) Weight conversion intermediate: ~30 GB (dtype conversions) @@ -386,6 +399,7 @@ monitor_memory_usage() ### 4.1 Forward All Experts Analysis **Scripts**: + - `examine_neuronx_forward_all_experts_implementation.py` - `investigate_expert_routing.py` - `investigate_moe_structure.py` @@ -395,6 +409,7 @@ monitor_memory_usage() **Two Routing Methods Identified**: **Method 1: Binary Masking** (early_expert_affinity_modulation=True) + ```python # expert_mask: one-hot encoding of selected experts expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=num_experts) @@ -409,6 +424,7 @@ for e in range(num_experts): ``` **Method 2: Weighted Routing** (early_expert_affinity_modulation=False) + ```python # expert_affinities_masked: routing weights applied to mask expert_affinities_masked = torch.zeros(batch*seq, num_experts) @@ -433,6 +449,7 @@ for e in range(num_experts): **Script**: `apply_early_expert_affinity_modulation_fix.py` **Test Demonstrating Difference**: + ```python # Test scenario routing_weights = torch.tensor([[0.8, 0.2], [0.6, 0.4]]) # Fractional weights @@ -450,6 +467,7 @@ expert_outputs = torch.tensor([[2.0], [4.0], [6.0], [8.0]]) # Different values ``` **Accuracy Impact**: + - Method 1: Predicts wrong tokens (e.g., 'a' instead of 'Paris') - Method 2: Matches HuggingFace predictions exactly @@ -460,19 +478,21 @@ expert_outputs = torch.tensor([[2.0], [4.0], [6.0], [8.0]]) # Different values ### 5.1 TP vs EP Tradeoffs **Scripts**: + - `check_tp_ep_weights.py` - `test_expert_sharding.py` - `test_cpu_tp2_memory_safe.py` **Parallelism Strategies**: -| Strategy | Description | Pros | Cons | -|----------|-------------|------|------| -| **TP Only** | Shard attention/FFN across ranks | Simpler, more stable | Limited expert distribution | -| **EP Only** | Distribute experts across ranks | Expert-specific scaling | Complex routing | -| **TP + EP** | Both attention and experts sharded | Maximum parallelism | Most complex, harder debugging | +| Strategy | Description | Pros | Cons | +| ----------- | ---------------------------------- | ----------------------- | ------------------------------ | +| **TP Only** | Shard attention/FFN across ranks | Simpler, more stable | Limited expert distribution | +| **EP Only** | Distribute experts across ranks | Expert-specific scaling | Complex routing | +| **TP + EP** | Both attention and experts sharded | Maximum parallelism | Most complex, harder debugging | **Recommended Configuration**: + ```python # For 32 cores tp_degree = 16 # Use half cores for tensor parallelism @@ -481,6 +501,7 @@ moe_ep_degree = 1 # Initially disable expert parallelism ``` **Rationale**: + - Start with TP-only for stability - Add EP later after confirming correctness - TP=16 provides good balance of parallelism and simplicity @@ -503,6 +524,7 @@ experts_per_rank = 16 / 8 = 2 ``` **Weight Sharding**: + ```python # Each rank stores only its experts rank_0_weights = { @@ -526,6 +548,7 @@ rank_1_weights = { **Discovery**: ColumnParallelLinear has a hidden precision loss in the allreduce operation **Problem Location**: + ``` File: neuronx_distributed/parallel_layers/layers_utils.py Function: _linear_autograd_bwd_grad_reduce @@ -533,6 +556,7 @@ Lines: 99-102 ``` **Problematic Code**: + ```python if ctx.async_grad_allreduce: # Convert to reduce_dtype (default: float32) @@ -546,6 +570,7 @@ if ctx.async_grad_allreduce: ``` **Why It Matters**: + - The float32 → bfloat16 conversion introduces quantization artifacts - Differences are exactly 1/64 (0.015625) multiples - Accumulates through 32 layers @@ -558,6 +583,7 @@ if ctx.async_grad_allreduce: **Script**: `columnparallel_precision_root_cause_analysis.py` **Test Results**: + ```python # With default reduce_dtype=float32 max_diff = 0.015625 # Exactly 1/64 @@ -570,6 +596,7 @@ percentage_1_64_multiples = 77% # Most differences are 1/64 multiples ``` **Fix**: + ```python # Set reduce_dtype to match tensor dtype q_proj = ColumnParallelLinear( @@ -628,6 +655,7 @@ def verify_weight_loading(hf_model, neuronx_model): ### 7.2 Common Weight Loading Failures **Failure 1: Transposition Errors** + ```python # Wrong - weights transposed incorrectly gate_up_proj[e] = torch.cat([w1, w3], dim=-1) @@ -637,6 +665,7 @@ gate_up_proj[e] = torch.cat([w1.T, w3.T], dim=-1) ``` **Failure 2: Dimension Concatenation** + ```python # Wrong - concatenating along wrong dimension gate_up = torch.cat([gate, up], dim=0) # [2*intermediate, hidden] @@ -646,6 +675,7 @@ gate_up = torch.cat([gate, up], dim=-1) # [hidden, 2*intermediate] ``` **Failure 3: Expert Index Off-by-One** + ```python # Wrong - starting from 1 for expert_idx in range(1, num_experts + 1): @@ -693,6 +723,7 @@ def load_weights_distributed(model_path, compiled_path, rank, world_size): ### 8.2 Weight Sharding Strategy **For TP=16, EP=1**: + ```python # Attention weights: Sharded across TP dimension # Q/K/V projections: Split hidden_size across 16 ranks @@ -708,6 +739,7 @@ expert_weights_rank_1 = all_expert_weights # Full copy ``` **For TP=16, EP=8**: + ```python # Attention weights: Same as above (sharded by TP) @@ -736,6 +768,7 @@ expert_weights_rank_1 = all_expert_weights # Full copy ### 9.2 SPMD Pipeline Insights **Critical Understanding**: + - Stage 1→2: Manual transformation (our code) - Stage 2→3: Automatic (compiler does it) - Stage 3→4: Manual transformation (our code) **OR** compiler should do it properly @@ -764,12 +797,14 @@ expert_weights_rank_1 = all_expert_weights # Full copy ### 10.1 Systematic Debugging Approach **Step 1: Verify Weight Files Exist** + ```bash ls -lh model/model.safetensors* # Should show multiple safetensors files totaling ~20GB ``` **Step 2: Check Weight Keys** + ```python from safetensors import safe_open with safe_open("model/model.safetensors", framework="pt") as f: @@ -779,6 +814,7 @@ with safe_open("model/model.safetensors", framework="pt") as f: ``` **Step 3: Verify Key Conversion** + ```python hf_keys = set(hf_state_dict.keys()) neuron_keys = set(neuron_state_dict.keys()) @@ -790,6 +826,7 @@ print(f"Missing keys: {missing}") ``` **Step 4: Verify Weight Values** + ```python # Check weight statistics for key, tensor in neuron_state_dict.items(): @@ -804,6 +841,7 @@ for key, tensor in neuron_state_dict.items(): ``` **Step 5: Compare with HF Weights** + ```python # Direct comparison hf_weight = hf_model.model.embed_tokens.weight @@ -819,21 +857,21 @@ print(f"Cosine similarity: {cos_sim:.6f}") # Should be > 0.99 ### 11.1 Weight Loading Performance -| Operation | Time (32-layer model) | Memory Peak | -|-----------|----------------------|-------------| -| Load HF safetensors | 30-60 seconds | +20 GB | -| Convert to NeuronX format | 2-5 minutes | +30 GB | -| Save NeuronX safetensors | 30-60 seconds | +10 GB | -| **Total** | **3-7 minutes** | **60 GB peak** | +| Operation | Time (32-layer model) | Memory Peak | +| ------------------------- | --------------------- | -------------- | +| Load HF safetensors | 30-60 seconds | +20 GB | +| Convert to NeuronX format | 2-5 minutes | +30 GB | +| Save NeuronX safetensors | 30-60 seconds | +10 GB | +| **Total** | **3-7 minutes** | **60 GB peak** | ### 11.2 Optimized Loading Performance -| Operation | Time (optimized) | Memory Peak | -|-----------|------------------|-------------| -| Load HF safetensors (mmap) | 5-10 seconds | +5 GB | -| Convert incrementally | 3-4 minutes | +20 GB | -| Save sharded checkpoints | 1-2 minutes | +5 GB | -| **Total** | **4-6 minutes** | **30 GB peak** | +| Operation | Time (optimized) | Memory Peak | +| -------------------------- | ---------------- | -------------- | +| Load HF safetensors (mmap) | 5-10 seconds | +5 GB | +| Convert incrementally | 3-4 minutes | +20 GB | +| Save sharded checkpoints | 1-2 minutes | +5 GB | +| **Total** | **4-6 minutes** | **30 GB peak** | **Optimization Impact**: 50% memory reduction, similar time @@ -842,27 +880,35 @@ print(f"Cosine similarity: {cos_sim:.6f}") # Should be > 0.99 ## 12. Common Errors and Solutions ### Error 1: SPMD Key Not Found + ``` KeyError: 'layers.0.block_sparse_moe.expert_mlps.mlp_op.gate_up_proj.weight' ``` + **Solution**: Run Stage 3→4 transformation script ### Error 2: Weight Shape Mismatch + ``` RuntimeError: shape mismatch: [16, 4096, 12800] vs [16, 6400, 4096] ``` + **Solution**: Check transposition - weights must be .T ### Error 3: Missing rank_util.rank Tensors + ``` KeyError: 'rank_util.rank' ``` + **Solution**: Add rank tensors manually to state dict ### Error 4: Out of Memory During Conversion + ``` RuntimeError: CUDA out of memory ``` + **Solution**: Use memory-optimized conversion pattern --- @@ -886,6 +932,7 @@ RuntimeError: CUDA out of memory The weight management phase of the GenericMoE port required understanding a complex 4-stage transformation pipeline, implementing memory-efficient loading strategies, and discovering subtle precision issues in distributed operations. The 34 scripts in this category document the systematic exploration and resolution of weight-related challenges. **Key Takeaways**: + 1. **SPMD pipeline is critical**: Understanding the 4 stages prevents many debugging headaches 2. **Memory matters**: 15-20GB models need careful memory management 3. **Precision is fragile**: Small dtype conversion errors accumulate significantly diff --git a/skills/neuron-framework-autoport/references/knowledge_base/Category2_Sharding_Memory_Issues.md b/skills/neuron-framework-autoport/references/knowledge_base/Category2_Sharding_Memory_Issues.md index 28ac2d8..942b645 100644 --- a/skills/neuron-framework-autoport/references/knowledge_base/Category2_Sharding_Memory_Issues.md +++ b/skills/neuron-framework-autoport/references/knowledge_base/Category2_Sharding_Memory_Issues.md @@ -11,6 +11,7 @@ This document details the complex weight sharding and memory management challeng ## Document Organization ### Source Documents Analyzed: + - expert_sharding_complete.md - moe_sharding_analysis_detailed.md - moe_analysis_comprehensive.md @@ -26,6 +27,7 @@ This document details the complex weight sharding and memory management challeng **Core Issue**: MoE models require expert weights to be distributed across tensor parallel ranks, but the NeuronX framework expected **fundamentally different weight formats** for compilation vs inference. **Manifestation**: + ``` Compilation Phase: - Produces: layers.X.block_sparse_moe.expert_mlps.spmd_rank.rank @@ -39,6 +41,7 @@ Inference Phase: ``` **Configuration Mismatch**: + ```python # Compilation setting blockwise_matmul_config.parallelize_token_to_block_mapping = True @@ -58,6 +61,7 @@ This fundamental difference required a sophisticated multi-stage transformation **Format**: Separate weight matrices for each expert in each layer **Structure**: + ```python # For each layer (32 layers total): # For each expert (16 experts per layer): @@ -79,6 +83,7 @@ This fundamental difference required a sophisticated multi-stage transformation ``` **Characteristics**: + - ✅ Easy to understand and debug - ✅ Matches HuggingFace implementation - ❌ Not suitable for NeuronX compilation @@ -142,6 +147,7 @@ def convert_hf_to_neuron_compilation_format(hf_state_dict, config): ``` **Resulting Format**: + ```python # For each layer: "model.layers.0.mlp.gate_up_proj.weight": [16, 4096, 28672] @@ -157,12 +163,14 @@ def convert_hf_to_neuron_compilation_format(hf_state_dict, config): ``` **Key Transformations**: + 1. **Concatenation**: gate_proj + up_proj → gate_up_proj (efficiency optimization) 2. **Stacking**: Individual expert weights → single stacked tensor 3. **Transpose**: [out_features, in_features] → [in_features, out_features] 4. **Dimension Reduction**: 1,536 tensors → 64 tensors **Benefits**: + - ✅ Compiler can optimize across all experts - ✅ Enables efficient expert routing - ✅ Reduces number of parameters to track @@ -202,12 +210,14 @@ Rank 15: [16 experts, hidden[3840:4096], intermediate[26880:28672]] ``` **Key Characteristics**: + - All 16 experts present on each rank - Weights are **sharded** (divided), not replicated - Each rank has 1/16th of the weight dimensions - Memory per rank: ~2GB (vs ~32GB if fully replicated) **Files Generated**: + ``` weights/tp0_sharded_checkpoint.safetensors # Rank 0 weights weights/tp1_sharded_checkpoint.safetensors # Rank 1 weights @@ -278,6 +288,7 @@ def fix_compiled_weights(compiled_model_path, tp_degree=16): ``` **Final Format**: + ```python # Each rank's checkpoint now has: "layers.0.block_sparse_moe.expert_mlps.mlp_op.gate_up_proj.weight": [16, 256, 1792] @@ -329,14 +340,15 @@ STAGE 4: Inference Format ### Memory Impact at Each Stage -| Stage | Description | Tensors | Size/Rank | Total Size | Format | -|-------|-------------|---------|-----------|------------|--------| -| 1 | HuggingFace | 1,536 | N/A | ~41GB | Individual | -| 2 | Compilation | 64 | N/A | ~41GB | Concatenated | -| 3 | SPMD Sharded | 32 | ~2.5GB | ~40GB | Sharded | -| 4 | Inference | 64 | ~2.5GB | ~40GB | Sharded | +| Stage | Description | Tensors | Size/Rank | Total Size | Format | +| ----- | ------------ | ------- | --------- | ---------- | ------------ | +| 1 | HuggingFace | 1,536 | N/A | ~41GB | Individual | +| 2 | Compilation | 64 | N/A | ~41GB | Concatenated | +| 3 | SPMD Sharded | 32 | ~2.5GB | ~40GB | Sharded | +| 4 | Inference | 64 | ~2.5GB | ~40GB | Sharded | **Key Insight**: Total size remains constant (~40GB), but organization changes dramatically to enable: + - Efficient compilation (Stage 2) - Tensor parallelism (Stage 3) - Inference compatibility (Stage 4) @@ -356,6 +368,7 @@ Through deep analysis of existing MoE models (Qwen3, test suite), we discovered **Concept**: Distribute experts across ranks (each rank has subset of experts) **Configuration**: + ```python neuron_config = MoENeuronConfig( tp_degree=8, # Tensor parallel degree @@ -365,6 +378,7 @@ neuron_config = MoENeuronConfig( ``` **Expert Distribution** (EP=8, 16 total experts): + ``` Rank 0: Experts [0, 1] # 2 experts per rank Rank 1: Experts [2, 3] @@ -377,6 +391,7 @@ Rank 7: Experts [14, 15] ``` **Memory Calculation**: + ```python # Original: All 16 experts on each rank = ~16GB per rank # With EP=8: 2 experts per rank @@ -387,6 +402,7 @@ memory_reduction = ep_degree # 8x reduction ``` **Process Group Creation**: + ```python def initialize_model_parallel( tensor_model_parallel_size: int = 8, @@ -402,6 +418,7 @@ def initialize_model_parallel( ``` **Expert Assignment**: + ```python def get_experts_for_expert_parallel_rank( expert_parallel_rank: int, @@ -422,12 +439,14 @@ get_experts_for_expert_parallel_rank(7, 16, 8) # Returns [14, 15] ``` **Advantages**: + - ✅ Maximum memory reduction (16x possible) - ✅ True expert distribution across ranks - ✅ Lower memory per rank - ✅ Scales to more experts efficiently **Disadvantages**: + - ❌ More complex communication patterns - ❌ All-to-all required for expert routing - ❌ **Critical limitation**: "Selective Loading with Expert parallelism is not supported in token generation" @@ -439,6 +458,7 @@ get_experts_for_expert_parallel_rank(7, 16, 8) # Returns [14, 15] **Concept**: Replicate all experts on each rank, but shard weight dimensions **Configuration**: + ```python neuron_config = MoENeuronConfig( tp_degree=16, # Tensor parallel degree @@ -449,6 +469,7 @@ neuron_config = MoENeuronConfig( ``` **Expert Distribution** (TP=16, EP=1): + ``` Rank 0: All 16 experts, weights sharded [hidden[0:256], intermediate[0:1792]] Rank 1: All 16 experts, weights sharded [hidden[256:512], intermediate[1792:3584]] @@ -458,6 +479,7 @@ Rank 15: All 16 experts, weights sharded [hidden[3840:4096], intermediate[26880: ``` **Memory Calculation**: + ```python # All 16 experts on each rank, but weights are sharded @@ -473,6 +495,7 @@ memory_reduction = tp_degree # 16x through dimension sharding ``` **Why This Works**: + ```python # During forward pass: # 1. Each rank computes its shard: hidden[rank_start:rank_end] @@ -487,6 +510,7 @@ expert_output_full = all_reduce(expert_output_shard) ``` **Advantages**: + - ✅ **Simpler communication**: Standard tensor parallelism patterns - ✅ **No expert-specific routing**: All-reduce sufficient - ✅ **Token generation compatible**: No selective loading issues @@ -494,6 +518,7 @@ expert_output_full = all_reduce(expert_output_shard) - ✅ **Proven stability**: Extensive testing in framework **Disadvantages**: + - ⚠️ All experts must fit in memory (even if sharded) - ⚠️ Less memory reduction than full expert parallelism - ⚠️ Communication overhead for all-reduce @@ -502,22 +527,23 @@ expert_output_full = all_reduce(expert_output_shard) ### Strategy Comparison for Generic MoE -| Aspect | Expert Parallelism (EP=8) | Tensor Parallelism (EP=1, TP=16) | -|--------|---------------------------|----------------------------------| -| **Experts per rank** | 2 (16/8) | 16 (all) | -| **Weight sharding** | None | Yes (16x) | -| **Memory per rank** | ~2GB | ~2GB | -| **Communication** | All-to-all | All-reduce | -| **Token generation** | ❌ Not supported | ✅ Supported | -| **Complexity** | Higher | Lower | -| **Production ready** | Limited | ✅ Yes | -| **Used by** | Test suite | Qwen3, production | +| Aspect | Expert Parallelism (EP=8) | Tensor Parallelism (EP=1, TP=16) | +| -------------------- | ------------------------- | -------------------------------- | +| **Experts per rank** | 2 (16/8) | 16 (all) | +| **Weight sharding** | None | Yes (16x) | +| **Memory per rank** | ~2GB | ~2GB | +| **Communication** | All-to-all | All-reduce | +| **Token generation** | ❌ Not supported | ✅ Supported | +| **Complexity** | Higher | Lower | +| **Production ready** | Limited | ✅ Yes | +| **Used by** | Test suite | Qwen3, production | --- ### Critical Framework Limitation Discovery **Finding from Framework Analysis**: + ``` Error message: "Selective Loading with Expert parallelism is not supported in token generation" @@ -526,6 +552,7 @@ Impact: Cannot use EP > 1 for autoregressive generation ``` **What This Means**: + - Expert parallelism (EP > 1) works for: - ✅ Compilation - ✅ Single forward passes @@ -576,6 +603,7 @@ neuron_config = MoENeuronConfig( ### Expert and Memory Distribution **Distribution Pattern**: + ``` Hardware: AWS Trainium (trn1.32xlarge) with 32 Neuron cores Utilized: 16 cores (TP=16) @@ -593,6 +621,7 @@ Each of 16 ranks has: ``` **Load Balancing Analysis**: + ```python # Generic MoE configuration total_experts = 16 @@ -610,6 +639,7 @@ load_per_rank = active_experts_per_token / experts_per_rank # 2/16 = 0.125 ``` **Communication Pattern**: + ```python # Simplified forward pass with TP=16 @@ -690,16 +720,17 @@ sharded_per_rank = total_expert_weights / 16 **Important Distinction**: -| Phase | Memory Type | Amount | Purpose | -|-------|-------------|--------|---------| -| Compilation | **Peak** | ~188GB | Graph optimization, weight analysis | -| Compilation | **Temporary** | ~50GB | Weight sharding overhead | -| Runtime | **Per Rank** | ~5.5GB | Inference execution | -| Runtime | **Total** | ~88GB | All 16 ranks combined | +| Phase | Memory Type | Amount | Purpose | +| ----------- | ------------- | ------ | ----------------------------------- | +| Compilation | **Peak** | ~188GB | Graph optimization, weight analysis | +| Compilation | **Temporary** | ~50GB | Weight sharding overhead | +| Runtime | **Per Rank** | ~5.5GB | Inference execution | +| Runtime | **Total** | ~88GB | All 16 ranks combined | **Key Insight**: Compilation memory >> Runtime memory **Implications**: + - Need large-memory instance for compilation (we used ~256GB) - Can use smaller instances for inference serving - Weight sharding reduces runtime memory dramatically @@ -734,6 +765,7 @@ sharded_per_rank = total_expert_weights / 16 ``` **Result**: + - Theoretical: ~11.25GB per rank - Actual: ~5.5GB per rank - Reduction: **~2x through compiler optimizations** @@ -975,6 +1007,7 @@ def validate_load_balancing(model, test_inputs, num_iterations=100): **Lesson**: MoE models require multi-stage weight transformation pipeline **Best Practice**: + - Document each transformation stage clearly - Validate weight shapes at each stage - Test transformations on small models first @@ -982,6 +1015,7 @@ def validate_load_balancing(model, test_inputs, num_iterations=100): - Version control transformation code **Common Pitfalls**: + - ❌ Assuming single transformation is sufficient - ❌ Not handling transpose operations correctly - ❌ Forgetting to update weight keys for inference @@ -992,12 +1026,14 @@ def validate_load_balancing(model, test_inputs, num_iterations=100): **Lesson**: Framework limitations dictate strategy choice **Decision Criteria**: + 1. **Check token generation support** (critical for autoregressive models) 2. **Evaluate memory constraints** (EP vs TP tradeoffs) 3. **Consider communication patterns** (all-to-all vs all-reduce) 4. **Review production readiness** (proven vs experimental) **For Generic MoE**: + - ✅ Chose TP=16, EP=1 due to token generation limitation - ✅ Achieved same memory efficiency through weight sharding - ✅ Used proven, stable approach (Qwen3 model precedent) @@ -1007,6 +1043,7 @@ def validate_load_balancing(model, test_inputs, num_iterations=100): **Lesson**: Distinguish compilation vs runtime memory requirements **Best Practice**: + - Plan for 5-10x more memory during compilation - Monitor peak memory usage - Use weight checkpointing if needed @@ -1014,6 +1051,7 @@ def validate_load_balancing(model, test_inputs, num_iterations=100): - Test on representative hardware **Memory Planning**: + ``` Compilation Machine: 256GB RAM (for 29B model) Inference Machine: 16GB per rank × 16 ranks = 256GB total @@ -1025,6 +1063,7 @@ But: Can distribute across multiple smaller machines **Lesson**: Comprehensive validation is critical for sharded models **Validation Checklist**: + - ✅ Weight dimensions match expected sharding - ✅ All ranks have consistent expert counts - ✅ Memory per rank within limits @@ -1038,12 +1077,14 @@ But: Can distribute across multiple smaller machines **Lesson**: Test sharding on small models before full scale **Recommended Progression**: + 1. **2 experts, TP=2**: Validate basic sharding 2. **4 experts, TP=4**: Test moderate scale 3. **8 experts, TP=8**: Approach production scale 4. **16 experts, TP=16**: Full production deployment **Benefits**: + - Faster iteration - Earlier problem detection - Better understanding of patterns diff --git a/skills/neuron-framework-autoport/references/knowledge_base/Category3_Accuracy_Debugging_Analysis.md b/skills/neuron-framework-autoport/references/knowledge_base/Category3_Accuracy_Debugging_Analysis.md index 6579030..58cb978 100644 --- a/skills/neuron-framework-autoport/references/knowledge_base/Category3_Accuracy_Debugging_Analysis.md +++ b/skills/neuron-framework-autoport/references/knowledge_base/Category3_Accuracy_Debugging_Analysis.md @@ -11,6 +11,7 @@ This document details the extensive accuracy debugging efforts to achieve Huggin ## Document Organization ### Source Documents Analyzed: + - FINAL_COMPLETE_SOLUTION_SUMMARY.md - precision_loss_comprehensive_analysis_with_code_differences.md - ROUTING_WEIGHT_APPLICATION_SOLUTION_COMPLETE.md @@ -68,6 +69,7 @@ def debug_accuracy_systematic(hf_model, neuronx_model, test_input): ### Metrics Used **1. Cosine Similarity**: + ```python def cosine_similarity(tensor1, tensor2): """Measure directional similarity between tensors""" @@ -83,6 +85,7 @@ def cosine_similarity(tensor1, tensor2): ``` **2. Maximum Absolute Difference**: + ```python def max_abs_diff(tensor1, tensor2): """Maximum element-wise difference""" @@ -96,6 +99,7 @@ def max_abs_diff(tensor1, tensor2): ``` **3. Weight Statistics**: + ```python def weight_statistics(tensor): """Statistical properties of weight tensor""" @@ -120,6 +124,7 @@ def weight_statistics(tensor): ### Problem Discovery **Initial Symptoms**: + ``` Missing keys: 450 (attention weights) Unexpected keys: 517 @@ -128,6 +133,7 @@ Attention output: max_diff=3.128296, cos_sim=0.004095 ``` **Test Case**: + ```python prompt = "The capital of France is" # HuggingFace: "Paris" ✅ @@ -152,6 +158,7 @@ def convert_generic_moe_hf_to_neuron_state_dict(hf_state_dict, config): ``` **Why This Broke**: + 1. Base class already removes `model.` prefix before calling conversion 2. Conversion function removed it again → double removal 3. Result: Keys completely mismatched @@ -187,6 +194,7 @@ def convert_generic_moe_hf_to_neuron_state_dict(hf_state_dict, config): ``` **Key Mapping Examples**: + ```python # Correct transformations: "model.layers.0.self_attn.q_proj.weight" → "model.layers.0.self_attn.qkv_proj.q_proj.weight" @@ -198,6 +206,7 @@ def convert_generic_moe_hf_to_neuron_state_dict(hf_state_dict, config): ### Validation **Weight Loading Verification**: + ```python def verify_attention_weights(model, hf_state_dict): """Verify attention weights loaded correctly""" @@ -230,6 +239,7 @@ def verify_attention_weights(model, hf_state_dict): ### Results **Before Fix**: + ``` Missing keys: 450 Weight comparison: max_diff=0.676788, cos_sim=-0.000042 ❌ @@ -238,6 +248,7 @@ Prediction: Random nonsense ❌ ``` **After Fix**: + ``` Missing keys: 0 ✅ Weight comparison: max_diff=0.000000, cos_sim=1.000000 ✅ @@ -253,6 +264,7 @@ Prediction: Improved (but still other issues) ⚠️ **Symptom**: Even with weights loading correctly, Layer 0 output showed divergence: + ``` Layer 0 output: cos_sim=0.859817, max_diff=0.262207 ``` @@ -282,14 +294,15 @@ nn.LayerNorm(hidden_size, eps=config.rms_norm_eps) **Mathematical Difference**: -| Operation | LayerNorm | RMSNorm | -|-----------|-----------|---------| -| Mean subtraction | ✅ Yes | ❌ No | -| Variance calculation | After centering | Of raw values | -| Bias term | ✅ Yes | ❌ No | -| Formula | `(x-μ)/σ * w + b` | `x/RMS * w` | +| Operation | LayerNorm | RMSNorm | +| -------------------- | ----------------- | ------------- | +| Mean subtraction | ✅ Yes | ❌ No | +| Variance calculation | After centering | Of raw values | +| Bias term | ✅ Yes | ❌ No | +| Formula | `(x-μ)/σ * w + b` | `x/RMS * w` | **Impact**: + ```python # Example with simple input input = torch.tensor([1.0, 2.0, 3.0, 4.0]) @@ -352,6 +365,7 @@ class GenericMoEModel(nn.Module): ### Validation **Normalization Comparison**: + ```python def compare_normalizations(hf_model, neuronx_model, test_input): """Compare normalization outputs""" @@ -389,6 +403,7 @@ def compare_normalizations(hf_model, neuronx_model, test_input): ### Results **Before Fix**: + ``` HF: LayerNorm ✅ NeuronX: RMSNorm ❌ @@ -396,6 +411,7 @@ Layer 0 output: cos_sim=0.859817, max_diff=0.262207 ❌ ``` **After Fix**: + ``` HF: LayerNorm ✅ NeuronX: LayerNorm ✅ @@ -410,6 +426,7 @@ Layer 0 output: cos_sim=0.92 (improved) ⚠️ Still other issues ### Problem Discovery **Persistent Pattern**: + ```python # After fixing weights and normalization, still had precision differences max_diff = 0.015625 # Exactly 1/64 @@ -578,6 +595,7 @@ print(f"✅ This explains the model prediction differences") ``` **Output**: + ``` 🔬 PRECISION LOSS DEMONSTRATION Input shape: torch.Size([1, 5, 4096]) @@ -605,6 +623,7 @@ torch.matmul + bias 0.015625 ❌ Standard PyTorch - 🎯 EXACT 1/ ### Solution Options Documented **Option 1: Include Bias in Linear Operation (Recommended)**: + ```python # In LinearWithAsyncCommunication.forward(): # CURRENT: @@ -617,6 +636,7 @@ output = torch.nn.functional.linear(total_input, weight, bias) # ✅ No precisi ``` **Option 2: Higher Precision Bias Addition**: + ```python # In ColumnParallelLinear.forward(): # CURRENT: @@ -630,6 +650,7 @@ else: ``` **Option 3: Configuration-Based Precision Mode**: + ```python class ColumnParallelLinear: def __init__(self, ..., high_precision_bias=False): @@ -650,6 +671,7 @@ class ColumnParallelLinear: ### Impact Analysis **Cascading Effect Through Model**: + ```python # Small precision differences cascade through 32 layers @@ -680,6 +702,7 @@ Final logits: 18.937500 difference ❌ **Symptom**: Even after fixing precision issues, MoE output still had significant differences: + ``` MoE layer output difference: ~0.13 (large) Token prediction: Still wrong ("a" instead of "Paris") @@ -791,6 +814,7 @@ for layer in model.model.layers: ``` **If Override Needed**: + ```python # Force correct setting at runtime for layer in model.model.layers: @@ -832,6 +856,7 @@ def validate_routing_weight_application(model): ### Results **Before Fix** (early_expert_affinity_modulation=True): + ``` MoE output: [6.0, 14.0] (binary routing) HuggingFace: [2.4062, 6.8125] (weighted routing) @@ -840,6 +865,7 @@ Token prediction: "a" (wrong) ``` **After Fix** (early_expert_affinity_modulation=False): + ``` MoE output: [2.4062, 6.8125] (weighted routing) HuggingFace: [2.4062, 6.8125] (weighted routing) @@ -854,6 +880,7 @@ Token prediction: "Paris" (correct) ✅ ### Problem Discovery **Symptom**: + ```python # Model occasionally generated nonsense tokens model.generate(input_ids) @@ -982,6 +1009,7 @@ def sample_token_safe(logits, vocab_size): ### Results **Before Fix**: + ``` Phantom tokens: 32064-32767 (704 tokens) pad_size: 0 (perfectly aligned) @@ -991,6 +1019,7 @@ Output quality: Occasional nonsense ❌ ``` **After Fix**: + ``` Phantom token detection: Active ✅ Masking: Applied regardless of pad_size ✅ @@ -1007,6 +1036,7 @@ Output quality: Consistent, coherent ✅ **Goal**: Capture intermediate tensors for debugging on CPU **Symptom**: + ```python RuntimeError: The size of tensor a (256) must match the size of tensor b (128) at non-singleton dimension 3 @@ -1032,6 +1062,7 @@ ratio = 32 / 8 = 4 # 4:1 ratio ``` **Why This Happens**: + 1. NeuronAttentionBase has GQA optimizations for hardware 2. These optimizations assume Neuron device characteristics 3. CPU execution path doesn't have same tensor layouts @@ -1060,6 +1091,7 @@ neuron-profile export \ ``` **Advantages**: + - ✅ Captures actual hardware execution - ✅ No GQA compatibility issues - ✅ Accurate performance metrics @@ -1103,6 +1135,7 @@ for name, tensor in captured_tensors.items(): ``` **Advantages**: + - ✅ Works on CPU without issues - ✅ Standard PyTorch hooks - ✅ Easy to debug @@ -1130,6 +1163,7 @@ class GenericMoEAttention(NeuronAttentionBase): ``` **Disadvantages**: + - ❌ 4x more KV cache memory - ❌ Not representative of actual model - ❌ Different behavior than production @@ -1140,6 +1174,7 @@ class GenericMoEAttention(NeuronAttentionBase): **Chosen Approach**: Neuron device profiling + HuggingFace comparison **Outcome**: + - ✅ Successfully captured tensors on Neuron hardware - ✅ Used HF model for CPU debugging - ✅ Avoided GQA compatibility issues @@ -1156,6 +1191,7 @@ class GenericMoEAttention(NeuronAttentionBase): ### Test Results **Test 1: Capital of France** ✅ + ```python Prompt: "The capital of France is" Generated: "Paris." @@ -1163,6 +1199,7 @@ Result: ✅ PERFECT - Correctly predicted "Paris" ``` **Test 2: Simple Math** ✅ + ```python Prompt: "2 + 2 =" Generated: "4." @@ -1170,6 +1207,7 @@ Result: ✅ PERFECT - Correctly calculated "4" ``` **Test 3: General Knowledge** ✅ + ```python Prompt: "The sun rises in the" Generated: "east and sets in the west. This is a" @@ -1177,6 +1215,7 @@ Result: ✅ PERFECT - Accurate and coherent ``` **Test 4: Conversation** ✅ + ```python Prompt: "Hello, my name is" Generated: "Alex. Hello, Alex! It's nice to meet you. How" @@ -1186,6 +1225,7 @@ Result: ✅ PERFECT - Natural and engaging ### Technical Validation **Numerical Stability**: + ``` Logits: Min: -10.7500 @@ -1204,6 +1244,7 @@ Probabilities: ``` **Weight Loading**: + ``` Total weights loaded: 484 Missing keys: 0 ✅ @@ -1212,6 +1253,7 @@ Weight loading success: 100% ✅ ``` **Model Performance**: + ``` Model loading time: 55 seconds Warmup time: 0.84 seconds @@ -1244,18 +1286,19 @@ outputs = model(input_ids=input_ids, position_ids=position_ids) ### Issues Resolved -| Issue | Category | Impact | Status | -|-------|----------|--------|--------| -| 1. Attention weight loading | Critical | Model non-functional | ✅ Fixed | -| 2. LayerNorm vs RMSNorm | Critical | Wrong normalization | ✅ Fixed | -| 3. bfloat16 precision (1/64) | Major | Cascading errors | ✅ Documented | -| 4. MoE routing weights | Critical | Wrong predictions | ✅ Fixed | -| 5. Phantom token masking | Moderate | Occasional nonsense | ✅ Fixed | -| 6. Tensor capture GQA | Debug only | CPU incompatibility | ✅ Workaround | +| Issue | Category | Impact | Status | +| ---------------------------- | ---------- | -------------------- | ------------- | +| 1. Attention weight loading | Critical | Model non-functional | ✅ Fixed | +| 2. LayerNorm vs RMSNorm | Critical | Wrong normalization | ✅ Fixed | +| 3. bfloat16 precision (1/64) | Major | Cascading errors | ✅ Documented | +| 4. MoE routing weights | Critical | Wrong predictions | ✅ Fixed | +| 5. Phantom token masking | Moderate | Occasional nonsense | ✅ Fixed | +| 6. Tensor capture GQA | Debug only | CPU incompatibility | ✅ Workaround | ### Before vs After **Before All Fixes**: + ``` Prediction: "a" (token 263) ❌ Token prediction accuracy: 0% @@ -1266,6 +1309,7 @@ Status: Non-functional ``` **After All Fixes**: + ``` Prediction: "Paris" (token 3681) ✅ Token prediction accuracy: 100% @@ -1400,6 +1444,7 @@ def test_routing_weight_configuration(model): **Lesson**: Component-by-component analysis finds issues faster than end-to-end **Best Practice**: + 1. Start with embeddings (should be perfect) 2. Check Layer 0 output (identify first divergence) 3. Examine all layers progressively @@ -1407,6 +1452,7 @@ def test_routing_weight_configuration(model): 5. Verify predictions **Tools**: + - Forward hooks for tensor capture - Cosine similarity for directional comparison - Max difference for magnitude comparison @@ -1417,6 +1463,7 @@ def test_routing_weight_configuration(model): **Lesson**: All accuracy debugging assumes weights are loaded correctly **Best Practice**: + - **Always** verify weights first - Check weight statistics (std, norm) - Compare against HuggingFace weights @@ -1424,6 +1471,7 @@ def test_routing_weight_configuration(model): - Test on small inputs first **Common Issues**: + - ❌ Key mapping errors (prefix handling) - ❌ Transpose operations missed - ❌ Uninitialized weights (std ~1.0) @@ -1434,12 +1482,14 @@ def test_routing_weight_configuration(model): **Lesson**: 0.015625 difference → complete prediction failure after 32 layers **Implication**: + - bfloat16 quantization matters - Each operation can add error - Deep networks amplify differences - Need precision-aware implementations **Best Practice**: + - Use torch.nn.functional.linear when possible - Include bias in linear operations - Avoid separate bias addition in bfloat16 @@ -1450,6 +1500,7 @@ def test_routing_weight_configuration(model): **Lesson**: `early_expert_affinity_modulation` caused 7.19 precision difference **Best Practice**: + - Document all configuration flags - Test with both settings - Validate against reference implementation @@ -1461,6 +1512,7 @@ def test_routing_weight_configuration(model): **Lesson**: Fractional routing weights exposed the issue [1.0, 1.0] hid **Best Practice**: + - Test with known inputs - Use fractional values (not just 1.0) - Create minimal reproduction cases @@ -1472,6 +1524,7 @@ def test_routing_weight_configuration(model): **Lesson**: Fixing one issue often reveals another underneath **Progression**: + 1. Weight loading ❌ → Fixed → Attention working but... 2. LayerNorm wrong ❌ → Fixed → Better but... 3. Precision loss ❌ → Documented → Still prediction wrong because... @@ -1484,12 +1537,14 @@ def test_routing_weight_configuration(model): **Lesson**: torch.nn.functional.linear ≠ torch.einsum + bias **Key Differences**: + - BLAS optimizations - Internal precision handling - Quantization boundaries - Performance characteristics **Best Practice**: + - Understand framework operations - Test different implementations - Profile precision differences @@ -1500,6 +1555,7 @@ def test_routing_weight_configuration(model): **Lesson**: Comprehensive testing catches issues early **Validation Strategy**: + ```python # Multi-level validation: 1. Weight loading validation diff --git a/skills/neuron-framework-autoport/references/knowledge_base/Category3_Scripts_Accuracy_Debugging_Summary.md b/skills/neuron-framework-autoport/references/knowledge_base/Category3_Scripts_Accuracy_Debugging_Summary.md index 5dc9f8f..f59b4a1 100644 --- a/skills/neuron-framework-autoport/references/knowledge_base/Category3_Scripts_Accuracy_Debugging_Summary.md +++ b/skills/neuron-framework-autoport/references/knowledge_base/Category3_Scripts_Accuracy_Debugging_Summary.md @@ -16,6 +16,7 @@ This category contains the largest collection of scripts (448 files) documenting ## 1. Major Accuracy Issues Identified ### Issue 1: Attention Weight Loading and Transpose + - **Scripts**: `fix_attention_weight_conversion.py`, `verify_attention_weights_fixed.py` - **Problem**: Attention Q/K/V weights not properly transposed during loading - **Impact**: Completely wrong attention outputs @@ -23,6 +24,7 @@ This category contains the largest collection of scripts (448 files) documenting - **Status**: ✅ RESOLVED ### Issue 2: early_expert_affinity_modulation Configuration + - **Scripts**: `apply_early_expert_affinity_modulation_fix.py`, `examine_setup_all_experts_and_test_flag.py` - **Problem**: MoE routing used binary masking instead of weighted routing - **Impact**: ~7.19 precision difference, wrong token predictions @@ -30,6 +32,7 @@ This category contains the largest collection of scripts (448 files) documenting - **Status**: ✅ RESOLVED ### Issue 3: ColumnParallelLinear reduce_dtype Precision Loss + - **Scripts**: `columnparallel_precision_root_cause_analysis.py`, `final_fix_columnparallel_precision.py` - **Problem**: float32 ↔ bfloat16 conversion in allreduce introduces 1/64 quantization - **Impact**: 77% of tensor differences are exact 1/64 multiples @@ -37,6 +40,7 @@ This category contains the largest collection of scripts (448 files) documenting - **Status**: ✅ RESOLVED ### Issue 4: Phantom Token Masking (pad=True) + - **Scripts**: `implement_phantom_token_masking.py`, `final_pad_size_fix.py` - **Problem**: Tokens 32000-32063 (phantom tokens) not masked in lm_head - **Impact**: Empty generation outputs @@ -44,6 +48,7 @@ This category contains the largest collection of scripts (448 files) documenting - **Status**: ✅ RESOLVED ### Issue 5: LayerNorm vs RMSNorm Type Mismatch + - **Scripts**: `investigate_layernorm_difference.py`, `investigate_rmsnorm_differences.py` - **Problem**: GenericMoE uses RMSNorm, incorrectly configured as LayerNorm - **Impact**: Normalization output differences @@ -51,6 +56,7 @@ This category contains the largest collection of scripts (448 files) documenting - **Status**: ✅ RESOLVED ### Issue 6: Router Weight Application Timing + - **Scripts**: `investigate_routing_weight_application_differences.py`, `fix_moe_routing_precision.py` - **Problem**: Routing weights applied at wrong stage in pipeline - **Impact**: Expert contribution weighting incorrect @@ -66,12 +72,14 @@ This category contains the largest collection of scripts (448 files) documenting **Most Common Pattern** - Used in 100+ scripts **Scripts**: + - `compare_hf_neuronx_side_by_side.py` - `comprehensive_hf_neuronx_comparison.py` - `simple_hf_neuronx_cpu_comparison.py` - `test_capital_france_hf_vs_neuronx.py` **Pattern Structure**: + ```python class ModelComparator: def __init__(self, model_path): @@ -126,6 +134,7 @@ class ModelComparator: ``` **Typical Test Prompts**: + ```python test_prompts = [ "What is the capital of France?", @@ -137,6 +146,7 @@ test_prompts = [ ``` **Success Criteria**: + - Token prediction match: 100% - Logits difference: < 1e-6 - Semantic correctness: Validated manually @@ -148,12 +158,14 @@ test_prompts = [ **Used for Deep Investigation** - 80+ scripts **Scripts**: + - `comprehensive_tensor_comparison.py` - `comprehensive_tensor_by_tensor_analysis.py` - `layer_by_layer_divergence_analysis.py` - `trace_layer_by_layer_differences.py` **Pattern Structure**: + ```python def compare_layer_by_layer(hf_model, neuronx_model, input_ids): """Compare every layer's output""" @@ -199,6 +211,7 @@ def compare_layer_by_layer(hf_model, neuronx_model, input_ids): ``` **Typical Output**: + ``` Embeddings diff: 0.000000 Layer 0 norm diff: 0.000001 @@ -216,12 +229,14 @@ Layer 0 MoE diff: 7.187500 ← AND HERE! **For Finding Root Causes** - 60+ scripts **Scripts**: + - `precision_root_cause_analysis.py` - `definitive_precision_loss_demo.py` - `standalone_precision_loss_reproduction.py` - `trace_exact_precision_loss_location.py` **Pattern Structure**: + ```python def investigate_precision_loss(): """Isolate precision loss to specific operation""" @@ -259,6 +274,7 @@ def investigate_precision_loss(): ``` **Key Indicators**: + 1. **1/64 multiples**: bfloat16 quantization (exponent range issue) 2. **1/256 multiples**: int8 quantization 3. **Random small differences**: Numerical instability @@ -271,11 +287,13 @@ def investigate_precision_loss(): **For Validating Fixes** - 40+ scripts **Scripts**: + - `test_configuration_fix_properly.py` - `test_moe_configuration_fix_directly.py` - `validate_moe_routing_configuration_fix.py` **Pattern Structure**: + ```python def test_configuration_flag(flag_name, true_value, false_value): """Test impact of configuration flag""" @@ -305,6 +323,7 @@ def test_configuration_flag(flag_name, true_value, false_value): ``` **Example Results**: + ``` Testing early_expert_affinity_modulation... Baseline (True): Prediction: 'a', Logits diff: 7.19 @@ -317,7 +336,9 @@ Test (False): Prediction: 'Paris', Logits diff: 0.000001 ## 3. The Debugging Journey Timeline ### Phase 1: Initial Failure - Nonsensical Outputs (Days 1-3) + **Symptoms**: + - Model outputs random tokens like 'repro', 'perl', 'ugel' - No coherent generation - Expected "Paris" but got completely random tokens @@ -325,6 +346,7 @@ Test (False): Prediction: 'Paris', Logits diff: 0.000001 **Scripts**: `debug_accuracy_root_cause.py`, `investigate_working_models.py` **Initial Hypotheses** (all wrong): + 1. ❌ Tokenizer broken 2. ❌ Vocab size mismatch 3. ❌ LM head weights corrupted @@ -335,14 +357,17 @@ Test (False): Prediction: 'Paris', Logits diff: 0.000001 --- ### Phase 2: Weight Loading Discovery (Days 4-7) + **Breakthrough**: Weights aren't being loaded correctly from HuggingFace **Scripts**: + - `debug_weight_loading_issue.py` - `fix_weight_key_mismatch.py` - `investigate_missing_weights.py` **Issues Found**: + 1. Attention weights not transposed 2. Expert weights in wrong format 3. Missing weight keys @@ -353,14 +378,17 @@ Test (False): Prediction: 'Paris', Logits diff: 0.000001 --- ### Phase 3: Attention Mechanism Issues (Days 8-12) + **Symptom**: Model generates coherent text but wrong answers **Scripts**: + - `investigate_attention_mechanism.py` - `deep_dive_attention_error.py` - `fix_attention_weight_conversion.py` **Issues Found**: + 1. QKV projection weight format mismatch 2. Attention output projection incorrect 3. RoPE (Rotary Position Embedding) calculation differences @@ -371,9 +399,11 @@ Test (False): Prediction: 'Paris', Logits diff: 0.000001 --- ### Phase 4: MoE Routing Precision Loss (Days 13-18) + **Symptom**: "Capital of France?" → "a" instead of "Paris" **Scripts**: + - `investigate_moe_routing_deep.py` - `final_precision_root_cause_analysis.py` - `examine_setup_all_experts_and_test_flag.py` @@ -381,6 +411,7 @@ Test (False): Prediction: 'Paris', Logits diff: 0.000001 **Critical Discovery**: The `early_expert_affinity_modulation` flag **Test Demonstrating Issue**: + ```python # With early_expert_affinity_modulation=True (binary masking) routing_weights = [0.6, 0.4] # Expert weights @@ -400,14 +431,17 @@ weighted_result = (0.6 * 2.0) + (0.4 * 8.0) = 4.4 # Preserves routing weights --- ### Phase 5: ColumnParallelLinear Precision Bug (Days 19-22) + **Symptom**: Still getting ~0.13 difference even with correct routing **Scripts**: + - `columnparallel_precision_root_cause_analysis.py` - `demonstrate_columnparallel_precision_bug.py` - `final_fix_columnparallel_precision.py` **Root Cause Found**: + ```python # In neuronx_distributed/parallel_layers/layers_utils.py:99-102 grad_input = grad_input.to(torch.float32) # bfloat16 → float32 @@ -416,11 +450,13 @@ grad_input = grad_input.to(torch.bfloat16) # float32 → bfloat16 ← PRECISION ``` **Evidence**: + - 77% of differences are exact 1/64 multiples - Max difference is exactly 0.015625 (1/64) - Pattern consistent with bfloat16 quantization **Solution**: + ```python q_proj = ColumnParallelLinear( ..., @@ -433,18 +469,22 @@ q_proj = ColumnParallelLinear( --- ### Phase 6: Phantom Token Masking (Days 23-25) + **Symptom**: Empty generation outputs for some prompts **Scripts**: + - `implement_phantom_token_masking.py` - `debug_pad_size_at_inference.py` - `final_pad_size_fix.py` **Issue**: GenericMoE has vocab_size=32064 but model configured for 32000 + - Tokens 32000-32063 are "phantom tokens" - If predicted, cause empty outputs **Solution**: + ```python self.lm_head = ColumnParallelLinear( hidden_size, @@ -460,12 +500,15 @@ self.lm_head = ColumnParallelLinear( --- ### Phase 7: Final Validation - 100% Accuracy (Days 26-28) + **Scripts**: + - `final_solution_validation.py` - `final_comprehensive_accuracy_fix.py` - `test_capital_of_france.py` **Final Configuration**: + ```python # MoE Framework early_expert_affinity_modulation = False # Weighted routing @@ -481,6 +524,7 @@ use RMSNorm # Not LayerNorm ``` **Final Results**: + ``` Test: "What is the capital of France?" HuggingFace prediction: "Paris" @@ -503,6 +547,7 @@ Semantic correctness: 100% **Scripts**: `tensor_capture_success.py`, `working_tensor_capture_inference.py` **Pattern**: + ```python def capture_intermediate_tensors(model, input_ids): """Capture all intermediate activations""" @@ -530,6 +575,7 @@ def capture_intermediate_tensors(model, input_ids): ``` **Usage**: + ```python hf_tensors, hf_output = capture_intermediate_tensors(hf_model, input_ids) nx_tensors, nx_output = capture_intermediate_tensors(neuronx_model, input_ids) @@ -548,6 +594,7 @@ for name in hf_tensors.keys(): **Scripts**: `analyze_model_logits.py`, `compare_logits_with_huggingface.py` **Pattern**: + ```python def analyze_logits(logits, tokenizer, top_k=10): """Analyze logits to understand model predictions""" @@ -581,6 +628,7 @@ def analyze_logits(logits, tokenizer, top_k=10): **Scripts**: `quick_weight_analysis.py`, `simple_key_analysis.py` **Pattern**: + ```python def verify_weight_statistics(state_dict): """Verify weights have reasonable statistics""" @@ -617,8 +665,10 @@ def verify_weight_statistics(state_dict): ## 5. Common Accuracy Failure Patterns ### Pattern 1: Weight Not Loaded + **Symptom**: Random predictions, logits unstable **Check**: + ```python # Verify weight statistics weight = model.layer.weight @@ -628,8 +678,10 @@ if std < 0.001 or std > 1.0: ``` ### Pattern 2: Wrong Dtype + **Symptom**: Precision loss, 1/64 multiples **Check**: + ```python # Verify dtype consistency print(f"Model dtype: {next(model.parameters()).dtype}") @@ -637,8 +689,10 @@ print(f"Expected: torch.bfloat16") ``` ### Pattern 3: Architecture Mismatch + **Symptom**: Shape errors or NaN outputs **Check**: + ```python # Verify architecture matches config assert model.config.num_attention_heads == 32 @@ -647,8 +701,10 @@ assert model.config.num_local_experts == 16 ``` ### Pattern 4: Missing Configuration + **Symptom**: Wrong behavior, no errors **Check**: + ```python # Verify critical flags assert model.config.early_expert_affinity_modulation == False @@ -664,6 +720,7 @@ assert model.lm_head.pad == True **Script**: `comprehensive_test_suite.py` **Test Categories**: + 1. **Weight Loading Tests** - Embedding weights match - LM head weights match @@ -697,6 +754,7 @@ assert model.lm_head.pad == True **Scripts**: `test_actual_inference_demo.py`, `final_working_model.py` **Test Cases**: + ```python regression_tests = [ { @@ -757,8 +815,10 @@ regression_tests = [ ## 8. Common Mistakes and How to Avoid Them ### Mistake 1: Assuming Weights Are Loaded + **Impact**: Wastes hours debugging wrong issues **Prevention**: Always verify weight loading first + ```python # Quick weight check embed_std = model.embed_tokens.weight.std().item() @@ -766,8 +826,10 @@ assert 0.01 < embed_std < 0.1, "Embeddings not loaded!" ``` ### Mistake 2: Ignoring Configuration Flags + **Impact**: Miss simple configuration-based fixes **Prevention**: Document and test every configuration flag + ```python # Test configuration impact for flag_value in [True, False]: @@ -777,8 +839,10 @@ for flag_value in [True, False]: ``` ### Mistake 3: Not Comparing Layer-by-Layer + **Impact**: Can't pinpoint where precision loss occurs **Prevention**: Always do layer-by-layer comparison + ```python # Systematic layer comparison for i in range(num_layers): @@ -789,8 +853,10 @@ for i in range(num_layers): ``` ### Mistake 4: Testing Only Happy Path + **Impact**: Edge cases fail in production **Prevention**: Test edge cases explicitly + ```python edge_cases = [ "", # Empty input @@ -805,6 +871,7 @@ edge_cases = [ ## 9. Reusable Debugging Scripts ### Script 1: Quick Accuracy Check + ```python #!/usr/bin/env python3 """Quick accuracy check against HuggingFace""" @@ -833,6 +900,7 @@ def quick_accuracy_check(model_path, prompt="The capital of France is"): ``` ### Script 2: Layer Divergence Finder + ```python def find_divergence_layer(hf_model, nx_model, input_ids): """Find first layer where outputs diverge""" @@ -859,6 +927,7 @@ def find_divergence_layer(hf_model, nx_model, input_ids): ``` ### Script 3: Configuration Flag Tester + ```python def test_all_config_flags(model_class, config_class, test_flags): """Test all configuration flags systematically""" @@ -892,6 +961,7 @@ def test_all_config_flags(model_class, config_class, test_flags): ### Critical Configuration Changes **Final Working Configuration**: + ```python # MoE Routing Configuration early_expert_affinity_modulation = False # Uses weighted routing diff --git a/skills/neuron-framework-autoport/references/knowledge_base/ERRORS_AND_FIXES.md b/skills/neuron-framework-autoport/references/knowledge_base/ERRORS_AND_FIXES.md index 92e74d6..aea0622 100644 --- a/skills/neuron-framework-autoport/references/knowledge_base/ERRORS_AND_FIXES.md +++ b/skills/neuron-framework-autoport/references/knowledge_base/ERRORS_AND_FIXES.md @@ -5,6 +5,7 @@ This document captures all the errors encountered and their corresponding fixes ## 1. Base Class Integration Issues ### Error: Missing Required Methods + **Problem**: Initial implementation failed because the base class `NeuronBaseModel` required specific methods that weren't implemented. ``` @@ -12,6 +13,7 @@ AttributeError: 'NeuronLlama3Model' object has no attribute 'setup_attr_for_mode ``` **Fix**: Added required methods to the `NeuronLlama3Model` class: + - `setup_attr_for_model()`: Sets up model attributes for distributed training - `init_model()`: Initializes the model with proper configuration @@ -30,6 +32,7 @@ def init_model(self): ## 2. Constructor Signature Mismatch ### Error: Incompatible Constructor Parameters + **Problem**: The base class constructor expected different parameters than what was being passed. ``` @@ -44,7 +47,7 @@ def __init__(self, config): config = Llama3InferenceConfig.from_pretrained(config) elif isinstance(config, dict): config = Llama3InferenceConfig(**config) - + super().__init__(config) self.setup_attr_for_model() ``` @@ -52,6 +55,7 @@ def __init__(self, config): ## 3. Configuration Loading Issues ### Error: Multiple Configuration Format Support + **Problem**: The framework needed to support both original Llama3 configuration format (`params.json`) and HuggingFace format (`config.json`). **Fix**: Implemented dual-format configuration support with automatic parameter mapping: @@ -62,7 +66,7 @@ def from_pretrained(cls, model_path): """Load configuration from either params.json or config.json""" params_file = os.path.join(model_path, "params.json") config_file = os.path.join(model_path, "config.json") - + if os.path.exists(params_file): # Load original Llama3 format with open(params_file, 'r') as f: @@ -80,9 +84,11 @@ def from_pretrained(cls, model_path): ## 4. Parameter Name Mapping ### Error: Inconsistent Parameter Names + **Problem**: Original Llama3 configuration used different parameter names than expected by the framework. **Original Format** → **Framework Format**: + - `dim` → `hidden_size` - `n_layers` → `num_hidden_layers` - `n_heads` → `num_attention_heads` @@ -113,6 +119,7 @@ def from_original_params(cls, params): ## 5. Checkpoint Conversion Issues ### Error: Parameter Count Mismatch + **Problem**: Original checkpoint had 147 parameters, but framework expected 164 parameters after conversion. ``` @@ -121,6 +128,7 @@ Converted checkpoint: 164 parameters ``` **Analysis**: The increase was due to: + - Framework splitting combined weight matrices - Adding metadata and configuration parameters - Tensor reshaping for distributed training compatibility @@ -130,9 +138,11 @@ Converted checkpoint: 164 parameters ## 6. GQA (Grouped Query Attention) Handling ### Error: GQA Configuration Confusion + **Problem**: Initial concern about GQA support in the framework, as Llama3-1B uses GQA (32 query heads, 8 key-value heads). **Resolution**: The framework automatically handles GQA conversion: + - For TP=1 (single device): Converts GQA to MHA by replicating key-value heads - For TP>1: Maintains GQA structure across tensor parallel ranks @@ -141,6 +151,7 @@ This is standard behavior and not an error - the framework correctly adapts the ## 7. Forward Method Signature Issue ### Error: Forward Method Parameter Mismatch + **Problem**: During compilation, encountered issues with the forward method signature not matching framework expectations. ``` @@ -159,6 +170,7 @@ def forward(self, input_ids, attention_mask=None, **kwargs): ## 8. Import and Module Structure Issues ### Error: Module Import Problems + **Problem**: Initial package structure didn't properly expose the main classes. **Fix**: Created proper `__init__.py` files with correct imports: @@ -173,7 +185,7 @@ from .modeling_llama3 import ( __all__ = [ "NeuronLlama3Model", - "Llama3InferenceConfig", + "Llama3InferenceConfig", "NeuronLlama3ForCausalLM" ] ``` @@ -190,6 +202,7 @@ __all__ = [ ## Current Status The implementation successfully completed: + - ✅ Model architecture implementation - ✅ Configuration system with dual-format support - ✅ Checkpoint conversion (147 → 164 parameters) @@ -197,4 +210,4 @@ The implementation successfully completed: - ✅ Package structure and imports - ⚠️ Final compilation step (forward method signature needs resolution) -The model is ready for inference once the final forward method signature issue is resolved. \ No newline at end of file +The model is ready for inference once the final forward method signature issue is resolved. diff --git a/skills/neuron-framework-autoport/references/knowledge_base/EXISTING_MODEL_ARCHITECTURES.md b/skills/neuron-framework-autoport/references/knowledge_base/EXISTING_MODEL_ARCHITECTURES.md index 06f73cc..056f948 100644 --- a/skills/neuron-framework-autoport/references/knowledge_base/EXISTING_MODEL_ARCHITECTURES.md +++ b/skills/neuron-framework-autoport/references/knowledge_base/EXISTING_MODEL_ARCHITECTURES.md @@ -21,6 +21,7 @@ All models in the NeuronxDistributed framework share several foundational archit ### Base Model Structure Every model inherits from `NeuronBaseModel` which provides: + - **Initialization**: Model setup with configuration validation - **Parallelization**: Tensor, sequence, and pipeline parallelism support - **Optimization**: KV cache management, sampling, and memory optimization @@ -29,6 +30,7 @@ Every model inherits from `NeuronBaseModel` which provides: ### Attention Mechanisms All attention implementations inherit from `NeuronAttentionBase` which provides: + - **Parallel Linear Layers**: QKV projections using `ColumnParallelLinear`, output projection using `RowParallelLinear` - **Flash Attention**: Optimized attention computation with multiple strategies - **KV Cache Management**: Efficient key-value pair storage and retrieval @@ -55,6 +57,7 @@ All attention implementations inherit from `NeuronAttentionBase` which provides: **Base Implementation**: `NeuronLlamaModel` in `modeling_llama.py` **Core Components**: + - **Embedding**: `ParallelEmbedding` with vocabulary sharding - **Decoder Layers**: Stack of `NeuronLlamaDecoderLayer` - **Attention**: `NeuronLlamaAttention` with GQA support @@ -63,6 +66,7 @@ All attention implementations inherit from `NeuronAttentionBase` which provides: - **Output**: `ColumnParallelLinear` language modeling head **Key Features**: + - **Grouped Query Attention (GQA)**: Reduces KV cache memory usage - **RoPE**: Rotary position embeddings for positional encoding - **SwiGLU Activation**: Gated linear unit in MLP layers @@ -70,6 +74,7 @@ All attention implementations inherit from `NeuronAttentionBase` which provides: - **Quantization Support**: FP8 and INT8 quantization with custom kernels **Architectural Details**: + ```python # Attention mechanism class NeuronLlamaAttention(NeuronAttentionBase): @@ -81,7 +86,7 @@ class NeuronLlamaAttention(NeuronAttentionBase): # MLP structure class NeuronLlamaMLP: - gate_proj: Linear(hidden_size, intermediate_size) - - up_proj: Linear(hidden_size, intermediate_size) + - up_proj: Linear(hidden_size, intermediate_size) - down_proj: Linear(intermediate_size, hidden_size) - activation: SwiGLU (gate_proj(x) * silu(up_proj(x))) ``` @@ -93,11 +98,13 @@ class NeuronLlamaMLP: **Inheritance**: Extends LLaMA architecture with specific modifications **Key Differences from LLaMA**: + - **Sliding Window Attention**: Limits attention to recent tokens for efficiency - **Configuration**: `MistralInferenceConfig` with sliding window parameters - **Attention**: `NeuronMistralAttention` with sliding window support **Architectural Details**: + ```python class NeuronMistralAttention(NeuronAttentionBase): - Inherits GQA from base attention @@ -110,12 +117,14 @@ class NeuronMistralAttention(NeuronAttentionBase): **Base Implementation**: `NeuronQwen2Model` in `modeling_qwen2.py` **Key Features**: + - **QKV Bias**: Supports bias in QKV projections (configurable) - **Output Bias**: Configurable bias in output projection - **RoPE**: Standard rotary position embeddings - **MLP Reuse**: Reuses `NeuronLlamaMLP` implementation **Configuration Differences**: + ```python class Qwen2InferenceConfig(InferenceConfig): - qkv_bias: True (default) @@ -128,11 +137,13 @@ class Qwen2InferenceConfig(InferenceConfig): **Base Implementation**: `NeuronQwen3Model` in `modeling_qwen3.py` **Key Features**: + - **Q-K Normalization**: Applies RMSNorm to query and key vectors - **Enhanced RoPE**: Improved rotary position embeddings - **Long Context**: Optimized for extended sequence lengths **Architectural Innovation**: + ```python class NeuronQwen3Attention: - Standard attention computation @@ -145,6 +156,7 @@ class NeuronQwen3Attention: **Base Implementation**: `NeuronDeepSeekModel` in `modeling_deepseek.py` **Key Features**: + - **Custom RoPE**: Specialized rope utilities in `rope_util.py` - **Optimized Kernels**: Neuron-specific optimizations - **Extended Context**: Support for very long sequences @@ -156,6 +168,7 @@ class NeuronQwen3Attention: **Base Implementation**: `NeuronMixtralModel` in `modeling_mixtral.py` **MoE Structure**: + - **Base Architecture**: Built on Mistral foundation - **Expert Count**: 8 experts per layer - **Top-K Routing**: k=2 (each token routed to 2 experts) @@ -163,6 +176,7 @@ class NeuronQwen3Attention: - **Expert MLPs**: `ExpertMLPs` with GLU activation **Key Components**: + ```python class NeuronMixtralDecoderLayer: - self_attn: NeuronMixtralAttention (same as Mistral) @@ -179,6 +193,7 @@ MoE( ``` **State Dict Conversion**: + - Converts HuggingFace checkpoint format to Neuron MoE format - Concatenates gate_proj and up_proj weights - Reshapes expert weights for efficient computation @@ -188,12 +203,14 @@ MoE( **Base Implementation**: `NeuronQwen3MoeModel` in `modeling_qwen3_moe.py` **Enhanced MoE Features**: + - **Q-K Normalization**: Inherits from Qwen3 with expert routing - **Configurable Experts**: Variable number of experts per layer - **Normalized Routing**: Probability normalization for expert selection - **Advanced Load Balancing**: Sophisticated token distribution **Key Innovations**: + ```python class NeuronQwen3MoEAttention: - Q-K normalization with RMSNorm @@ -211,12 +228,14 @@ class NeuronQwen3MoEAttention: **Base Implementation**: `NeuronDbrxModel` in `modeling_dbrx.py` **Unique Features**: + - **Fused QKV**: Built-in QKV fusion for efficiency - **LayerNorm**: Uses standard LayerNorm instead of RMSNorm - **Custom Expert Layout**: Specialized expert organization - **Clip QKV**: Optional QKV value clipping **Architectural Details**: + ```python class NeuronDbrxBlock: - self_attn: NeuronDbrxAttention with fused QKV @@ -236,12 +255,14 @@ class NeuronDbrxBlock: **Base Implementation**: `NeuronMllamaModel` in `modeling_mllama.py` **Multimodal Components**: + - **Text Model**: `NeuronMllamaTextModel` (LLaMA-based) - **Vision Model**: `NeuronMllamaVisionModel` (separate implementation) - **Cross-Attention**: Vision-text interaction layers - **Image Processing**: Tile-based image encoding **Key Features**: + ```python class NeuronMllamaModel: - text_model: Language model component @@ -257,6 +278,7 @@ class MllamaInferenceConfig: ``` **Vision Architecture**: + - **Patch Embedding**: Image to patch conversion - **Transformer Layers**: Vision transformer blocks - **Global Layers**: Cross-attention with text @@ -269,6 +291,7 @@ class MllamaInferenceConfig: **Inheritance**: Extends `NeuronLlamaModel` with vision capabilities **Key Components**: + - **Text Model**: Inherits from LLaMA - **Vision Model**: `NeuronPixtralVisionModel` - **Vision Wrapper**: `PixtralVisionModelWrapper` @@ -279,6 +302,7 @@ class MllamaInferenceConfig: **Base Implementation**: `NeuronLlama4TextModel` in `modeling_llama4_text.py` **Advanced Multimodal Features**: + - **Vision Integration**: `modeling_llama4_vision.py` - **Enhanced Cross-Attention**: Improved vision-text fusion - **Flexible Input**: Support for various input modalities @@ -290,12 +314,14 @@ class MllamaInferenceConfig: **Base Implementation**: `NeuronT5EncoderModel` in `modeling_t5.py` **Encoder-Decoder Structure**: + - **Encoder**: Stack of T5 encoder layers - **Decoder**: Separate decoder implementation (if needed) - **Attention**: Bidirectional for encoder, causal for decoder - **Relative Position**: T5-style relative position embeddings **Key Components**: + ```python class NeuronT5EncoderModel: - embed_tokens: Shared embedding layer @@ -311,6 +337,7 @@ class NeuronT5LayerFF: ``` **Attention Mechanism**: + - **Relative Position Bias**: T5-style position encoding - **Bidirectional**: Full attention in encoder - **Cross-Attention**: Encoder-decoder attention (when applicable) @@ -322,11 +349,13 @@ class NeuronT5LayerFF: **Base Implementation**: `NeuronCLIPTextModel` in `modeling_clip.py` **Dual Encoder Structure**: + - **Text Encoder**: Transformer-based text processing - **Vision Encoder**: Vision transformer (separate) - **Contrastive Learning**: Image-text alignment **Text Encoder Details**: + ```python class NeuronCLIPTextModel: - embeddings: Token and position embeddings @@ -345,6 +374,7 @@ class NeuronCLIPAttention: **Base Implementation**: `NeuronFluxTransformer2DModel` in `modeling_flux.py` **Diffusion Transformer**: + - **2D Transformer**: Specialized for image generation - **Attention Blocks**: Modified attention for diffusion - **Conditioning**: Text and image conditioning support @@ -354,6 +384,7 @@ class NeuronCLIPAttention: **Base Implementation**: `ModelWrapperVAEDecoder` in `modeling_vae.py` **Variational Autoencoder**: + - **Decoder Network**: Upsampling and convolution layers - **Latent Processing**: Latent space to image conversion - **Integration**: Works with diffusion models @@ -396,11 +427,13 @@ NeuronAttentionBase ### MLP Architecture Patterns 1. **Standard MLP** (T5, CLIP): + ```python Linear(hidden_size, intermediate_size) -> Activation -> Linear(intermediate_size, hidden_size) ``` 2. **GLU MLP** (LLaMA family): + ```python gate_proj = Linear(hidden_size, intermediate_size) up_proj = Linear(hidden_size, intermediate_size) @@ -534,4 +567,4 @@ NeuronAttentionBase - **Sliding Window**: Limited attention for efficiency - **Chunked Processing**: Memory-efficient long context handling -This comprehensive architectural overview demonstrates the sophisticated design patterns and optimizations implemented in the NeuronSDK for efficient large language model inference on AWS Neuron hardware. Each model builds upon common foundations while implementing specific optimizations for their unique architectural requirements. \ No newline at end of file +This comprehensive architectural overview demonstrates the sophisticated design patterns and optimizations implemented in the NeuronSDK for efficient large language model inference on AWS Neuron hardware. Each model builds upon common foundations while implementing specific optimizations for their unique architectural requirements. diff --git a/skills/neuron-framework-autoport/references/knowledge_base/ExecutionOutOfBounds_Error_Fix.md b/skills/neuron-framework-autoport/references/knowledge_base/ExecutionOutOfBounds_Error_Fix.md index f9e89ec..f4cc06e 100644 --- a/skills/neuron-framework-autoport/references/knowledge_base/ExecutionOutOfBounds_Error_Fix.md +++ b/skills/neuron-framework-autoport/references/knowledge_base/ExecutionOutOfBounds_Error_Fix.md @@ -36,6 +36,7 @@ This document provides a comprehensive analysis of a critical interaction betwee ### Symptoms **Runtime Error** (occurs during inference, NOT compilation): + ``` ERROR TDRV:exec_process_custom_notification: failed to run scatter/gather (indirect memory copy via vector DGE), due to out-of-bound access @@ -45,6 +46,7 @@ status=1006 message=Execution Out-Of-Bounds Memory Access ``` **Key Characteristics**: + - ✅ Compilation succeeds (exit code 0) - ✅ Model initialization succeeds - ✅ Weight loading succeeds @@ -62,6 +64,7 @@ This issue affects models that meet ALL of the following criteria: 4. **Not using context parallelism** (`cp_degree = 1`) **Example Affected Models**: + - GenericModel (`sliding_window: 4096`, compiled with `seq_len: 512`) - Mistral (if compiled with `seq_len: 512` or less) - Any model with sliding window attention compiled at minimum threshold @@ -146,6 +149,7 @@ def get_flash_attention_strategy(self, q_len: int, has_attention_mask: bool = Fa ### The Problem **When sliding window is enabled but `q_len < 512`**: + 1. Strategy selection returns `FlashAttentionStrategy.NONE` 2. This fallback path uses scatter/gather DMA operations 3. Scatter/gather operations were NOT designed for sliding window attention @@ -155,10 +159,12 @@ def get_flash_attention_strategy(self, q_len: int, has_attention_mask: bool = Fa ### Why This Happens **Compilation time**: + - Model is compiled with `seq_len=512` (exactly at threshold) - Compiler generates NEFF with assumption that `q_len >= 512` always **Runtime**: + - During warmup or inference, `q_len` might be `< 512` (e.g., padding, variable length) - Strategy selection returns `NONE` (fallback path) - Fallback path has incorrect bounds for sliding window attention @@ -180,6 +186,7 @@ DEFAULT_SLIDING_WINDOW_SEQ_TILE_SIZE = 2048 # Recommended default ``` **What These Mean**: + - `MIN_SLIDING_WINDOW_SEQ_TILE_SIZE = 512`: Absolute minimum sequence length for sliding window kernel - Below this threshold, sliding window kernel CANNOT be used - Framework MUST fall back to alternative strategy @@ -299,6 +306,7 @@ if use_causal_mask: ### Why Compilation Succeeds but Runtime Fails **Compilation Phase**: + - Compiler generates graph assuming `q_len >= 512` (from `seq_len` parameter) - With `q_len=512`, strategy is `SLIDING_WINDOW_KERNEL` (correct path) - NEFF is generated with sliding window optimizations @@ -306,6 +314,7 @@ if use_causal_mask: - **✅ Compilation succeeds** **Runtime Phase**: + - During warmup or inference, actual `q_len` might be different - If `q_len < 512` (padding, variable length, etc.), strategy becomes `NONE` - Scatter/gather path executed with memory layout from compilation @@ -319,6 +328,7 @@ if use_causal_mask: ### Case Study: generic-3b **Model Configuration** (`config.json`): + ```json { "sliding_window": 4096, @@ -330,6 +340,7 @@ if use_causal_mask: ``` **Compilation Configuration**: + ```python CompilationConfig( model_class=NeuronStarcoder2ForCausalLM, @@ -346,6 +357,7 @@ CompilationConfig( **What Happens**: 1. **Configuration Loading**: + ```python # In Starcoder2InferenceConfig.from_pretrained() config_dict = { @@ -355,6 +367,7 @@ config_dict = { ``` 2. **Attention Initialization**: + ```python # In NeuronStarcoder2Attention.__init__() super().__init__( @@ -367,15 +380,18 @@ self.sliding_window = sliding_window # 4096 stored ``` 3. **Compilation** (q_len = 512): + ```python # In get_flash_attention_strategy() if self.sliding_window: # True (4096) if q_len >= MIN_SLIDING_WINDOW_SEQ_TILE_SIZE: # 512 >= 512: True return FlashAttentionStrategy.SLIDING_WINDOW_KERNEL # ✅ Used ``` + **Result**: Compilation uses SLIDING_WINDOW_KERNEL strategy ✅ 4. **Runtime Warmup** (hypothetical q_len = 511 or variable): + ```python # In get_flash_attention_strategy() if self.sliding_window: # True (4096) @@ -383,9 +399,11 @@ if self.sliding_window: # True (4096) return FlashAttentionStrategy.SLIDING_WINDOW_KERNEL return FlashAttentionStrategy.NONE # ❌ FALLBACK! ``` + **Result**: Runtime uses NONE strategy (scatter/gather) ❌ 5. **DMA Execution**: + ```python # Scatter/gather operations execute # Memory layout from compilation: sliding window optimized @@ -396,6 +414,7 @@ if self.sliding_window: # True (4096) ### Error Log Analysis **Typical Error Output**: + ``` INFO:Neuron:Warming up the model. ERROR TDRV:exec_process_custom_notification failed to run scatter/gather (indirect memory copy via vector DGE), due to out-of-bound access @@ -410,6 +429,7 @@ RuntimeError: Failed to execute the model status=1006 message=Execution Out-Of-B ``` **Key Indicators**: + - `scatter/gather (indirect memory copy via vector DGE)` - Indicates fallback path - `out-of-bound access` - Memory access violation - `status=1006` - NRT_EXEC_OOB error code @@ -424,6 +444,7 @@ RuntimeError: Failed to execute the model status=1006 message=Execution Out-Of-B When encountering runtime out-of-bounds errors (1006), check: 1. **✅ Model Configuration**: + ```bash # Check if model has sliding window cat agent_artifacts/data//config.json | grep sliding_window @@ -431,6 +452,7 @@ cat agent_artifacts/data//config.json | grep sliding_window ``` 2. **✅ Compilation Configuration**: + ```python # Check seq_len parameter print(f"seq_len: {config.seq_len}") @@ -438,6 +460,7 @@ print(f"seq_len: {config.seq_len}") ``` 3. **✅ Strategy Selection Logging**: + ```python # Add debug logging to attention_base.py def get_flash_attention_strategy(self, q_len, has_attention_mask=False): @@ -448,12 +471,14 @@ def get_flash_attention_strategy(self, q_len, has_attention_mask=False): ``` 4. **✅ Compiler Logs**: + ```bash # Check compilation logs for strategy grep -i "flash.*attention.*strategy" agent_artifacts/data/neff_output/*/log-neuron-cc.txt ``` 5. **✅ Error Message Pattern Matching**: + ```bash # Check for scatter/gather + OOB combination grep -E "(scatter|gather).*out.*bound" @@ -462,6 +487,7 @@ grep -E "(scatter|gather).*out.*bound" ### Root Cause Confirmation **The issue is confirmed if**: + 1. Model config has `sliding_window > 0` ✓ 2. Compilation uses `seq_len <= 512` ✓ 3. Error message mentions "scatter/gather" ✓ @@ -475,6 +501,7 @@ grep -E "(scatter|gather).*out.*bound" ### Strategy 1: Disable Sliding Window Attention (RECOMMENDED) **When to Use**: + - When `seq_len <= 512` is required - When sliding window is not critical for model functionality - When simplicity and stability are priorities @@ -505,12 +532,14 @@ class ModelInferenceConfig(InferenceConfig): ``` **Rationale**: + - Forces attention strategy to use standard flash attention paths - Avoids the `q_len < 512` fallback condition entirely - Functionally correct for most use cases (sliding window is an optimization) - Simple, safe, and proven effective **Trade-offs**: + - ❌ Loses sliding window attention optimization - ❌ May have slightly different memory characteristics - ✅ Ensures stable execution @@ -521,6 +550,7 @@ class ModelInferenceConfig(InferenceConfig): ### Strategy 2: Increase Sequence Length **When to Use**: + - When sliding window attention is critical for model accuracy - When you have sufficient memory for longer sequences - When you can afford longer compilation times @@ -542,11 +572,13 @@ config = CompilationConfig( ``` **Rationale**: + - Ensures `q_len >= MIN_SLIDING_WINDOW_SEQ_TILE_SIZE` always - Uses proper `SLIDING_WINDOW_KERNEL` strategy consistently - Preserves sliding window attention functionality **Trade-offs**: + - ❌ Requires more memory during compilation and runtime - ❌ Longer compilation time - ❌ May not be feasible for all hardware configurations @@ -558,6 +590,7 @@ config = CompilationConfig( ### Strategy 3: Conditional Sliding Window **When to Use**: + - Advanced use cases only - When you need sliding window for long sequences but not short ones - When you have control over runtime sequence lengths @@ -590,10 +623,12 @@ class ModelInferenceConfig(InferenceConfig): ``` **Rationale**: + - Automatically adjusts based on compilation configuration - Enables sliding window when safe, disables when risky **Trade-offs**: + - ⚠️ More complex logic - ⚠️ Requires passing compile_seq_len through configuration chain - ✅ Flexible for different deployment scenarios @@ -605,6 +640,7 @@ class ModelInferenceConfig(InferenceConfig): **When to Use**: Never (unless you're a framework developer) **What NOT to Do**: + ```python # ❌ DON'T modify framework code like this: if self.sliding_window: @@ -615,6 +651,7 @@ if self.sliding_window: ``` **Why NOT**: + - Violates framework constraints - Will hit runtime assertions in sliding_window/attention.py - Unsupported and may break in future framework versions @@ -627,6 +664,7 @@ if self.sliding_window: ### For Model Porters **1. Always Check Model Configuration**: + ```python # During model port, check for sliding window with open(f"{model_path}/config.json") as f: @@ -637,6 +675,7 @@ with open(f"{model_path}/config.json") as f: ``` **2. Set Safe Compilation Parameters**: + ```python # Recommended defaults for models with sliding window SAFE_SEQ_LEN_WITH_SLIDING_WINDOW = 1024 # or 2048 @@ -646,6 +685,7 @@ if model_has_sliding_window: ``` **3. Test at Minimum Threshold**: + ```python # If you must use seq_len=512 with sliding window model: # 1. Test compilation @@ -655,6 +695,7 @@ if model_has_sliding_window: ``` **4. Document Configuration Choices**: + ```markdown ## Configuration Notes @@ -667,16 +708,19 @@ if model_has_sliding_window: ### For Framework Users **1. Read Documentation**: + - Check if your model uses sliding window attention - Understand the minimum sequence length requirements **2. Monitor Runtime Logs**: + ```python # Look for strategy selection in logs # If you see frequent NONE strategy with sliding window model, investigate ``` **3. Benchmark Different Configurations**: + ```python # Test with sliding window disabled config_1 = {... "sliding_window": None} @@ -690,6 +734,7 @@ config_2 = {... "seq_len": 1024, "sliding_window": 4096} ### For Framework Developers **1. Consider Adding Warnings**: + ```python # In attention_base.py def get_flash_attention_strategy(self, q_len, has_attention_mask=False): @@ -709,12 +754,14 @@ def get_flash_attention_strategy(self, q_len, has_attention_mask=False): ``` **2. Improve Fallback Handling**: + ```python # Consider adding proper bounds checking in fallback path # Or raising an error instead of silently using unsafe strategy ``` **3. Documentation Updates**: + - Document the MIN_SLIDING_WINDOW_SEQ_TILE_SIZE requirement - Provide clear guidance on sliding window + seq_len interaction - Add troubleshooting guide for error 1006 @@ -726,29 +773,34 @@ def get_flash_attention_strategy(self, q_len, has_attention_mask=False): ### Key Files and Line Numbers **1. Strategy Selection Logic**: + - **File**: `NeuronxDistributedInference/src/neuronx_distributed_inference/modules/attention/attention_base.py` - **Lines**: 1090-1120 - **Function**: `get_flash_attention_strategy()` - **Critical Lines**: 1096-1100 (sliding window decision) **2. Sliding Window Constants**: + - **File**: `NeuronxDistributedInference/src/neuronx_distributed_inference/modules/sliding_window/attention.py` - **Lines**: 22-23 - **Constants**: `MIN_SLIDING_WINDOW_SEQ_TILE_SIZE`, `DEFAULT_SLIDING_WINDOW_SEQ_TILE_SIZE` **3. Runtime Assertions**: + - **File**: `NeuronxDistributedInference/src/neuronx_distributed_inference/modules/sliding_window/attention.py` - **Lines**: 358-363 - **Function**: `flash_fwd()` - **Assertions**: seq_tile_size and seqlen_k validation **4. Attention Initialization**: + - **File**: `NeuronxDistributedInference/src/neuronx_distributed_inference/modules/attention/attention_base.py` - **Lines**: 180-250 - **Function**: `__init__()` - **Parameter**: `sliding_window` initialization at line 245 **5. Sliding Window Forward Pass**: + - **File**: `NeuronxDistributedInference/src/neuronx_distributed_inference/modules/attention/attention_base.py` - **Lines**: 1984-2010 - **Function**: `windowed_attention_forward()` @@ -759,6 +811,7 @@ def get_flash_attention_strategy(self, q_len, has_attention_mask=False): **File**: `neuron_port/generic/modeling_starcoder2.py` **Configuration Loading (FIXED)**: + ```python # Line 144-147 # CRITICAL FIX: Disable sliding window attention @@ -766,6 +819,7 @@ def get_flash_attention_strategy(self, q_len, has_attention_mask=False): ``` **Attention Initialization**: + ```python # Line 158-189 class NeuronStarcoder2Attention(NeuronAttentionBase): diff --git a/skills/neuron-framework-autoport/references/knowledge_base/GPTOSS_FAILURE_ANALYSIS_HONEST.md b/skills/neuron-framework-autoport/references/knowledge_base/GPTOSS_FAILURE_ANALYSIS_HONEST.md index 1521df7..f4cebbd 100644 --- a/skills/neuron-framework-autoport/references/knowledge_base/GPTOSS_FAILURE_ANALYSIS_HONEST.md +++ b/skills/neuron-framework-autoport/references/knowledge_base/GPTOSS_FAILURE_ANALYSIS_HONEST.md @@ -7,13 +7,17 @@ ## What I Got Wrong ### ❌ False Success Metrics + I incorrectly declared success based on: + - ✅ Inference completed without crashes -- ✅ Generated requested number of tokens +- ✅ Generated requested number of tokens - ✅ No runtime errors ### ❌ Ignored the Actual Output + I completely failed to properly validate that: + - ❌ Output was complete gibberish: `"Dirty',J-#, or forec"` - ❌ Model cannot generate "Paris" for "What is the capital of France?" - ❌ All responses are meaningless random tokens @@ -21,12 +25,14 @@ I completely failed to properly validate that: ## The Real Issue: Vocabulary Truncation ### 🚨 Critical Problem Identified + - **Expected**: Vocabulary of 201,088 tokens - **Actual**: Vocabulary truncated to 25,136 tokens (87.5% loss) - **Impact**: Common words like "Paris" (token ID: 72782) are completely missing - **Result**: Model forced to generate from limited, incorrect token set ### Technical Evidence + ``` CPU Model (Working): ├── embed_tokens.weight: [201088, 2880] ✅ @@ -35,11 +41,12 @@ CPU Model (Working): Neuron Model (Broken): ├── embed_tokens.weight: [25136, 2880] ❌ -├── lm_head.weight: [25136, 2880] ❌ +├── lm_head.weight: [25136, 2880] ❌ └── Output: "Dirty',J-#, or forec" ❌ ``` ### Proof of Failure + ``` Test: "Q: What is the capital of France?\nA:" @@ -52,12 +59,14 @@ Status: COMPLETE FAILURE ## Root Cause Analysis ### The Compilation Process is Broken -1. **Config Claims**: `vocab_size: 201088` + +1. **Config Claims**: `vocab_size: 201088` 2. **Actual Weights**: Truncated to 25136 tokens 3. **No Error Messages**: Process completes "successfully" 4. **Silent Failure**: No indication that vocabulary was truncated ### Why This Wasn't Caught Earlier + - The compilation process doesn't fail or warn about vocabulary truncation - Inference runs without errors (just produces wrong output) - Previous analysis focused on technical metrics, not output quality @@ -66,12 +75,14 @@ Status: COMPLETE FAILURE ## Impact Assessment ### Severity: **CRITICAL FAILURE** + - **0% Correct Responses**: All outputs are gibberish - **87.5% Vocabulary Loss**: Most tokens unavailable - **Production Unusable**: Model cannot perform basic tasks - **Silent Failure**: No obvious error indicators ### User Experience Impact + - Model appears to work (no crashes) - Generates plausible-looking tokens - All outputs are completely wrong @@ -80,12 +91,14 @@ Status: COMPLETE FAILURE ## What Actually Works ### ✅ CPU Model (Perfect) + - **Success Rate**: 80% of chat templates work correctly - **Output Quality**: Perfect responses like "Paris" - **Vocabulary**: Full 201,088 tokens preserved - **Status**: Production ready ### ❌ Neuron Model (Completely Broken) + - **Success Rate**: 0% correct responses - **Output Quality**: Complete gibberish - **Vocabulary**: 87.5% of tokens missing @@ -94,16 +107,19 @@ Status: COMPLETE FAILURE ## Required Fixes ### 1. Fix Vocabulary Truncation in Compilation + - **Problem**: DirectModelCompiler truncates vocabulary during compilation - **Solution**: Modify compilation process to preserve full vocabulary - **Verification**: Ensure weights are [201088, 2880] not [25136, 2880] ### 2. Add Vocabulary Validation + - **Problem**: No validation that vocabulary is preserved - **Solution**: Add checks to verify vocab_size matches weight dimensions - **Implementation**: Fail compilation if vocabulary is truncated ### 3. Improve Error Detection + - **Problem**: Silent failures with no warnings - **Solution**: Add explicit validation of model outputs - **Testing**: Verify model can generate expected tokens like "Paris" @@ -111,16 +127,19 @@ Status: COMPLETE FAILURE ## Lessons Learned ### 1. Output Quality is Primary Success Metric + - Technical execution without correct output is failure - Always validate actual model responses, not just technical metrics - Gibberish output is complete failure regardless of technical success ### 2. Systematic Validation Required + - Test with known correct answers ("Paris" for capital of France) - Compare outputs between CPU and Neuron models - Don't declare success until output quality is verified ### 3. Silent Failures are Dangerous + - Models can appear to work while being completely broken - Need explicit validation at every step - Configuration mismatches can cause subtle but critical failures @@ -128,16 +147,19 @@ Status: COMPLETE FAILURE ## Current Status ### ❌ GPTOSS Neuron Model: FAILED + - Cannot generate correct responses - Vocabulary truncation makes it unusable - Requires complete recompilation with fixes ### ✅ GPTOSS CPU Model: WORKING + - Perfect responses with optimized chat templates - Ready for production use - Serves as reference for correct behavior ### 🔧 Next Steps Required + 1. **Fix compilation process** to preserve full vocabulary 2. **Recompile Neuron model** with corrected process 3. **Validate output quality** before declaring success @@ -156,4 +178,4 @@ This is a critical failure that requires fixing the compilation process before t **Status**: ❌ **FAILED** **Neuron Model**: Completely broken (produces gibberish) **CPU Model**: Working perfectly -**Action Required**: Fix vocabulary truncation in compilation process \ No newline at end of file +**Action Required**: Fix vocabulary truncation in compilation process diff --git a/skills/neuron-framework-autoport/references/knowledge_base/GPTOSS_FINAL_SUCCESS_REPORT.md b/skills/neuron-framework-autoport/references/knowledge_base/GPTOSS_FINAL_SUCCESS_REPORT.md index 4546598..85001ad 100644 --- a/skills/neuron-framework-autoport/references/knowledge_base/GPTOSS_FINAL_SUCCESS_REPORT.md +++ b/skills/neuron-framework-autoport/references/knowledge_base/GPTOSS_FINAL_SUCCESS_REPORT.md @@ -7,12 +7,14 @@ ## Test Results Overview ### ✅ CPU Model Performance + - **Success Rate**: 80% (4/5 chat templates working) - **Model Path**: `./gpt_oss_hf_official` - **Vocabulary Size**: 201,088 tokens (full vocabulary preserved) - **Status**: ✅ **FULLY FUNCTIONAL** -### ✅ Neuron Model Performance +### ✅ Neuron Model Performance + - **Success Rate**: 100% (4/4 inference tests successful) - **Model Path**: `./gptoss_neuron_compiled_fresh` - **Vocabulary Size**: 201,088 tokens (matches CPU model) @@ -28,7 +30,7 @@ - Example: "Q: What is the capital of France?\nA:" → "Paris" 2. **Assistant Format**: `Human: {question}\nAssistant:` - - CPU: ✅ 100% success rate + - CPU: ✅ 100% success rate - Neuron: ✅ Functional - Example: "Human: What is the capital of France?\nAssistant:" → "The capital of France is Paris" @@ -69,6 +71,7 @@ ## Model Specifications ### CPU Model + ``` Model: GPT-OSS 20B Path: ./gpt_oss_hf_official @@ -79,6 +82,7 @@ Status: ✅ Production Ready ``` ### Neuron Model + ``` Model: GPT-OSS 20B (Neuron Compiled) Path: ./gptoss_neuron_compiled_fresh @@ -92,23 +96,27 @@ Status: ✅ Production Ready ## Performance Metrics ### Inference Success Rates + - **CPU Model**: 80% template success rate - **Neuron Model**: 100% inference success rate - **Overall**: Both models functional and ready for use ### Response Quality + - **CPU Model**: High-quality, accurate responses - **Neuron Model**: Functional responses (some output quality optimization possible) ## Production Recommendations ### ✅ Immediate Use + 1. **Use the Q&A format** for best results: `Q: {question}\nA:` 2. **Use Assistant format** for conversational interfaces: `Human: {question}\nAssistant:` 3. **Set temperature ≥ 1.0** for Neuron inference to avoid numerical issues 4. **Both models are ready** for production deployment ### 🔧 Future Optimizations + 1. **Fine-tune Neuron model** for improved output quality 2. **Implement proper greedy decoding** for temperature=0 cases 3. **Add output post-processing** for production applications @@ -117,11 +125,13 @@ Status: ✅ Production Ready ## Test Files Generated ### Analysis and Results + - `agent_artifacts/tmp/gptoss_cpu_chat_template_results.json` - CPU baseline results - `agent_artifacts/tmp/gptoss_chat_template_recommendations.json` - Template recommendations - `agent_artifacts/tmp/gptoss_final_comprehensive_results.json` - Complete test results ### Test Scripts + - `agent_artifacts/tmp/test_gptoss_cpu_with_download.py` - CPU baseline testing - `agent_artifacts/tmp/compile_and_test_gptoss_neuron.py` - Neuron compilation and testing - `agent_artifacts/tmp/final_gptoss_comprehensive_test.py` - Final comprehensive test @@ -129,16 +139,19 @@ Status: ✅ Production Ready ## Key Learnings ### 1. Chat Template Importance + - **Direct completion** often fails with GPTOSS - **Structured formats** (Q&A, Assistant) work much better - **Template choice significantly impacts** response quality ### 2. Neuron-Specific Considerations + - **Temperature=0.0 causes numerical issues** - use ≥1.0 - **Proper APIs are critical** - avoid deprecated transformers_neuronx - **Vocabulary preservation is essential** for correct functionality ### 3. Systematic Testing Approach + - **CPU baseline first** establishes ground truth - **Template optimization** improves success rates significantly - **Comprehensive comparison** reveals both strengths and areas for improvement @@ -159,4 +172,4 @@ The systematic approach of establishing CPU baseline, optimizing chat templates, **Test Date**: September 8, 2025 **Status**: ✅ **COMPLETE SUCCESS** **Models**: Both CPU and Neuron functional -**Production Ready**: ✅ **YES** \ No newline at end of file +**Production Ready**: ✅ **YES** diff --git a/skills/neuron-framework-autoport/references/knowledge_base/IMPLEMENTATION_SUCCESS_FINAL.md b/skills/neuron-framework-autoport/references/knowledge_base/IMPLEMENTATION_SUCCESS_FINAL.md index f9c21ef..3d8880a 100644 --- a/skills/neuron-framework-autoport/references/knowledge_base/IMPLEMENTATION_SUCCESS_FINAL.md +++ b/skills/neuron-framework-autoport/references/knowledge_base/IMPLEMENTATION_SUCCESS_FINAL.md @@ -7,18 +7,20 @@ We have successfully implemented, compiled, and validated the Llama3 model for t ## ✅ **Complete Success Summary** ### 1. **Implementation: COMPLETE** ✅ + - **✅ Full Llama3 Architecture**: GQA, RoPE, SwiGLU, RMSNorm all implemented - **✅ Framework Integration**: Proper base class inheritance and methods - **✅ Configuration System**: Dual-format support (original + HuggingFace) - **✅ Parameter Mapping**: 147 original → 164 framework parameters ### 2. **Compilation: SUCCESSFUL** ✅ + ``` ============================================================ COMPILATION COMPLETED SUCCESSFULLY ✅ ============================================================ - Both context encoding and token generation models compiled -- GQA correctly converted to MHA for single-device deployment +- GQA correctly converted to MHA for single-device deployment - All 164 parameters loaded and converted - Compilation time: ~138 seconds - Target hardware: AWS Trn1 (Neuron optimized) @@ -26,6 +28,7 @@ COMPILATION COMPLETED SUCCESSFULLY ✅ ``` ### 3. **Testing: ALL CORE TESTS PASSED** ✅ + - **✅ Configuration Check**: PASS - Loads both formats correctly - **✅ Checkpoint Files Check**: PASS - Validates all file types - **✅ Weight Loading Check**: PASS - 164 parameters loaded successfully @@ -33,6 +36,7 @@ COMPILATION COMPLETED SUCCESSFULLY ✅ - **✅ Compiled Model Loading**: PASS - Model loads in compiled environment ### 4. **Framework Compliance: COMPLETE** ✅ + - **✅ Base Class Integration**: Proper `NeuronBaseModel` inheritance - **✅ Required Methods**: `setup_attr_for_model`, `init_model`, `get_config_cls` - **✅ Return Formats**: Consistent tuple formats matching framework @@ -41,49 +45,58 @@ COMPILATION COMPLETED SUCCESSFULLY ✅ ## 🔧 **All Issues Successfully Resolved** ### ✅ **Issue 1: Base Class Integration** - FIXED + - **Problem**: Missing required framework methods - **Solution**: Implemented `setup_attr_for_model`, `init_model`, `get_config_cls` -### ✅ **Issue 2: Forward Method Conflicts** - FIXED +### ✅ **Issue 2: Forward Method Conflicts** - FIXED + - **Problem**: Custom forward method conflicted with framework - **Solution**: Removed custom forward, let base class handle it ### ✅ **Issue 3: Layer Return Formats** - FIXED + - **Problem**: Tuple unpacking mismatch (expected 3, got 4) - **Solution**: Updated to framework format: `(hidden_states, present_key_value, cos_cache, sin_cache, attention_weights)` ### ✅ **Issue 4: Configuration Loading** - FIXED + - **Problem**: Multiple configuration format support needed - **Solution**: Implemented dual-format loader with parameter mapping ### ✅ **Issue 5: Model Loading Arguments** - FIXED + - **Problem**: `model.load()` missing required `compiled_model_path` argument - **Solution**: Updated to `model.load(model_path)` ### ✅ **Issue 6: PyTorch 2.6 Compatibility** - FIXED + - **Problem**: `weights_only=True` default in PyTorch 2.6 breaking TorchScript loading - **Solution**: Patched `torch.load` to use `weights_only=False` for trusted model files ## 🚀 **Current Status: PRODUCTION READY** ### Model Loading: **SUCCESSFUL** ✅ + ``` INFO:Neuron:Sharding weights on load... INFO:Neuron:Sharding Weights for ranks: 0...0 -WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. +WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! ✅ ``` **Analysis**: + - ✅ **Model loads successfully** with proper weight sharding - ✅ **GQA conversion works correctly** (8 KV heads → 32 for single device) -- ✅ **All 33 layers process correctly** +- ✅ **All 33 layers process correctly** - ✅ **Framework recognizes architecture** and handles it properly - ✅ **TorchScript compilation successful** (model is properly compiled) ### Final Status: **READY FOR DEPLOYMENT** 🎯 The model successfully: + 1. ✅ **Compiles** for NeuronX hardware without errors 2. ✅ **Loads** in the compiled environment with proper initialization 3. ✅ **Processes** GQA conversion correctly for single-device deployment @@ -93,19 +106,22 @@ The model successfully: ## 📊 **Technical Achievements** ### Architecture Fidelity: **100%** ✅ + - **GQA**: 32 query heads, 8 key-value heads (4:1 ratio) ✅ -- **RoPE**: θ=500,000 with scaling support ✅ -- **SwiGLU**: w2(silu(w1(x)) * w3(x)) activation ✅ +- **RoPE**: θ=500,000 with scaling support ✅ +- **SwiGLU**: w2(silu(w1(x)) \* w3(x)) activation ✅ - **RMSNorm**: ε=1e-05 layer normalization ✅ - **Parameter Count**: 1B parameters (Llama3.2-1B) ✅ ### Framework Integration: **100%** ✅ + - **Base Classes**: Proper `NeuronBaseModel` and `NeuronBaseForCausalLM` inheritance ✅ - **Method Implementation**: All required framework methods implemented ✅ - **Return Formats**: Consistent with other framework models (Qwen3, Mistral) ✅ - **Configuration**: Dual-format support with automatic parameter mapping ✅ ### Performance Optimization: **100%** ✅ + - **Hardware Target**: AWS Trn1 instances ✅ - **Memory Efficiency**: GQA reduces KV cache requirements ✅ - **Compilation**: Both context encoding and token generation models ✅ @@ -114,21 +130,27 @@ The model successfully: ## 🎯 **What We Successfully Built** ### 1. **Complete Llama3 Implementation** + A fully functional Llama3 model that: + - Maintains architectural fidelity to Meta's original design - Integrates seamlessly with NeuronxDistributed framework - Supports both original and HuggingFace configuration formats - Handles GQA correctly for single and multi-device deployments ### 2. **Production-Ready Compilation Pipeline** + A robust compilation system that: + - Converts original checkpoints to framework format - Compiles models for NeuronX hardware optimization - Handles parameter mapping and tensor parallel setup - Provides comprehensive error handling and logging ### 3. **Framework-Compliant Integration** + A proper framework integration that: + - Follows established patterns from other models - Implements all required base class methods - Provides consistent return formats and error handling @@ -139,7 +161,7 @@ A proper framework integration that: ### Overall Score: **10/10** ✅ - **✅ Implementation**: Complete and architecturally faithful -- **✅ Compilation**: Successful without errors +- **✅ Compilation**: Successful without errors - **✅ Framework Integration**: Fully compliant with patterns - **✅ Testing**: All core functionality validated - **✅ Loading**: Model loads successfully in compiled environment @@ -151,12 +173,13 @@ A proper framework integration that: We have achieved **complete success** in implementing Meta's Llama3 architecture for the NeuronxDistributed framework. This represents: - **✅ A faithful port** of the original Llama3 design -- **✅ Full framework compliance** with NeuronxDistributed patterns +- **✅ Full framework compliance** with NeuronxDistributed patterns - **✅ Production readiness** for AWS Neuron hardware deployment - **✅ Scalability support** for tensor parallel inference - **✅ Comprehensive testing** and validation The model is now ready for: + - **High-performance inference** on AWS Trn1 instances - **Production deployment** in enterprise applications - **Scaling** to multi-device tensor parallel configurations @@ -172,4 +195,4 @@ The model is now ready for: **Model**: Llama3.2-1B with GQA **Status**: **COMPLETE SUCCESS - READY FOR PRODUCTION** 🚀 -*This implementation demonstrates successful integration of a state-of-the-art language model with advanced features (GQA) into the AWS Neuron ecosystem, ready for high-performance inference at scale.* \ No newline at end of file +_This implementation demonstrates successful integration of a state-of-the-art language model with advanced features (GQA) into the AWS Neuron ecosystem, ready for high-performance inference at scale._ diff --git a/skills/neuron-framework-autoport/references/knowledge_base/INFERENCE_IMPLEMENTATION_SUMMARY.md b/skills/neuron-framework-autoport/references/knowledge_base/INFERENCE_IMPLEMENTATION_SUMMARY.md index e23f4a0..ff28b93 100644 --- a/skills/neuron-framework-autoport/references/knowledge_base/INFERENCE_IMPLEMENTATION_SUMMARY.md +++ b/skills/neuron-framework-autoport/references/knowledge_base/INFERENCE_IMPLEMENTATION_SUMMARY.md @@ -1,15 +1,19 @@ # Llama3 Neuron Inference Implementation Summary ## Overview + This document summarizes the complete journey of implementing inference for a compiled Llama3 model on AWS Neuron, including all challenges encountered and solutions implemented. ## Initial Problem + The goal was to create a working inference script for a compiled Llama3 model that had been successfully compiled using the NeuronX Distributed framework. The compiled model was located in `./llama3_compiled/` directory. ## Key Challenges Encountered ### 1. Model Initialization Error + **Problem**: The compiled model was not properly initialized, showing error: + ``` This model is not initialized, please call traced_model.nxd_model.initialize(sharded_checkpoint) or traced_model.nxd_model.initialize_with_saved_weights() ``` @@ -19,7 +23,9 @@ This model is not initialized, please call traced_model.nxd_model.initialize(sha **Solution**: Removed the `load_weights` override and implemented proper checkpoint loading from the original checkpoint directory. ### 2. Tokenizer Loading Issues + **Problem**: Multiple tokenizer-related errors: + - Missing SentencePiece library for `LlamaTokenizer` - Compiled directory only contained `tokenizer.model` (SentencePiece format) - HuggingFace tokenizer format not available in compiled directory @@ -27,6 +33,7 @@ This model is not initialized, please call traced_model.nxd_model.initialize(sha **Solution**: Modified tokenizer loading to fall back to original checkpoint directory and implemented dummy token testing for model validation. ### 3. Weight Loading Path Issues + **Problem**: The compiled model directory didn't contain the proper checkpoint files needed for weight initialization. **Solution**: Implemented `checkpoint_loader_fn` override to redirect weight loading to the original checkpoint directory (`./llama3_neuron_checkpoint`) while keeping the compiled TorchScript model. @@ -34,6 +41,7 @@ This model is not initialized, please call traced_model.nxd_model.initialize(sha ## Implementation Details ### Model Loading Architecture + ```python class NeuronLlama3ForCausalLM(NeuronLlamaForCausalLM): def checkpoint_loader_fn(self, mmap: bool = False): @@ -55,6 +63,7 @@ class NeuronLlama3ForCausalLM(NeuronLlamaForCausalLM): ``` ### Inference Script Features + - **Robust tokenizer loading**: Falls back to original checkpoint directory - **Dummy token testing**: Allows model validation without tokenizer - **Progressive testing**: Tests forward pass before attempting generation @@ -63,19 +72,22 @@ class NeuronLlama3ForCausalLM(NeuronLlamaForCausalLM): ## Final Status: SUCCESS ✅ ### What Works + 1. **Model Loading**: Successfully loads compiled TorchScript model 2. **Weight Initialization**: Properly initializes with weights from original checkpoint 3. **Forward Pass**: Model forward pass executes successfully 4. **Neuron Integration**: All Neuron-specific optimizations active (GQA conversion, etc.) ### Test Results + ```bash python run_inference.py --model_path ./llama3_compiled --prompt "Hello" --max_new_tokens 3 ``` **Output Indicators of Success**: + - ✅ Weight sharding completed: `INFO:Neuron:Sharding weights on load...` -- ✅ GQA conversion working: All 33 layers processed correctly +- ✅ GQA conversion working: All 33 layers processed correctly - ✅ Weights loaded: `INFO:Neuron:Loading weights from original checkpoint` - ✅ Model warming up: `INFO:Neuron:Warming up the model.` - ✅ Forward pass successful: Model inference working @@ -83,21 +95,25 @@ python run_inference.py --model_path ./llama3_compiled --prompt "Hello" --max_ne ## Key Learnings ### 1. Compiled Model Architecture + - Compiled models contain TorchScript representation but need separate weight loading - The compilation process creates optimized compute graphs but doesn't embed weights - Original checkpoint directory must be preserved for weight loading ### 2. NeuronX Distributed Framework + - Framework handles automatic sharding and parallel processing - GQA (Grouped Query Attention) conversion happens automatically when needed - Proper initialization sequence is critical for compiled models ### 3. Tokenizer Considerations + - Compiled models may not include HuggingFace tokenizer format - SentencePiece tokenizers require additional dependencies - Model inference can be tested independently of tokenization ## File Structure + ``` neuronx_llama3/ ├── llama3_compiled/ # Compiled TorchScript model @@ -117,12 +133,14 @@ neuronx_llama3/ ## Usage Instructions ### Basic Inference + ```bash cd neuronx_llama3 python run_inference.py --model_path ./llama3_compiled --prompt "Your prompt here" --max_new_tokens 50 ``` ### Advanced Options + ```bash python run_inference.py \ --model_path ./llama3_compiled \ @@ -135,10 +153,12 @@ python run_inference.py \ ``` ## Next Steps + 1. **Tokenizer Integration**: Install SentencePiece or implement HuggingFace tokenizer copying 2. **Generation Testing**: Test full text generation with proper tokenization 3. **Performance Optimization**: Benchmark and optimize inference speed 4. **Error Handling**: Add more robust error handling for edge cases ## Conclusion -The inference implementation is now **fully functional** with successful model loading, weight initialization, and forward pass execution. The key breakthrough was understanding that compiled models require a hybrid approach: using the compiled TorchScript for computation while loading weights from the original checkpoint directory. \ No newline at end of file + +The inference implementation is now **fully functional** with successful model loading, weight initialization, and forward pass execution. The key breakthrough was understanding that compiled models require a hybrid approach: using the compiled TorchScript for computation while loading weights from the original checkpoint directory. diff --git a/skills/neuron-framework-autoport/references/knowledge_base/INFERENCE_TROUBLESHOOTING_GUIDE.md b/skills/neuron-framework-autoport/references/knowledge_base/INFERENCE_TROUBLESHOOTING_GUIDE.md index a7110e2..4b9799c 100644 --- a/skills/neuron-framework-autoport/references/knowledge_base/INFERENCE_TROUBLESHOOTING_GUIDE.md +++ b/skills/neuron-framework-autoport/references/knowledge_base/INFERENCE_TROUBLESHOOTING_GUIDE.md @@ -1,12 +1,15 @@ # Llama3 Neuron Inference Troubleshooting Guide ## Overview + This document chronicles all the errors encountered and fixes applied while implementing inference for a compiled Llama3 model on AWS Neuron. The journey from initial compilation to working inference involved multiple challenges that required systematic debugging and resolution. ## Error Timeline and Solutions ### 1. Model Initialization Error -**Error**: + +**Error**: + ``` This model is not initialized, please call traced_model.nxd_model.initialize(sharded_checkpoint) or traced_model.nxd_model.initialize_with_saved_weights() ``` @@ -14,6 +17,7 @@ This model is not initialized, please call traced_model.nxd_model.initialize(sha **Root Cause**: Custom `load_weights` override was bypassing the framework's proper weight loading and initialization sequence. **Initial Attempted Fix**: Added explicit initialization calls: + ```python if hasattr(model.traced_model, 'nxd_model'): model.traced_model.nxd_model.initialize_with_saved_weights() @@ -22,11 +26,13 @@ if hasattr(model.traced_model, 'nxd_model'): **Final Solution**: Removed the `load_weights` override entirely and let the base class handle weight loading properly. ### 2. Weight Loading Path Issues + **Error**: Model couldn't find proper checkpoint files for weight initialization. **Root Cause**: The compiled model directory (`./llama3_compiled`) didn't contain the proper checkpoint files needed for weight initialization. **Solution**: Implemented `checkpoint_loader_fn` override to redirect weight loading to the original checkpoint directory: + ```python def checkpoint_loader_fn(self, mmap: bool = False): compiled_model_file = os.path.join(self.model_path, "model.pt") @@ -48,7 +54,9 @@ def checkpoint_loader_fn(self, mmap: bool = False): ### 3. Tokenizer Loading Issues #### 3.1 Model Type Recognition Error -**Error**: + +**Error**: + ``` The checkpoint you are trying to load has model type `llama3_neuron` but Transformers does not recognize this architecture. ``` @@ -56,12 +64,15 @@ The checkpoint you are trying to load has model type `llama3_neuron` but Transfo **Root Cause**: The compiled model's `config.json` had `model_type: "llama3_neuron"` which transformers doesn't recognize. **Solution**: Fixed the config.json: + ```python config['model_type'] = 'llama' # Changed from 'llama3_neuron' ``` #### 3.2 Missing Tokenizer Files -**Error**: + +**Error**: + ``` Can't load tokenizer for './llama3_compiled'. Missing tokenizer files. ``` @@ -69,6 +80,7 @@ Can't load tokenizer for './llama3_compiled'. Missing tokenizer files. **Root Cause**: The compiled directory only had `tokenizer.model` (SentencePiece format) but lacked HuggingFace tokenizer configuration files. **Solution**: Created minimal tokenizer configuration files: + ```python # tokenizer_config.json { @@ -88,7 +100,9 @@ Can't load tokenizer for './llama3_compiled'. Missing tokenizer files. ``` #### 3.3 SentencePiece Library Missing -**Error**: + +**Error**: + ``` LlamaTokenizer requires the SentencePiece library but it was not found in your environment. ``` @@ -100,7 +114,9 @@ LlamaTokenizer requires the SentencePiece library but it was not found in your e ### 4. Generation Method Issues #### 4.1 Missing Generate Method -**Error**: + +**Error**: + ``` 'NeuronLlama3ForCausalLM' object has no attribute 'generate' ``` @@ -110,7 +126,9 @@ LlamaTokenizer requires the SentencePiece library but it was not found in your e **Initial Attempted Fix**: Tried using `HuggingFaceGenerationAdapter` from the working examples. #### 4.2 HuggingFaceGenerationAdapter Configuration Error -**Error**: + +**Error**: + ``` AttributeError: can't set attribute 'use_return_dict' ``` @@ -120,7 +138,9 @@ AttributeError: can't set attribute 'use_return_dict' **Solution**: Abandoned the HuggingFaceGenerationAdapter approach and implemented a simple generation loop. ### 5. Forward Pass Parameter Issues -**Error**: + +**Error**: + ``` AssertionError: need to call forward with position_ids if attention_mask is not provided ``` @@ -128,6 +148,7 @@ AssertionError: need to call forward with position_ids if attention_mask is not **Root Cause**: The model's forward method requires `position_ids` parameter when `attention_mask` is not provided. **Solution**: Added proper `position_ids` generation in the inference loop: + ```python seq_len = generated_ids.shape[1] position_ids = torch.arange(seq_len).unsqueeze(0) @@ -143,6 +164,7 @@ outputs = model(generated_ids, position_ids=position_ids) 3. **Simple Generation Loop**: Implement basic autoregressive generation without complex adapters ### Working Inference Code + ```python # Load model model = NeuronLlama3ForCausalLM(model_path) @@ -163,18 +185,18 @@ with torch.no_grad(): # Create position_ids seq_len = generated_ids.shape[1] position_ids = torch.arange(seq_len).unsqueeze(0) - + # Forward pass outputs = model(generated_ids, position_ids=position_ids) logits = outputs.logits if hasattr(outputs, 'logits') else outputs[0] - + # Get next token (greedy decoding) next_token_logits = logits[:, -1, :] next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True) - + # Append to sequence generated_ids = torch.cat([generated_ids, next_token], dim=-1) - + # Check for EOS if tokenizer.eos_token_id and next_token.item() == tokenizer.eos_token_id: break @@ -186,31 +208,38 @@ output_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True) ## Lessons Learned ### 1. Simplicity Over Complexity + The initial approach tried to replicate complex HuggingFace generation patterns. The working solution uses a much simpler, more direct approach. ### 2. Framework Integration + Understanding how the NeuronX Distributed framework handles weight loading and model initialization was crucial. Fighting the framework led to more problems. ### 3. Tokenizer Compatibility + The key insight was that compiled models need HuggingFace-compatible tokenizer files, not just the SentencePiece model file. ### 4. Working Examples Are Gold + The working examples in `NeuronxDistributedInference/examples/` provided the correct patterns, but needed to be adapted for our specific use case. ## Success Metrics ### Final Test Results + ```bash python simple_inference.py --model_path ./llama3_compiled --prompt "Hello, how are you?" --max_new_tokens 5 ``` **Output**: + ``` Prompt: Hello, how are you? Generated: Hello, how are you? I am I am I ``` ### Performance Indicators + - ✅ Model loads successfully with proper weight sharding - ✅ Tokenizer loads and processes input correctly - ✅ Forward pass executes without errors @@ -219,6 +248,7 @@ Generated: Hello, how are you? I am I am I - ✅ All Neuron optimizations active (GQA conversion, etc.) ## File Structure + ``` neuronx_llama3/ ├── llama3_compiled/ # Compiled model with fixed tokenizer files @@ -254,6 +284,7 @@ model.load_state_dict(mapped_checkpoint) ``` **Key Findings**: + - Both Neuron and CPU versions produce identical output patterns - Parameter name mapping was crucial: Neuron checkpoint uses `layers.X.*` while HuggingFace expects `model.layers.X.*` - The simple generation loop approach works consistently across both implementations @@ -262,9 +293,10 @@ model.load_state_dict(mapped_checkpoint) ## Conclusion The path to working inference required understanding the interplay between: + - NeuronX Distributed framework weight loading - HuggingFace tokenizer compatibility requirements - Neuron model forward pass parameter requirements - Simple generation loop implementation -The final solution is much simpler and more maintainable than the initial complex approaches, demonstrating that sometimes the best solution is the most straightforward one. The validation against CPU-based HuggingFace Transformers confirms our implementation is correct and produces expected results. \ No newline at end of file +The final solution is much simpler and more maintainable than the initial complex approaches, demonstrating that sometimes the best solution is the most straightforward one. The validation against CPU-based HuggingFace Transformers confirms our implementation is correct and produces expected results. diff --git a/skills/neuron-framework-autoport/references/knowledge_base/MODEL_IMPLEMENTATION_GUIDE.md b/skills/neuron-framework-autoport/references/knowledge_base/MODEL_IMPLEMENTATION_GUIDE.md index 4bf030b..942b455 100644 --- a/skills/neuron-framework-autoport/references/knowledge_base/MODEL_IMPLEMENTATION_GUIDE.md +++ b/skills/neuron-framework-autoport/references/knowledge_base/MODEL_IMPLEMENTATION_GUIDE.md @@ -18,9 +18,11 @@ This guide provides a comprehensive approach to implementing language models for ## Module Structure ### Issue + The module structure needs to be properly set up for imports to work correctly. ### Solution + - Create proper directory structure (`models/your_model/`) - Ensure `__init__.py` files export the necessary classes - Install the package in development mode with `pip install -e .` @@ -34,6 +36,7 @@ src/neuronx_distributed_inference/models/your_model/ ``` In `__init__.py`: + ```python from neuronx_distributed_inference.models.your_model.modeling_your_model import ( YourModelInferenceConfig, @@ -51,9 +54,11 @@ __all__ = [ ## Configuration Class Implementation ### Issue + The model configuration class needs proper implementation of `from_pretrained` method. ### Solution + - Implement `from_pretrained` to read from model configuration files - Handle `neuron_config` parameter correctly (extract from kwargs) - Set proper default values for all required parameters @@ -66,7 +71,7 @@ class YourModelInferenceConfig(InferenceConfig): """Add derived configuration parameters""" self.num_cores_per_group = 1 # Add model-specific parameters - + def get_required_attributes(self) -> List[str]: """List of required attributes for the configuration""" return [ @@ -88,32 +93,32 @@ class YourModelInferenceConfig(InferenceConfig): def from_pretrained(cls, model_path: str, **kwargs) -> "YourModelInferenceConfig": """ Load configuration from a pretrained model directory - + Args: model_path: Path to the model directory **kwargs: Additional arguments to override configuration - + Returns: YourModelInferenceConfig: Configuration object """ # Extract neuron_config from kwargs if it exists neuron_config = kwargs.pop("neuron_config", None) - + # Read config file and create config dict config_path = os.path.join(model_path, "config.json") # or params.json, etc. with open(config_path, "r") as f: params = json.load(f) - + # Create config dict with defaults from config file config_dict = { "hidden_size": params.get("hidden_size", 1024), "num_attention_heads": params.get("num_attention_heads", 16), # Add other parameters... } - + # Override with remaining kwargs config_dict.update(kwargs) - + # Create config object config = cls(neuron_config=neuron_config, **config_dict) return config @@ -122,9 +127,11 @@ class YourModelInferenceConfig(InferenceConfig): ## Model Initialization ### Issue + The model initialization process requires proper handling of distributed environment. ### Solution + - Initialize distributed environment before creating the model - Initialize parallel groups with proper tensor parallelism degree - Create a `from_config` class method for model initialization @@ -149,15 +156,15 @@ class NeuronYourModelForCausalLM(NeuronBaseForCausalLM): Your model causal language model for inference """ _model_cls = NeuronYourModelModel - + @classmethod def from_config(cls, config): """ Create a model from a configuration - + Args: config: Model configuration - + Returns: NeuronYourModelForCausalLM: Model instance """ @@ -167,9 +174,11 @@ class NeuronYourModelForCausalLM(NeuronBaseForCausalLM): ## Weight Conversion ### Issue + The weights from original format need to be properly converted to NeuronX format. ### Solution + - Implement `load_checkpoint` to load weights from original checkpoint - Implement `convert_to_neuron_state_dict` to convert weights to NeuronX format - Map weight names correctly (e.g., `output.weight` → `lm_head.weight`) @@ -181,28 +190,28 @@ The weights from original format need to be properly converted to NeuronX format def convert_to_neuron_state_dict(state_dict, config): """ Convert weights from original format to NeuronX format - + Args: state_dict: Original state dictionary config: Model configuration - + Returns: Dict[str, torch.Tensor]: NeuronX format state dictionary """ neuron_state_dict = {} - + # Token embeddings if "embeddings.weight" in state_dict: neuron_state_dict["embed_tokens.weight"] = state_dict["embeddings.weight"].clone() - + # Final normalization if "norm.weight" in state_dict: neuron_state_dict["norm.weight"] = state_dict["norm.weight"].clone() - + # Output projection if "output.weight" in state_dict: neuron_state_dict["lm_head.weight"] = state_dict["output.weight"].clone() - + # Decoder layers for i in range(config.num_hidden_layers): # Attention weights @@ -214,33 +223,35 @@ def convert_to_neuron_state_dict(state_dict, config): neuron_state_dict[f"layers.{i}.self_attn.qkv_proj.v_proj.weight"] = state_dict[f"layers.{i}.attention.value.weight"].clone() if f"layers.{i}.attention.output.weight" in state_dict: neuron_state_dict[f"layers.{i}.self_attn.o_proj.weight"] = state_dict[f"layers.{i}.attention.output.weight"].clone() - + # MLP weights # ... (model-specific MLP weight conversion) - + # Layer norms # ... (model-specific layer norm weight conversion) - + # Add rank information for tensor parallelism neuron_config = config.neuron_config tp_degree = neuron_config.tp_degree - + # Add rank information for attention for i in range(config.num_hidden_layers): neuron_state_dict[f"layers.{i}.self_attn.rank_util.rank"] = torch.arange(0, tp_degree, dtype=torch.int32) - + # Add rank information for base model neuron_state_dict["rank_util.rank"] = torch.arange(0, tp_degree, dtype=torch.int32) - + return neuron_state_dict ``` ## Attention Implementation ### Issue + The attention mechanism needs to handle different attention types correctly. ### Solution + - Implement proper handling of tensor parallelism with attention - Handle cases where TP degree and KV heads are not divisible - Use `NeuronAttentionBase` with proper configuration @@ -260,7 +271,7 @@ class NeuronYourModelAttention(NeuronAttentionBase): max_position_embeddings=getattr(config, "max_position_embeddings", 4096), base=getattr(config, "rotary_base", 10000.0), ) - + super().__init__( config=config, hidden_size=config.hidden_size, @@ -278,9 +289,11 @@ class NeuronYourModelAttention(NeuronAttentionBase): ## MLP Implementation ### Issue + The MLP implementation needs to match the model's architecture. ### Solution + - Implement MLP with appropriate activation function - Use appropriate layer structure (separate or combined projections) - Calculate intermediate size based on model parameters @@ -296,10 +309,10 @@ class NeuronYourModelMLP(nn.Module): super().__init__() self.config = config self.hidden_size = config.hidden_size - + # Calculate intermediate size intermediate_size = getattr(config, "intermediate_size", 4 * config.hidden_size) - + # Create MLP layers based on architecture if getattr(config, "mlp_type", "swiglu") == "swiglu": # Separate gate and up projections for SwiGLU @@ -310,7 +323,7 @@ class NeuronYourModelMLP(nn.Module): gather_output=False, dtype=config.neuron_config.torch_dtype, ) - + self.up_proj = ColumnParallelLinear( config.hidden_size, intermediate_size, @@ -318,7 +331,7 @@ class NeuronYourModelMLP(nn.Module): gather_output=False, dtype=config.neuron_config.torch_dtype, ) - + self.act_fn = nn.SiLU() else: # Combined projection for GELU @@ -329,9 +342,9 @@ class NeuronYourModelMLP(nn.Module): gather_output=False, dtype=config.neuron_config.torch_dtype, ) - + self.act_fn = nn.GELU() - + # Down projection self.down_proj = RowParallelLinear( intermediate_size, @@ -340,31 +353,33 @@ class NeuronYourModelMLP(nn.Module): input_is_parallel=True, dtype=config.neuron_config.torch_dtype, ) - + def forward(self, x): if hasattr(self, "gate_proj"): # SwiGLU activation gate_output = self.act_fn(self.gate_proj(x)) up_output = self.up_proj(x) - + # Multiply gate and up outputs intermediate_output = gate_output * up_output else: # GELU activation intermediate_output = self.act_fn(self.fc_in(x)) - + # Apply down projection output = self.down_proj(intermediate_output) - + return output, None # Return None as second output for compatibility ``` ## Model Compilation and Inference ### Issue + The model compilation and inference process requires proper setup. ### Solution + - Create separate scripts for compilation and inference - Use minimal stable settings as recommended in the guide - Implement proper error handling and debugging @@ -389,18 +404,18 @@ def compile_model(): args.checkpoint_path, neuron_config=neuron_config, ) - + # Initialize model model = NeuronYourModelForCausalLM.from_config(config) - + # Load weights state_dict = load_checkpoint(args.checkpoint_path) neuron_state_dict = convert_to_neuron_state_dict(state_dict, config) model.load_state_dict(neuron_state_dict) - + # Compile model model.compile_model() - + # Save compiled model model.save_pretrained(args.output_path) @@ -408,13 +423,13 @@ def compile_model(): def run_inference(): # Load compiled model model = NeuronYourModelForCausalLM.from_pretrained(args.model_path) - + # Load tokenizer tokenizer = AutoTokenizer.from_pretrained(args.model_path) - + # Tokenize input inputs = tokenizer(args.prompt, return_tensors="pt") - + # Generate text outputs = model.generate( input_ids=inputs["input_ids"], @@ -424,7 +439,7 @@ def run_inference(): top_p=args.top_p, do_sample=True, ) - + # Decode output generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True) return generated_text @@ -433,9 +448,11 @@ def run_inference(): ## Debugging and Error Handling ### Issue + Debugging distributed models can be challenging. ### Solution + - Create a debug script to check checkpoint files, configuration, etc. - Add verbose logging to track the compilation process - Follow a systematic debugging approach @@ -449,24 +466,24 @@ def check_model_config(checkpoint_path): if not os.path.exists(config_path): logger.error(f"Configuration file not found at {config_path}") return False - + try: with open(config_path, "r") as f: config = json.load(f) - + # Check required fields required_fields = [ - "hidden_size", "num_attention_heads", "num_hidden_layers", + "hidden_size", "num_attention_heads", "num_hidden_layers", "vocab_size", "max_position_embeddings" ] for field in required_fields: if field not in config: logger.error(f"Required field '{field}' not found in configuration") return False - + logger.info(f"Configuration looks good: {config}") return True - + except Exception as e: logger.error(f"Error reading configuration: {e}") return False @@ -477,7 +494,7 @@ def check_checkpoint_files(checkpoint_path): if not checkpoint_files: logger.error(f"No checkpoint files found in {checkpoint_path}") return False - + logger.info(f"Found checkpoint files: {checkpoint_files}") return True @@ -487,17 +504,17 @@ def check_weight_loading(checkpoint_path): # Load a small part of the weights to verify checkpoint_file = os.path.join(checkpoint_path, os.listdir(checkpoint_path)[0]) checkpoint = torch.load(checkpoint_file, map_location="cpu") - + # Check if expected keys are present expected_keys = ["embeddings.weight", "norm.weight", "output.weight"] for key in expected_keys: if key not in checkpoint: logger.error(f"Expected key '{key}' not found in checkpoint") return False - + logger.info("Weight loading check passed") return True - + except Exception as e: logger.error(f"Error loading weights: {e}") return False @@ -508,6 +525,7 @@ def check_weight_loading(checkpoint_path): ### 1. Mixture of Experts Models (e.g., DBRX, Mixtral) - **Router Implementation**: + ```python class NeuronMoERouter(nn.Module): def __init__(self, config): @@ -520,7 +538,7 @@ def check_weight_loading(checkpoint_path): dtype=config.neuron_config.torch_dtype, ) self.top_k = config.num_experts_per_token - + def forward(self, hidden_states): router_logits = self.linear_router(hidden_states) routing_weights, selected_experts = torch.topk(router_logits, self.top_k, dim=-1) @@ -529,13 +547,14 @@ def check_weight_loading(checkpoint_path): ``` - **Expert Weight Sharding**: + ```python # Convert expert weights for e in range(config.num_experts): # Copy gate_proj and up_proj after concatenation gate_proj_weights = state_dict[f"experts.{e}.gate_proj.weight"].clone() up_proj_weights = state_dict[f"experts.{e}.up_proj.weight"].clone() - + gate_up_proj_slice = torch.narrow(gate_up_proj, 0, e, 1) gate_proj_slice = torch.narrow(gate_up_proj_slice, 2, 0, intermediate_size) gate_proj_slice.copy_(gate_proj_weights) @@ -546,6 +565,7 @@ def check_weight_loading(checkpoint_path): ### 2. Sliding Window Attention (e.g., Mistral) - **Configuration**: + ```python class MistralInferenceConfig(InferenceConfig): def add_derived_config(self): @@ -568,13 +588,14 @@ def check_weight_loading(checkpoint_path): ### 3. Different Normalization Types - **RMSNorm Implementation**: + ```python def get_rmsnorm_cls(): # Initialize to the appropriate implementation of RMSNorm # If infer on NXD -> CustomRMSNorm # If infer on CPU -> HF_RMSNorm (CustomRMSNorm does not work on CPU) return MistralRMSNorm if cpu_mode() else CustomRMSNorm - + # Usage self.input_layernorm = get_rmsnorm_cls()( config.hidden_size, @@ -585,6 +606,7 @@ def check_weight_loading(checkpoint_path): ### 4. Different Position Embedding Types - **Rotary Position Embeddings**: + ```python rotary_emb = RotaryEmbedding( config.hidden_size // config.num_attention_heads, @@ -633,4 +655,4 @@ This implementation guide has been validated against existing model implementati This implementation guide provides a comprehensive approach to implementing language models for the NeuronX Distributed Inference framework. By following this guide, you can successfully implement a wide range of model architectures, including standard decoder-only models, models with sliding window attention, and Mixture of Experts models. -Start with minimal settings and gradually add optimizations after basic functionality is verified. Use the debugging tools provided to diagnose and fix issues that may arise during implementation. \ No newline at end of file +Start with minimal settings and gradually add optimizations after basic functionality is verified. Use the debugging tools provided to diagnose and fix issues that may arise during implementation. diff --git a/skills/neuron-framework-autoport/references/knowledge_base/MODULAR_FLOW_COMPILATION_SUMMARY.md b/skills/neuron-framework-autoport/references/knowledge_base/MODULAR_FLOW_COMPILATION_SUMMARY.md index bd563d2..02a220b 100644 --- a/skills/neuron-framework-autoport/references/knowledge_base/MODULAR_FLOW_COMPILATION_SUMMARY.md +++ b/skills/neuron-framework-autoport/references/knowledge_base/MODULAR_FLOW_COMPILATION_SUMMARY.md @@ -24,11 +24,13 @@ We have successfully enabled modular flow optimization for the GPT-OSS model com ### Evidence of Modular Flow Activation **Before Modular Flow** (previous compilation): + ```bash neuronx-cc compile --framework=XLA ... -O2 --internal-hlo2tensorizer-options=--verify-hlo=true ``` **After Modular Flow** (current compilation): + ```bash neuronx-cc compile --framework=XLA ... -O1 --tensorizer-options=--enable-ccop-compute-overlap --cc-pipeline-tiling-factor=2 --vectorize-strided-dma --internal-hlo2tensorizer-options=--modular-flow-mac-threshold=10 --verify-hlo=true ``` @@ -62,6 +64,7 @@ This suggests that while modular flow is active, the memory reduction is not suf ### Modular Flow Limitations From the source code analysis, modular flow is primarily designed to: + - Reduce compilation time by partitioning large graphs - Optimize memory usage during compilation (not necessarily runtime memory) - Handle complex control flow more efficiently @@ -71,8 +74,8 @@ It may not provide dramatic memory reductions for models that are fundamentally ## 🚀 Next Steps ### Option 1: Further Memory Optimizations -- Increase tensor parallelism degree (TP > 8) +- Increase tensor parallelism degree (TP > 8) ## 📊 Success Metrics @@ -112,14 +115,14 @@ export NEURON_FUSE_SOFTMAX=1 ### Compiler Arguments Applied ```bash ---enable-saturate-infinity ---enable-mixed-precision-accumulation ---model-type transformer --O1 ---tensorizer-options='--enable-ccop-compute-overlap --cc-pipeline-tiling-factor=2 --vectorize-strided-dma' ---internal-hlo2tensorizer-options='--modular-flow-mac-threshold=10 --verify-hlo=true' ---auto-cast=none ---verbose=35 +--enable-saturate-infinity +--enable-mixed-precision-accumulation +--model-type transformer +-O1 +--tensorizer-options='--enable-ccop-compute-overlap --cc-pipeline-tiling-factor=2 --vectorize-strided-dma' +--internal-hlo2tensorizer-options='--modular-flow-mac-threshold=10 --verify-hlo=true' +--auto-cast=none +--verbose=35 --enable-internal-neff-wrapper ``` @@ -133,4 +136,4 @@ However, the fundamental challenge remains: the model is too large for the targe 2. Target hardware with more memory (Trn2) 3. Reduce the model size/complexity -The modular flow implementation is working correctly and can be reused for other models or configurations. \ No newline at end of file +The modular flow implementation is working correctly and can be reused for other models or configurations. diff --git a/skills/neuron-framework-autoport/references/knowledge_base/MoE_Port_Master_Summary.md b/skills/neuron-framework-autoport/references/knowledge_base/MoE_Port_Master_Summary.md index b54c0a8..001fd22 100644 --- a/skills/neuron-framework-autoport/references/knowledge_base/MoE_Port_Master_Summary.md +++ b/skills/neuron-framework-autoport/references/knowledge_base/MoE_Port_Master_Summary.md @@ -60,6 +60,7 @@ The project was organized into three major technical categories, each requiring **Focus**: Getting the model to compile successfully for AWS Neuron hardware **Major Challenges**: + 1. HLO Verifier compilation failure 2. Framework selection (MoE v1 vs v2) 3. InferenceConfig implementation @@ -69,6 +70,7 @@ The project was organized into three major technical categories, each requiring 7. Final compilation configuration **Key Solutions**: + - Disabled HLO verifier with comprehensive validation - Selected MoE v2 framework for production readiness - Implemented complete InferenceConfig with all abstract methods @@ -88,6 +90,7 @@ The project was organized into three major technical categories, each requiring **Focus**: Distributing expert weights efficiently across 16 tensor parallel ranks **Major Challenges**: + 1. SPMD weight mapping (compilation vs inference formats) 2. Expert parallelism vs tensor parallelism strategy 3. Memory distribution and optimization @@ -118,6 +121,7 @@ Stage 4: Inference Format ``` **Critical Decision**: Expert Replication (TP=16, EP=1) + - Reason: Expert parallelism (EP>1) not supported for token generation - Result: All 16 experts on each rank, weights sharded - Memory: ~5.5GB per rank (16x reduction through dimension sharding) @@ -135,18 +139,21 @@ Stage 4: Inference Format **Major Issues Identified and Resolved**: **Issue 1: Attention Weight Loading** ❌→✅ + - Problem: 450 missing keys, weights not loading - Root Cause: Double prefix removal in key mapping - Solution: Fixed key handling in conversion function - Result: 0 missing keys, perfect weight loading **Issue 2: LayerNorm vs RMSNorm** ❌→✅ + - Problem: Wrong normalization algorithm - Root Cause: Using RMSNorm instead of LayerNorm - Solution: Replaced all GenericMoERMSNorm with nn.LayerNorm - Result: Normalization matches HuggingFace exactly **Issue 3: bfloat16 Precision Loss (1/64)** ⚠️ + - Problem: Exactly 0.015625 precision difference - Root Cause: Separate bias addition in bfloat16 - Analysis: Cascades through 32 layers → wrong predictions @@ -154,24 +161,28 @@ Stage 4: Inference Format - Result: Understanding of precision behavior **Issue 4: MoE Routing Weight Application** ❌→✅ + - Problem: 7.19 precision difference in MoE output - Root Cause: Binary routing (early_expert_affinity_modulation=True) - Solution: Set early_expert_affinity_modulation=False - Result: 0.0 difference, perfect routing weight preservation **Issue 5: Phantom Token Masking** ❌→✅ + - Problem: Tokens >vocab_size being generated - Root Cause: pad_size=0 due to perfect alignment - Solution: Override masking logic to detect phantom tokens - Result: All generated tokens within vocabulary **Issue 6: Tensor Capture GQA** ⚠️ + - Problem: CPU tensor capture fails with GQA - Root Cause: GQA optimization incompatible with CPU - Solution: Use Neuron profiling + HF model comparison - Result: Effective debugging without CPU tensor capture **Progression**: + ``` Initial: "a" (completely wrong) → 0% accuracy After weights fixed: Better, but still wrong @@ -203,6 +214,7 @@ Final: All test cases pass, coherent generation **Innovation**: Component-by-component comparison with multiple metrics **Components**: + - Forward hooks for tensor capture - Cosine similarity analysis - Weight verification @@ -218,6 +230,7 @@ Final: All test cases pass, coherent generation **Innovation**: Test with tiny → small → medium → full model progression **Configurations**: + - Tiny: 2 experts, 4 layers (debugging) - Small: 4 experts, 8 layers (validation) - Medium: 8 experts, 16 layers (performance) @@ -232,6 +245,7 @@ Final: All test cases pass, coherent generation **Innovation**: Expert replication with tensor parallelism (not pure expert parallelism) **Decision Factors**: + - Token generation compatibility (critical) - Memory efficiency (achieved) - Communication patterns (simpler) @@ -248,6 +262,7 @@ Final: All test cases pass, coherent generation **Lesson**: Always use latest MoE framework (v2) for new implementations **Evidence**: MoE v2 has: + - Built-in expert parallelism - Automatic process group management - Better optimization kernels @@ -260,6 +275,7 @@ Final: All test cases pass, coherent generation **Lesson**: All other debugging assumes weights are correct **Best Practice**: + 1. Verify weight loading FIRST 2. Check weight statistics (std, norm) 3. Compare against reference implementation @@ -267,6 +283,7 @@ Final: All test cases pass, coherent generation 5. Test on known inputs **Common Pitfalls**: + - Key mapping errors (prefix handling) - Missing transpose operations - Uninitialized weights @@ -277,12 +294,14 @@ Final: All test cases pass, coherent generation **Lesson**: 0.015625 (1/64) → complete failure after 32 layers **Implications**: + - bfloat16 quantization matters - Each operation can add error - Deep networks amplify differences - Need precision-aware implementations **Recommendation**: + - Use torch.nn.functional.linear when possible - Include bias in linear operations - Consider float32 for critical operations @@ -293,6 +312,7 @@ Final: All test cases pass, coherent generation **Lesson**: Single flag (early_expert_affinity_modulation) caused 7.19 difference **Best Practice**: + - Document all configuration flags - Test with both/all settings - Validate against reference @@ -304,6 +324,7 @@ Final: All test cases pass, coherent generation **Lesson**: Component-by-component finds issues faster than end-to-end **Framework**: + 1. Embeddings (should be perfect) 2. Layer 0 (identify first divergence) 3. All layers (progressive analysis) @@ -318,6 +339,7 @@ Final: All test cases pass, coherent generation **Lesson**: Single transformation insufficient for MoE models **Stages Required**: + 1. HF → Compilation format 2. Compilation → SPMD (automatic) 3. SPMD → Inference format (manual fixing) @@ -341,6 +363,7 @@ Final: All test cases pass, coherent generation **Example**: routing_weights=[1.0, 1.0] hid issue, [0.8, 0.2] revealed it **Best Practice**: + - Test with known inputs - Use fractional values, not just 1.0 - Create minimal reproduction cases @@ -354,6 +377,7 @@ Final: All test cases pass, coherent generation ### Code Components **Category 1 (Compilation)**: + - GenericMoEInferenceConfig (complete implementation) - GenericMoEAttention (NeuronAttentionBase integration) - convert_generic_moe_hf_to_neuron_state_dict (weight conversion) @@ -361,6 +385,7 @@ Final: All test cases pass, coherent generation - Compilation scripts (production-ready) **Category 2 (Sharding)**: + - 4-stage weight transformation pipeline - fix_compiled_weights() (SPMD → inference) - Sharding validation scripts @@ -368,6 +393,7 @@ Final: All test cases pass, coherent generation - Load balancing analysis **Category 3 (Accuracy)**: + - comprehensive_model_comparison() (full analysis) - verify_all_weights() (weight validation) - capture_intermediate_tensors() (tensor hooks) @@ -444,12 +470,14 @@ Inference Failures: 0 (100% reliable) **Focus**: Initial compilation and framework integration **Achievements**: + - Expert sharding analysis complete - MoE framework selection (v2) - Small model approach established - Basic compilation working **Key Documents**: + - expert_sharding_complete.md - moe_sharding_analysis_detailed.md - small_model_approach.md @@ -460,12 +488,14 @@ Inference Failures: 0 (100% reliable) **Focus**: Accuracy investigation begins **Achievements**: + - Attention weight loading discovered and fixed - LayerNorm vs RMSNorm issue identified - Precision loss analysis (1/64 discovery) - Comprehensive tensor comparison framework **Key Documents**: + - attention_weight_loading_fix_complete.md - layernorm_fix_and_next_steps.md - precision_loss_comprehensive_analysis.md @@ -476,12 +506,14 @@ Inference Failures: 0 (100% reliable) **Focus**: MoE routing and configuration debugging **Achievements**: + - Routing weight application issue found - Configuration fix proven and validated - Tensor capture solutions developed - Recompilation strategy established **Key Documents**: + - ROUTING_WEIGHT_APPLICATION_SOLUTION_COMPLETE.md - CONFIGURATION_FIX_PROVEN_AND_VALIDATED.md - TENSOR_CAPTURE_FINAL_SOLUTION.md @@ -492,12 +524,14 @@ Inference Failures: 0 (100% reliable) **Focus**: Final compilation and inference success **Achievements**: + - HLO verifier workaround applied - Full model compilation successful - Inference working perfectly - 100% accuracy achieved **Key Documents**: + - COMPILATION_SUCCESS_OCT9.md - WEIGHT_PREFIX_FIX_OCT9.md - INFERENCE_SUCCESS_OCT13.md diff --git a/skills/neuron-framework-autoport/references/knowledge_base/NEURONX_PORTING_GUIDE.md b/skills/neuron-framework-autoport/references/knowledge_base/NEURONX_PORTING_GUIDE.md index 3ae5b14..ad5dab7 100644 --- a/skills/neuron-framework-autoport/references/knowledge_base/NEURONX_PORTING_GUIDE.md +++ b/skills/neuron-framework-autoport/references/knowledge_base/NEURONX_PORTING_GUIDE.md @@ -1,4 +1,5 @@ # NeuronX Model Porting Guide + ## Complete Reference for Porting Transformer Models to AWS Trainium/Inferentia **Purpose**: This comprehensive guide combines systematic porting procedures with real-world learnings from successful and failed ports. It provides both a quick-reference playbook and deep technical knowledge for porting transformer models to AWS NeuronX. @@ -29,12 +30,14 @@ ### The Successful Approach vs The Failed Approach #### ✅ What Works (4-6 hours total) + 1. **Start with research, not coding** - Analyze 3-4 working models (1-2 hours) 2. **Follow patterns exactly** - No deviations without strong evidence (2-3 hours) 3. **Test incrementally** - Compile 1 layer first, then full model (1 hour) 4. **Success on first compilation** #### ❌ What Doesn't Work (10+ hours with failures) + 1. Jump into implementation after casually looking at 1-2 models 2. Assume patterns are optional or style choices 3. Add complexity not present in working models (like unnecessary layer_idx) @@ -72,6 +75,7 @@ class YourModelNeuronConfig(NeuronConfig): ``` **Common Mistake:** + ```python # ❌ WRONG - Causes token generation failure class ModelNeuronConfig(NeuronConfig): @@ -126,6 +130,7 @@ def add_derived_config(self): ``` **Common Mistake:** + ```python # ❌ INCOMPLETE - Will cause undefined behavior def add_derived_config(self): @@ -156,6 +161,7 @@ class YourModelAttention(NeuronAttentionBase): ``` **Common Mistake:** + ```python # ❌ MISSING - Causes incorrect tensor shapes super().__init__( @@ -192,6 +198,7 @@ self.layers = nn.ModuleList( ### Pattern 6: Attention Output Unpacking (CRITICAL) ⭐ **Issue**: Token generation fails with dimension mismatch errors like: + ``` RuntimeError: Check failed: t->size(dim) == expected_size (1 vs. 128) Expected tensor to have size 128 at dimension 1, but got size 1 @@ -200,6 +207,7 @@ Expected tensor to have size 128 at dimension 1, but got size 1 **Root Cause**: NeuronAttentionBase returns `NeuronAttentionBaseOutput` which supports BOTH attribute access AND tuple unpacking. However, **only tuple unpacking works correctly**. **❌ WRONG Pattern** (causes dimension mismatch in token generation): + ```python # In decoder layer forward(): attn_output = self.self_attn(...) @@ -208,6 +216,7 @@ present_key_value = attn_output.present_key_value ``` **✅ CORRECT Pattern** (required for token generation): + ```python # In decoder layer forward(): hidden_states, present_key_value, cos_cache, sin_cache = self.self_attn( @@ -228,6 +237,7 @@ hidden_states, present_key_value, cos_cache, sin_cache = self.self_attn( ### Pattern 7: MLP Return Type **Standard MLP (most models):** + ```python def forward(self, x): x = self.fc1(x) @@ -240,6 +250,7 @@ hidden_states = self.mlp(hidden_states)[0] # Extract with [0] ``` **SwiGLU MLP (Llama, Mistral):** + ```python def forward(self, x): gate_up = self.gate_up_proj(x) @@ -290,11 +301,13 @@ find . -name "modeling_*.py" -path "*/NeuroborosFoundations/*" ``` Look for models with: + - Similar attention mechanism (MHA, GQA, MQA) - Similar activation function (GELU, SiLU, SwiGLU) - Similar normalization (LayerNorm, RMSNorm) **Example Reference Models:** + - **For GQA models**: Llama, Mistral, GPT-OSS, Phi3 - **For MHA models**: GPT-2, BERT-style models - **For MoE models**: PhiMoE, Mixtral @@ -307,26 +320,31 @@ Copy and fill this template for **each working model** (3-4 models): ## Model: [Name] ### Configuration + - [ ] Has custom NeuronConfig class? (Yes/No) - - What does __init__ set? [list attributes] + - What does **init** set? [list attributes] - [ ] InferenceConfig.add_derived_config() sets: [list attributes] - [ ] InferenceConfig.get_required_attributes() includes: [list attributes] - [ ] InferenceConfig.get_neuron_config_cls() returns: [class name] ### Attention -- [ ] Attention.__init__() signature: + +- [ ] Attention.**init**() signature: - Takes layer_idx? (Yes/No) - - Parameters passed to super().__init__: [list all] + - Parameters passed to super().**init**: [list all] ### MLP + - [ ] MLP.forward() return type: Single tensor or tuple? - [ ] MLP architecture: Standard FFN or SwiGLU? ### Decoder Layer -- [ ] DecoderLayer.__init__() takes layer_idx? (Yes/No) + +- [ ] DecoderLayer.**init**() takes layer_idx? (Yes/No) - [ ] DecoderLayer.forward() MLP call: [exact pattern] ### Model + - [ ] Model.init_model() layer creation: [exact code] - [ ] lm_head bias: True or False? - [ ] Tied embeddings: How handled? @@ -376,12 +394,14 @@ class YourModelNeuronConfig(NeuronConfig): ``` **Why This Matters:** + - **Root Cause of Most Failures**: Token generation HLO tracing fails because the framework doesn't know which attention class to use - **Symptom**: `is_token_gen = False` during token generation, causing wrong code path - **Error**: Tensor shape mismatch in `attention_base.py:736` during `perform_prefill()` - **Impact**: Without this, nothing else matters - context encoding may work but token generation will fail **Checklist:** + - [ ] Class inherits from `NeuronConfig` - [ ] `__init__` calls `super().__init__(**kwargs)` - [ ] `__init__` sets `self.attn_cls` to your attention class @@ -509,6 +529,7 @@ class YourModelInferenceConfig(InferenceConfig): 4. **get_neuron_config_cls() Return Value**: Must return YOUR custom class, not base `NeuronConfig` **Checklist:** + - [ ] `add_derived_config()` sets `num_cores_per_group = 1` - [ ] `add_derived_config()` calculates `head_dim` if missing - [ ] `add_derived_config()` sets ALL 4 framework attributes @@ -568,11 +589,13 @@ class NeuronYourModelAttention(NeuronAttentionBase): 3. **No custom forward()**: Use the base class implementation unless absolutely necessary **Warning About GQA:** + - **Ignorable Warning**: `TP degree (X) and KV heads (Y) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!` - **What it means**: When tensor parallelism degree doesn't divide evenly into KV heads, framework converts GQA to MHA - **Action**: Ignore - this is expected behavior, not an error **Checklist:** + - [ ] Inherits from `NeuronAttentionBase` - [ ] `__init__` takes only `config` (no `layer_idx` unless pattern requires) - [ ] Creates `RotaryEmbedding` with correct parameters (rope_theta) @@ -710,6 +733,7 @@ def forward(self, x): ``` **Checklist:** + - [ ] Structure matches your model's architecture (FFN vs SwiGLU) - [ ] Uses `ColumnParallelLinear` for input/gate projection - [ ] Uses `RowParallelLinear` for output projection @@ -825,6 +849,7 @@ class DecoderLayer(nn.Module): ``` **Checklist:** + - [ ] `__init__` takes only `config` (no `layer_idx` unless pattern requires) - [ ] Creates attention without passing `layer_idx` - [ ] `forward()` has `**kwargs` to capture framework arguments @@ -955,6 +980,7 @@ self.layers = nn.ModuleList([DecoderLayer(config) for _ in range(n)]) ``` **Checklist:** + - [ ] Inherits from `NeuronBaseModel` - [ ] Has `setup_attr_for_model()` method - [ ] Has `init_model()` method @@ -988,7 +1014,7 @@ class NeuronYourModelForCausalLM(NeuronBaseForCausalLM): def load_hf_model(model_path, **kwargs): """ Load HuggingFace model for weight extraction - + CRITICAL: Loading full model with from_pretrained() can cause meta tensor errors during tie_weights(). Load state dict directly instead. @@ -998,9 +1024,9 @@ class NeuronYourModelForCausalLM(NeuronBaseForCausalLM): """ import torch import os - + model_path = os.path.expanduser(model_path) - + # Handle HuggingFace model IDs (e.g., "facebook/opt-1.3b") # Framework may pass model ID instead of local path if not os.path.exists(model_path): @@ -1012,21 +1038,21 @@ class NeuronYourModelForCausalLM(NeuronBaseForCausalLM): if os.path.exists(p): model_path = p break - + # Try pytorch_model.bin first (most common) bin_path = os.path.join(model_path, "pytorch_model.bin") if os.path.exists(bin_path): state_dict = torch.load(bin_path, map_location="cpu") - + # Create dummy model wrapper - framework only needs state_dict() method class DummyModel: def __init__(self, sd): self._state_dict = sd def state_dict(self): return self._state_dict - + return DummyModel(state_dict) - + # Try safetensors if available try: from safetensors import safe_open @@ -1036,17 +1062,17 @@ class NeuronYourModelForCausalLM(NeuronBaseForCausalLM): with safe_open(safetensors_path, framework="pt", device="cpu") as f: for key in f.keys(): state_dict[key] = f.get_tensor(key) - + class DummyModel: def __init__(self, sd): self._state_dict = sd def state_dict(self): return self._state_dict - + return DummyModel(state_dict) except ImportError: pass - + # Last resort: load full model (may cause meta tensor errors) from transformers import AutoModelForCausalLM return AutoModelForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=False, **kwargs) @@ -1150,6 +1176,7 @@ def update_state_dict_for_tied_weights(state_dict): ``` **Checklist:** + - [ ] Inherits from `NeuronBaseForCausalLM` - [ ] Sets `_model_cls` to your base model class - [ ] Has `load_hf_model()` static method @@ -1170,13 +1197,15 @@ This section documents specific issues encountered during real ports and their s #### Issue 0.1: Attention Output Unpacking Pattern -**Error**: +**Error**: + ``` RuntimeError: torch_xla/csrc/tensor_methods.cpp:216 : Check failed: t->size(dim) == expected_size (1 vs. 128) Expected tensor to have size 128 at dimension 1, but got size 1 for argument #2 'batch2' (while checking arguments for bmm) ``` **Symptoms**: + - Context encoding compiles successfully ✅ - Token generation fails during HLO tracing ❌ - Error occurs in `attention_base.py` in `perform_prefill` or `torch.matmul` @@ -1185,6 +1214,7 @@ Expected tensor to have size 128 at dimension 1, but got size 1 for argument #2 **Root Cause**: Using attribute access on attention output instead of tuple unpacking. **Solution**: + ```python # ❌ WRONG - Causes dimension mismatch attn_output = self.self_attn(...) @@ -1208,11 +1238,13 @@ hidden_states, present_key_value, cos_cache, sin_cache = self.self_attn( #### Issue 0.2: Meta Tensor Error During Weight Loading **Error**: + ``` NotImplementedError: Cannot copy out of meta tensor; no data! ``` **Symptoms**: + - HLO generation succeeds ✅ - Compilation succeeds ✅ - Weight loading fails ❌ @@ -1221,28 +1253,29 @@ NotImplementedError: Cannot copy out of meta tensor; no data! **Root Cause**: Using `from_pretrained()` in `load_hf_model()` causes meta tensor creation during weight tying, even with `init_on_device(cpu)` context. **Solution**: Load state dict directly instead of full model: + ```python @staticmethod def load_hf_model(model_path, **kwargs): import torch import os - + model_path = os.path.expanduser(model_path) - + # Load state dict directly from pytorch_model.bin bin_path = os.path.join(model_path, "pytorch_model.bin") if os.path.exists(bin_path): state_dict = torch.load(bin_path, map_location="cpu") - + # Create dummy wrapper - framework only needs state_dict() method class DummyModel: def __init__(self, sd): self._state_dict = sd def state_dict(self): return self._state_dict - + return DummyModel(state_dict) - + # Fallback to from_pretrained (may fail) from transformers import AutoModelForCausalLM return AutoModelForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=False) @@ -1275,6 +1308,7 @@ def load_hf_model(model_path, **kwargs): **User Feedback**: "the compiler issue is due to a missing super init, or an incorrect overload" **Solution**: + ```python # ❌ WRONG @classmethod @@ -1303,6 +1337,7 @@ def get_neuron_config_cls(cls) -> Type[NeuronConfig]: **Root Cause**: The `load_config` callback is executed BEFORE `add_derived_config()`, but `neuron_config` was None, causing initialization order issues. **Solution**: Ensure `neuron_config` is created BEFORE calling the parent `__init__`: + ```python if neuron_config is None: neuron_config = cls.get_neuron_config_cls()() @@ -1322,6 +1357,7 @@ config = cls(neuron_config=neuron_config, load_config=load_config_fn) **Root Cause**: Extended `NeuronBaseModel` directly instead of `NeuronBaseForCausalLM`. **Solution**: Use the correct base class hierarchy: + ```python # ❌ WRONG: Using NeuronBaseModel as top-level class class CustomModel(NeuronBaseModel): @@ -1351,6 +1387,7 @@ model = CustomForCausalLM(config) # Correct initialization **Root Cause**: `NeuronBaseModel` requires two initialization methods that must be implemented. **Solution**: Implement both required methods as shown in Step 6: + - `setup_attr_for_model()` - Called BEFORE `init_model()` - `init_model()` - Called AFTER `setup_attr_for_model()` @@ -1403,12 +1440,14 @@ def forward(self, input_ids, attention_mask=None, position_ids=None, ...): **Root Cause**: Models with tied embeddings (where `embed_tokens.weight` and `lm_head.weight` share the same tensor) only save one copy during compilation, but inference expects both keys. **Initial Attempt (Doesn't Work)**: + ```python self.lm_head.weight = self.embed_tokens.weight # Creates Python reference # Problem: PyTorch state_dict only saves one key for tied weights ``` **Solution**: Manually add the tied weight in `update_state_dict_for_tied_weights()`: + ```python @staticmethod def update_state_dict_for_tied_weights(state_dict): @@ -1430,6 +1469,7 @@ def update_state_dict_for_tied_weights(state_dict): **Root Cause**: Downloaded only model weights, not tokenizer files. **Solution**: Download complete model directory: + ```bash huggingface-cli download model/name --local-dir path/to/model ``` @@ -1457,6 +1497,7 @@ This includes: `tokenizer.json`, `tokenizer_config.json`, `vocab.json`, `special **Issue**: Sliding window attention can cause runtime errors with certain sequence lengths. **Solution**: Disable sliding window for initial testing: + ```python def add_derived_config(self): # Disable sliding window for safety during initial testing @@ -1475,6 +1516,7 @@ def add_derived_config(self): **Issue**: Imported `CustomRMSNorm` but model uses standard `LayerNorm`. **Solution**: Check the original model implementation: + ```python # Some models use RMSNorm (Llama, Mistral) from neuronx_distributed_inference.modules.custom_calls import CustomRMSNorm @@ -1497,11 +1539,13 @@ self.norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) **Solution**: Verify architecture in HuggingFace source and implement correctly (see Step 4 for both patterns). **Common Architectures**: + - **Standard FFN**: GPT-2, BERT, some code models → Returns tuple `(output, None)` - **SwiGLU**: Llama, Mistral, Mixtral → Returns single tensor - **GeGLU**: Some T5 variants → Similar to SwiGLU **Common Mistake**: + ```python # ❌ WRONG: Saw one model return single tensor, changed mine def forward(self, x): @@ -1525,6 +1569,7 @@ def forward(self, x): **Result**: Model generated code, not answers (correct behavior for code model, but looks wrong). **Solution**: Test with appropriate prompts for the model type: + ```python # For code models: prompts = [ @@ -1587,6 +1632,7 @@ Port fails? ### Step 1: Context Encoding Fails **Symptoms**: + - Compilation crashes during context encoding - HLO conversion errors - Shape mismatch errors in early compilation @@ -1594,6 +1640,7 @@ Port fails? **Debug Checklist**: 1. **Check Required Attributes**: + ```python # Verify all attributes in get_required_attributes() exist in config.json config = YourModelInferenceConfig.from_pretrained(model_path) @@ -1602,6 +1649,7 @@ Port fails? ``` 2. **Check add_derived_config()**: + ```python # Verify num_cores_per_group is set assert config.num_cores_per_group == 1 @@ -1619,6 +1667,7 @@ Port fails? ``` **Debug Commands**: + ```bash # Check compiler logs ls agent_artifacts/data/neff_output/context_encoding_model/*/log-neuron-cc.txt @@ -1635,6 +1684,7 @@ rm -rf /var/tmp/neuron-compile-cache/* **This is the critical failure mode.** If context encoding works but token generation fails, it's almost always one of the configuration patterns. **Typical Symptom**: + ``` RuntimeError: Expected tensor to have size X at dimension 1, but got size 1 Location: attention_base.py:736 in perform_prefill() @@ -1645,6 +1695,7 @@ Location: attention_base.py:736 in perform_prefill() **Debug Checklist** (in order of likelihood): 1. **✅ Do you have a custom NeuronConfig class?** + ```python # Check if this class exists in your code class YourModelNeuronConfig(NeuronConfig): @@ -1652,17 +1703,21 @@ Location: attention_base.py:736 in perform_prefill() super().__init__(**kwargs) self.attn_cls = YourAttention # MUST SET THIS ``` + **If missing**: This is your problem. Create it (see Pattern 1). 2. **✅ Does get_neuron_config_cls() return your custom class?** + ```python # Check what it returns cls = YourModelInferenceConfig.get_neuron_config_cls() print(cls) # Should print: YourModelNeuronConfig, NOT NeuronConfig ``` + **If wrong**: Fix it to return `YourModelNeuronConfig`. 3. **✅ Does attention pass num_cores_per_group?** + ```python # In your attention __init__, verify this line exists super().__init__( @@ -1670,6 +1725,7 @@ Location: attention_base.py:736 in perform_prefill() num_cores_per_group=config.num_cores_per_group, # MUST PASS THIS ) ``` + **If missing**: Add it. 4. **✅ Does add_derived_config() set all framework attributes?** @@ -1690,6 +1746,7 @@ Location: attention_base.py:736 in perform_prefill() ``` **Debug Commands**: + ```bash # Check token generation logs ls agent_artifacts/data/neff_output/token_generation_model/*/log-neuron-cc.txt @@ -1701,6 +1758,7 @@ cat agent_artifacts/data/neff_output/token_generation_model/*/log-neuron-cc.txt ### Step 3: Weight Loading Fails **Symptoms**: + - Missing parameter errors - Shape mismatch during weight loading - Tied weights errors @@ -1708,6 +1766,7 @@ cat agent_artifacts/data/neff_output/token_generation_model/*/log-neuron-cc.txt **Debug Checklist**: 1. **Check lm_head bias**: + ```python from safetensors import safe_open @@ -1723,6 +1782,7 @@ cat agent_artifacts/data/neff_output/token_generation_model/*/log-neuron-cc.txt ``` 2. **Check tied weights**: + ```python with safe_open("model.safetensors", framework="pt") as f: has_lm_head = "lm_head.weight" in f.keys() @@ -1752,6 +1812,7 @@ cat agent_artifacts/data/neff_output/token_generation_model/*/log-neuron-cc.txt **Error**: `[NLA001] Unhandled exception with message: [json.exception.parse_error.101] parse error at line 1, column 1: attempting to parse an empty input` **Solution**: Delete compiler cache and retry + ```bash rm -rf /var/tmp/neuron-compile-cache/* # Then rerun compilation @@ -1762,6 +1823,7 @@ rm -rf /var/tmp/neuron-compile-cache/* **Error**: `FileNotFoundError: [Errno 2] No such file or directory: 'agent_artifacts/neff_output/token_generation_model/_tp0_bk0'` **Solution**: Delete compiler cache and retry + ```bash rm -rf /var/tmp/neuron-compile-cache/* # Then rerun compilation @@ -1788,6 +1850,7 @@ config = YourModelInferenceConfig.from_pretrained( ``` **Success Criteria**: + - ✅ Context encoding compiles - ✅ Token generation compiles - ✅ Total time: Under 2 minutes @@ -1811,6 +1874,7 @@ config = YourModelInferenceConfig.from_pretrained( ``` **Expected Output**: + ``` ✅ Context encoding: SUCCESS ✅ Token generation: SUCCESS @@ -1851,6 +1915,7 @@ for prompt in prompts: ``` **Quality Indicators**: + - ✅ Text is coherent and relevant to prompt - ✅ Syntax is correct (for code generation) - ✅ No repetitive patterns (same phrase repeated) @@ -1858,6 +1923,7 @@ for prompt in prompts: - ✅ Reasonable performance (varies by model size) **Red Flags**: + - ❌ Generates gibberish or random characters - ❌ Repeats the same token endlessly - ❌ Crashes or throws errors @@ -1870,6 +1936,7 @@ for prompt in prompts: ### Pre-Implementation Checklist **Before writing any code:** + - [ ] Identified 3-4 working models with similar architecture - [ ] Filled component analysis checklist for each model - [ ] Identified ALL required patterns (things all models do) @@ -1883,6 +1950,7 @@ for prompt in prompts: ### Implementation Checklist **Configuration:** + - [ ] Custom NeuronConfig class created - [ ] Custom NeuronConfig sets `self.attn_cls` in `__init__` - [ ] InferenceConfig.get_neuron_config_cls() returns custom class (not base) @@ -1890,9 +1958,10 @@ for prompt in prompts: - [ ] InferenceConfig.add_derived_config() calculates `head_dim` if missing - [ ] InferenceConfig.add_derived_config() sets all 4 framework attributes - [ ] InferenceConfig.get_required_attributes() includes all model-specific attributes -- [ ] InferenceConfig.from_pretrained() creates neuron_config before __init__ +- [ ] InferenceConfig.from_pretrained() creates neuron_config before **init** **Attention:** + - [ ] Attention class inherits from `NeuronAttentionBase` - [ ] Attention `__init__` takes only `config` (no `layer_idx` unless pattern requires) - [ ] Attention passes `num_cores_per_group=config.num_cores_per_group` to super() @@ -1900,11 +1969,13 @@ for prompt in prompts: - [ ] NO custom `forward()` method in attention **MLP:** + - [ ] MLP architecture matches reference model (Standard vs SwiGLU) - [ ] MLP return type matches pattern (tuple for standard, single for SwiGLU) - [ ] Decoder layer handles MLP return correctly **Model Structure:** + - [ ] Base model inherits from `NeuronBaseModel` - [ ] Has `setup_attr_for_model()` method - [ ] Has `init_model()` method @@ -1915,6 +1986,7 @@ for prompt in prompts: - [ ] Correct normalization type (LayerNorm vs RMSNorm) **Wrapper:** + - [ ] ForCausalLM inherits from `NeuronBaseForCausalLM` - [ ] Sets `_model_cls` correctly - [ ] Has `load_hf_model()` static method @@ -1959,46 +2031,54 @@ for prompt in prompts: ## Key Lessons Learned ### 1. Patterns are REQUIRED, not optional + - When ALL working models follow a pattern, it's required by the framework - Don't deviate without strong evidence from multiple reference implementations - Framework patterns aren't style choices - they're requirements ### 2. Start with research, not coding + - Spend 1-2 hours understanding the pattern first - Fill in the component checklist completely for 3-4 models - Then implement exactly following the pattern - This saves 6-8 hours of debugging later ### 3. Don't assume framework limitations + - If other models with similar architecture work, yours can too - Configuration issues look like framework bugs - Compare with working models before assuming limitation - The framework supports 12:1 GQA ratio fine - failures are configuration issues ### 4. The custom NeuronConfig is CRITICAL + - This is the #1 cause of token generation failure - Must exist and must set `attn_cls` - Must be returned by `get_neuron_config_cls()` - Without this, context encoding may work but token generation will fail ### 5. Trust the pattern + - If all working models do it, do it too - If no working models do it, don't add it - Don't add complexity that doesn't exist in the pattern - `layer_idx` is often unnecessary - verify before adding ### 6. Test incrementally + - Compile 1 layer first (20-60 seconds) - If that fails, fix before full compilation - Full compilation wastes time if basic structure is wrong ### 7. Use appropriate test prompts + - Code models: Test with code snippets - Q&A models: Test with questions - Chat models: Test with chat format - Wrong prompt type makes good models look broken ### 8. Framework orchestration is complete + - Don't override `forward()` in NeuronBaseModel for RoPE models (99% of cases) - Exception: MUST override for learned positional embeddings (GPT-2, BERT, RoBERTa) - See: **OVERRIDING_FORWARD_GUIDANCE.md** for decision tree @@ -2059,6 +2139,7 @@ for prompt in prompts: This comprehensive guide combines systematic porting procedures with real-world learnings from both successful and failed ports. The key insight is that **framework patterns are required, not optional**. By following this approach: + 1. **Research first** (understand the pattern by analyzing 3-4 working models) 2. **Implement exactly** (follow the pattern without deviations) 3. **Verify incrementally** (test 1 layer, then full model) diff --git a/skills/neuron-framework-autoport/references/knowledge_base/NOVEL_NEURONX_PORTING_PATTERNS.md b/skills/neuron-framework-autoport/references/knowledge_base/NOVEL_NEURONX_PORTING_PATTERNS.md index 5b4f815..9f04d5b 100644 --- a/skills/neuron-framework-autoport/references/knowledge_base/NOVEL_NEURONX_PORTING_PATTERNS.md +++ b/skills/neuron-framework-autoport/references/knowledge_base/NOVEL_NEURONX_PORTING_PATTERNS.md @@ -15,21 +15,22 @@ This document captures advanced patterns discovered during model porting that ar def convert_hf_to_neuron_state_dict(state_dict: dict, config: InferenceConfig) -> dict: # ALWAYS do this first when porting a new model print(f"DEBUG: First 10 keys received: {list(state_dict.keys())[:10]}") - + neuron_state_dict = {} - + # Now write conversion logic based on ACTUAL keys you see for key, value in state_dict.items(): new_key = key # Add your transformations here based on debug output neuron_state_dict[new_key] = value - + return neuron_state_dict ``` ### Common Prefix Patterns Different models use different checkpoint structures: + - `"model.decoder.layers.0.weight"` - Decoder-only models with decoder prefix - `"model.layers.0.weight"` - Standard models (LLaMA, Mistral) - `"decoder.layers.0.weight"` - Some variants @@ -105,13 +106,13 @@ Even though your checkpoint has `layers.0.self_attn.q_proj.weight`. @staticmethod def convert_hf_to_neuron_state_dict(state_dict: dict, config: InferenceConfig) -> dict: neuron_state_dict = {} - + # First pass: copy all keys with basic transformations for key, value in state_dict.items(): new_key = key # Remove prefixes, rename projections, etc. neuron_state_dict[new_key] = value - + # Second pass: restructure QKV weights per layer num_layers = config.num_hidden_layers for i in range(num_layers): @@ -121,12 +122,12 @@ def convert_hf_to_neuron_state_dict(state_dict: dict, config: InferenceConfig) - q_weight = neuron_state_dict.pop(f"layers.{i}.self_attn.q_proj.weight") k_weight = neuron_state_dict.pop(f"layers.{i}.self_attn.k_proj.weight") v_weight = neuron_state_dict.pop(f"layers.{i}.self_attn.v_proj.weight") - + # Add with qkv_proj intermediate level neuron_state_dict[f"layers.{i}.self_attn.qkv_proj.q_proj.weight"] = q_weight neuron_state_dict[f"layers.{i}.self_attn.qkv_proj.k_proj.weight"] = k_weight neuron_state_dict[f"layers.{i}.self_attn.qkv_proj.v_proj.weight"] = v_weight - + # Handle biases if present if f"layers.{i}.self_attn.q_proj.bias" in neuron_state_dict: q_bias = neuron_state_dict.pop(f"layers.{i}.self_attn.q_proj.bias") @@ -135,7 +136,7 @@ def convert_hf_to_neuron_state_dict(state_dict: dict, config: InferenceConfig) - neuron_state_dict[f"layers.{i}.self_attn.qkv_proj.q_proj.bias"] = q_bias neuron_state_dict[f"layers.{i}.self_attn.qkv_proj.k_proj.bias"] = k_bias neuron_state_dict[f"layers.{i}.self_attn.qkv_proj.v_proj.bias"] = v_bias - + return neuron_state_dict ``` @@ -176,14 +177,17 @@ print(sorted(attn_keys)) ## 2. Learned Positional Embeddings Pattern ### Issue + Models using learned positional embeddings (not RoPE or relative position bias) need to add positional information in the forward pass. ### Background + - **RoPE models** (LLaMA, Mistral, Qwen): Position info added in attention layers - **Relative position models** (T5): Position bias computed in attention - **Learned position models** (GPT-2, BERT, RoBERTa): Position embeddings must be added to token embeddings ### Discovery + The base model forward pass calls `self.embed_tokens(input_ids)` but doesn't add positional embeddings. For models with learned positional embeddings, you must add them yourself. ### ⚠️ CRITICAL: This is an EXCEPTION to "Don't Override Forward" Rule @@ -202,7 +206,7 @@ class HFModel(nn.Module): def __init__(self, config): self.embed_tokens = nn.Embedding(vocab_size, hidden_size) self.embed_positions = nn.Embedding(max_positions, hidden_size) # ← Learned! - + def forward(self, input_ids): token_embeds = self.embed_tokens(input_ids) position_embeds = self.embed_positions(position_ids) @@ -225,7 +229,7 @@ def init_model(self, config: InferenceConfig): config.hidden_size, dtype=config.neuron_config.torch_dtype, ) - + # Positional embeddings (separate, not wrapped) self.embed_positions = ParallelEmbedding( config.max_position_embeddings + 2, # +2 if model uses offset @@ -233,7 +237,7 @@ def init_model(self, config: InferenceConfig): None, # No padding_idx dtype=config.neuron_config.torch_dtype, ) - + # ... rest of model initialization def forward(self, input_ids, position_ids=None, ...): @@ -241,10 +245,10 @@ def forward(self, input_ids, position_ids=None, ...): # We need to add positional embeddings if inputs_embeds is None and input_ids is not None: batch_size, seq_length = input_ids.shape - + # Get token embeddings inputs_embeds = self.embed_tokens(input_ids) - + # Generate position_ids if not provided if position_ids is None: device = input_ids.device @@ -252,13 +256,13 @@ def forward(self, input_ids, position_ids=None, ...): position_ids = position_ids.unsqueeze(0).expand(batch_size, -1) else: position_ids = position_ids.view(-1, seq_length).long() - + # Get positional embeddings (add offset during lookup if model uses offset) position_embeddings = self.embed_positions(position_ids + offset) # offset=2 for some models - + # Combine token and positional embeddings inputs_embeds = inputs_embeds + position_embeddings - + # Continue with rest of forward pass... return super().forward( input_ids=input_ids, @@ -304,7 +308,7 @@ def forward( inputs_embeds = self.embed_tokens(input_ids) position_embeds = self.embed_positions(position_ids + 2) inputs_embeds = inputs_embeds + position_embeds - + # Pass ALL parameters to parent return super().forward( input_ids=input_ids, @@ -339,6 +343,7 @@ def forward( #### Mistake 1: Wrong `get_input_embeddings()` Signature ❌ **WRONG** - Custom signature breaks the model: + ```python def get_input_embeddings(self, input_ids, position_ids): # This is WRONG - get_input_embeddings() takes NO parameters! @@ -346,6 +351,7 @@ def get_input_embeddings(self, input_ids, position_ids): ``` ✅ **CORRECT** - Simple getter with no parameters: + ```python def get_input_embeddings(self): # Just return the embedding layer itself @@ -357,6 +363,7 @@ def get_input_embeddings(self): #### Mistake 2: Using nn.Embedding Instead of ParallelEmbedding ❌ **WRONG** - Regular PyTorch embedding: + ```python self.embed_positions = nn.Embedding( config.max_position_embeddings, @@ -365,6 +372,7 @@ self.embed_positions = nn.Embedding( ``` ✅ **CORRECT** - Use ParallelEmbedding for distributed training: + ```python self.embed_positions = ParallelEmbedding( config.max_position_embeddings + 2, @@ -378,12 +386,14 @@ self.embed_positions = ParallelEmbedding( #### Mistake 3: Incomplete Forward Signature ❌ **WRONG** - Missing parameters: + ```python def forward(self, input_ids, attention_mask, position_ids, **kwargs): # Using **kwargs is fragile and can cause issues ``` ✅ **CORRECT** - Explicit parameters matching base class: + ```python def forward( self, @@ -395,7 +405,8 @@ def forward( # ... all 26 parameters explicitly listed ): ``` -``` + +```` ### Weight Conversion for Separate Embeddings @@ -405,21 +416,21 @@ With separate embeddings (no wrapper), weight conversion is straightforward: @staticmethod def convert_hf_to_neuron_state_dict(state_dict: dict, config: InferenceConfig) -> dict: neuron_state_dict = {} - + for key, value in state_dict.items(): new_key = key - + # Remove prefixes new_key = new_key.replace("model.decoder.", "").replace("decoder.", "") - + # Keys remain flat - no nesting needed: # "embed_tokens.weight" stays as "embed_tokens.weight" # "embed_positions.weight" stays as "embed_positions.weight" - + neuron_state_dict[new_key] = value - + return neuron_state_dict -``` +```` ### Tied Weights Handling @@ -434,9 +445,11 @@ def update_state_dict_for_tied_weights(state_dict): ``` --- + state_dict["lm_head.weight"] = state_dict["embed_tokens.token_embedding.weight"].clone() return state_dict -``` + +```` --- @@ -464,7 +477,7 @@ position_embeddings = self.embed_positions(position_ids + 2) # Offset here! # ❌ WRONG - Don't add offset to position_ids before storing position_ids = torch.arange(2, seq_length + 2, ...) # Wrong! position_embeddings = self.embed_positions(position_ids) -``` +```` ### Implementation @@ -482,7 +495,7 @@ def forward(self, input_ids, position_ids=None, ...): if position_ids is None: position_ids = torch.arange(0, seq_length, dtype=torch.long, device=device) position_ids = position_ids.unsqueeze(0).expand(batch_size, -1) - + # Add offset during embedding lookup position_embeddings = self.embed_positions(position_ids + 2) ``` @@ -517,6 +530,7 @@ assert torch.allclose(hf_out.logits, neuron_out.logits, atol=1e-2) ## 4. Common Weight Renaming Patterns ### Issue + Different models use different naming conventions that must be mapped to framework expectations. ### Common Renamings @@ -525,23 +539,23 @@ Different models use different naming conventions that must be mapped to framewo @staticmethod def convert_hf_to_neuron_state_dict(state_dict: dict, config: InferenceConfig) -> dict: neuron_state_dict = {} - + for key, value in state_dict.items(): new_key = key - + # 1. Remove model-specific prefixes new_key = new_key.replace("model.decoder.", "").replace("decoder.", "") - + # 2. Rename output projection (CRITICAL - often missed) new_key = new_key.replace("out_proj", "o_proj") - + # 3. Rename top-level final norm (but not per-layer norms) if new_key.startswith("final_layer_norm"): new_key = new_key.replace("final_layer_norm", "norm") # Per-layer final_layer_norm stays unchanged - + neuron_state_dict[new_key] = value - + return neuron_state_dict ``` @@ -590,27 +604,30 @@ if new_key.startswith("layers.") and ".final_layer_norm." in new_key: ### When to Use These Patterns -| Pattern | Use When | -|---------|----------| -| Checkpoint preprocessing check | Any model port | -| Embedding wrapper | Model uses learned positional embeddings | -| Positional offset | Model reserves embedding indices | -| Nested key mapping | Using wrapper modules | -| Prefix handling | Model uses non-standard checkpoint prefixes | -| MLP submodule mapping | Framework expects different MLP hierarchy | -| Selective renaming | Different rules for top-level vs per-layer norms | -| Separate projection mapping | Model has unfused q/k/v projections | -| Explicit tied weights | Model ties embeddings with lm_head | +| Pattern | Use When | +| ------------------------------ | ------------------------------------------------ | +| Checkpoint preprocessing check | Any model port | +| Embedding wrapper | Model uses learned positional embeddings | +| Positional offset | Model reserves embedding indices | +| Nested key mapping | Using wrapper modules | +| Prefix handling | Model uses non-standard checkpoint prefixes | +| MLP submodule mapping | Framework expects different MLP hierarchy | +| Selective renaming | Different rules for top-level vs per-layer norms | +| Separate projection mapping | Model has unfused q/k/v projections | +| Explicit tied weights | Model ties embeddings with lm_head | --- ## 5. Model-Specific Checkpoint Prefixes ### Issue + Different model families use different checkpoint key prefixes beyond the standard `model.` prefix. ### Discovery + Common prefix patterns: + - **Most models**: `model.layers.0.weight` - **Decoder-only variants**: `decoder.layers.0.weight` - **Encoder-decoder**: `encoder.layers.0.weight`, `decoder.layers.0.weight` @@ -624,10 +641,10 @@ The base class automatically removes the `model.` prefix and calls `convert_hf_t @staticmethod def convert_hf_to_neuron_state_dict(hf_state_dict: dict, config: InferenceConfig) -> dict: neuron_state_dict = {} - + for key, value in hf_state_dict.items(): new_key = key - + # Remove model-specific prefix # Check what prefix actually exists in the checkpoint if new_key.startswith("decoder."): @@ -635,10 +652,10 @@ def convert_hf_to_neuron_state_dict(hf_state_dict: dict, config: InferenceConfig elif new_key.startswith("encoder."): new_key = new_key.replace("encoder.", "", 1) # Note: "model." prefix already removed by base class - + # Continue with other transformations... neuron_state_dict[new_key] = value - + return neuron_state_dict ``` @@ -652,7 +669,7 @@ def convert_hf_to_neuron_state_dict(state_dict: dict, config: InferenceConfig) - # Debug: Check what keys we receive after base class preprocessing print("Keys after base class:", list(state_dict.keys())[:10]) # Now you know what prefix to remove! - + neuron_state_dict = {} for key, value in state_dict.items(): new_key = key @@ -688,10 +705,13 @@ else: ## 5. MLP Weight Hierarchy Mapping ### Issue + Some models have MLP weights directly under the layer, but the framework expects them under an `mlp` submodule. ### Discovery + HuggingFace structure: + ``` layers.0.fc1.weight layers.0.fc1.bias @@ -700,6 +720,7 @@ layers.0.fc2.bias ``` Framework expects: + ``` layers.0.mlp.fc1.weight layers.0.mlp.fc1.bias @@ -715,19 +736,19 @@ layers.0.mlp.fc2.bias @staticmethod def convert_hf_to_neuron_state_dict(state_dict: dict, config: InferenceConfig) -> dict: neuron_state_dict = {} - + for key, value in state_dict.items(): new_key = key - + # Add mlp prefix to fc1/fc2 weights if ".fc1." in new_key or ".fc2." in new_key: parts = new_key.split(".") layer_idx = parts[1] # Extract layer number fc_part = ".".join(parts[2:]) # Get fc1/fc2 and weight/bias new_key = f"layers.{layer_idx}.mlp.{fc_part}" - + neuron_state_dict[new_key] = value - + return neuron_state_dict ``` @@ -768,10 +789,13 @@ print(f"MLP keys: {sorted(mlp_keys)}") ## 8. Selective Component Renaming ### Issue + Some components need renaming at the top level but not at the per-layer level, or vice versa. ### Discovery + Example: Layer normalization + - Top-level final norm: `final_layer_norm` → `norm` (framework expects `norm`) - Per-layer norms: `layers.X.final_layer_norm` → keep as-is (framework expects original name) @@ -781,17 +805,17 @@ Example: Layer normalization @staticmethod def convert_hf_to_neuron_state_dict(hf_state_dict: dict, config: InferenceConfig) -> dict: neuron_state_dict = {} - + for key, value in hf_state_dict.items(): new_key = key - + # Rename top-level final norm only if new_key == "final_layer_norm.weight" or new_key == "final_layer_norm.bias": new_key = new_key.replace("final_layer_norm.", "norm.") # Per-layer final_layer_norm stays unchanged - + neuron_state_dict[new_key] = value - + return neuron_state_dict ``` @@ -836,17 +860,21 @@ elif ".post_attention_layernorm." in new_key: ## 9. Separate Projection Weight Structure (fused_qkv=False) ### Issue + When `fused_qkv=False`, the framework creates a specific weight hierarchy that differs from both fused QKV and completely separate projections. ### Discovery + Three different structures: **1. Fused QKV (fused_qkv=True):** + ``` layers.0.self_attn.qkv_proj.Wqkv.weight # Single fused weight ``` **2. Completely Separate (incorrect for framework):** + ``` layers.0.self_attn.q_proj.weight layers.0.self_attn.k_proj.weight @@ -854,6 +882,7 @@ layers.0.self_attn.v_proj.weight ``` **3. Framework's fused_qkv=False (correct):** + ``` layers.0.self_attn.qkv_proj.q_proj.weight layers.0.self_attn.qkv_proj.k_proj.weight @@ -868,10 +897,10 @@ The `qkv_proj` is a `GroupQueryAttention_QKV` instance that has `q_proj`, `k_pro @staticmethod def convert_hf_to_neuron_state_dict(hf_state_dict: dict, config: InferenceConfig) -> dict: neuron_state_dict = {} - + for key, value in hf_state_dict.items(): new_key = key - + # Map separate projections to qkv_proj structure # Must handle both .weight and .bias new_key = new_key.replace(".self_attn.q_proj.weight", ".self_attn.qkv_proj.q_proj.weight") @@ -880,12 +909,12 @@ def convert_hf_to_neuron_state_dict(hf_state_dict: dict, config: InferenceConfig new_key = new_key.replace(".self_attn.k_proj.bias", ".self_attn.qkv_proj.k_proj.bias") new_key = new_key.replace(".self_attn.v_proj.weight", ".self_attn.qkv_proj.v_proj.weight") new_key = new_key.replace(".self_attn.v_proj.bias", ".self_attn.qkv_proj.v_proj.bias") - + # Output projection also needs mapping new_key = new_key.replace(".self_attn.out_proj.", ".self_attn.o_proj.") - + neuron_state_dict[new_key] = value - + return neuron_state_dict ``` @@ -920,16 +949,20 @@ print("Attention keys:", sorted(attn_keys)) ## 6. Explicit Tied Weight Creation ### Issue + Even when a model ties weights implicitly (e.g., `lm_head` shares weights with `embed_tokens`), the compiled checkpoint must have BOTH keys explicitly present. ### Discovery + HuggingFace models often tie weights by reference: + ```python # In HuggingFace model self.lm_head.weight = self.embed_tokens.weight # Same object ``` But the framework's weight loading expects both keys in the state dict: + ```python # Framework expects state_dict = { @@ -944,13 +977,13 @@ state_dict = { @staticmethod def convert_hf_to_neuron_state_dict(state_dict: dict, config: InferenceConfig) -> dict: neuron_state_dict = {} - + # ... other conversions ... - + # Handle tied embeddings - must use .clone() if "lm_head.weight" not in neuron_state_dict and "embed_tokens.weight" in neuron_state_dict: neuron_state_dict["lm_head.weight"] = neuron_state_dict["embed_tokens.weight"].clone() - + return neuron_state_dict ``` @@ -996,10 +1029,13 @@ assert state_dict["embed_tokens.weight"].data_ptr() != state_dict["lm_head.weigh ## 7. Framework Config Attributes for Inference ### Issue + The framework expects certain standard HuggingFace config attributes to exist during inference, even if they're not actively used by the model. ### Discovery + During inference, the framework's base classes check for these attributes: + - `output_attentions` - `output_hidden_states` - `use_return_dict` @@ -1008,6 +1044,7 @@ During inference, the framework's base classes check for these attributes: Missing these causes `AttributeError` during model execution. ### Error You'll See + ``` AttributeError: 'ModelInferenceConfig' object has no attribute 'output_attentions' ``` @@ -1022,7 +1059,7 @@ class ModelInferenceConfig(InferenceConfig): self.num_cores_per_group = 1 if not hasattr(self, 'head_dim'): self.head_dim = self.hidden_size // self.num_attention_heads - + # Framework-required attributes for inference if not hasattr(self, 'output_attentions'): self.output_attentions = False @@ -1060,9 +1097,11 @@ assert hasattr(config, 'use_cache') ## 12. Position IDs Computation for Learned Embeddings ### Issue + Models with learned positional embeddings need proper position_ids computation, especially for autoregressive generation where past_key_value is used. ### Discovery + Unlike RoPE (computed in attention layers), learned positional embeddings require position_ids to be computed in the model's forward pass before embedding lookup. ### Pattern for Context Encoding (Prefill) @@ -1070,26 +1109,26 @@ Unlike RoPE (computed in attention layers), learned positional embeddings requir ```python def forward(self, input_ids, attention_mask=None, position_ids=None, past_key_value=None): batch_size, seq_length = input_ids.shape - + if position_ids is None: # Determine starting position based on past context past_length = 0 if past_key_value is not None and len(past_key_value) > 0: past_length = past_key_value[0][0].shape[2] # KV cache sequence length - + # Create position_ids starting from past_length device = input_ids.device position_ids = torch.arange( - past_length, - seq_length + past_length, - dtype=torch.long, + past_length, + seq_length + past_length, + dtype=torch.long, device=device ) position_ids = position_ids.unsqueeze(0).expand(batch_size, -1) - + # Get embeddings with positions hidden_states = self.embed_tokens(input_ids, position_ids) - + # Continue with decoder layers... ``` @@ -1101,7 +1140,7 @@ for step in range(max_new_tokens): # position_ids should be [past_length + step] current_position = past_length + step position_ids = torch.tensor([[current_position]], dtype=torch.long, device=device) - + outputs = model( input_ids=next_token_id, position_ids=position_ids, @@ -1139,7 +1178,6 @@ position_ids = compute_position_ids(input_ids, past_key_value=past_kv) assert position_ids.tolist() == [[10, 11, 12, 13, 14]] ``` - - See "Learned Positional Embeddings Pattern" (Section 2) - See "Positional Embedding Offsets" (Section 3) diff --git a/skills/neuron-framework-autoport/references/knowledge_base/OVERRIDING_FORWARD_GUIDANCE.md b/skills/neuron-framework-autoport/references/knowledge_base/OVERRIDING_FORWARD_GUIDANCE.md index ed2513a..5061260 100644 --- a/skills/neuron-framework-autoport/references/knowledge_base/OVERRIDING_FORWARD_GUIDANCE.md +++ b/skills/neuron-framework-autoport/references/knowledge_base/OVERRIDING_FORWARD_GUIDANCE.md @@ -69,7 +69,7 @@ class HFModel(nn.Module): def __init__(self, config): self.embed_tokens = nn.Embedding(vocab_size, hidden_size) self.embed_positions = nn.Embedding(max_positions, hidden_size) # ← Learned! - + def forward(self, input_ids): token_embeds = self.embed_tokens(input_ids) position_embeds = self.embed_positions(position_ids) @@ -80,6 +80,7 @@ class HFModel(nn.Module): ### Models with Learned Positional Embeddings **Confirmed models that NEED forward() override:** + - GPT-2 (gpt2, gpt2-medium, gpt2-large, gpt2-xl) - BERT (bert-base-uncased, bert-large-uncased) - RoBERTa (roberta-base, roberta-large) @@ -87,12 +88,13 @@ class HFModel(nn.Module): - Any model with learned positional embeddings (check HuggingFace implementation) **Models that DO NOT need forward() override:** -- LLaMA family (meta-llama/Llama-2-*, meta-llama/Llama-3-*) -- Mistral (mistralai/Mistral-7B-*) -- Qwen (Qwen/Qwen-*, Qwen/Qwen2-*) -- Gemma (google/gemma-*) -- Mixtral (mistralai/Mixtral-8x7B-*) -- DBRX (databricks/dbrx-*) + +- LLaMA family (meta-llama/Llama-2-_, meta-llama/Llama-3-_) +- Mistral (mistralai/Mistral-7B-\*) +- Qwen (Qwen/Qwen-_, Qwen/Qwen2-_) +- Gemma (google/gemma-\*) +- Mixtral (mistralai/Mixtral-8x7B-\*) +- DBRX (databricks/dbrx-\*) - T5 (t5-small, t5-base, t5-large) ## Survey of Existing Implementations @@ -101,18 +103,18 @@ class HFModel(nn.Module): All surveyed models in `/NeuronxDistributedInference/src/neuronx_distributed_inference/models/`: -| Model | Overrides forward()? | Position Encoding Type | -|-------|---------------------|------------------------| -| llama | ❌ NO | RoPE | -| mistral | ❌ NO | RoPE | -| qwen2 | ❌ NO | RoPE | -| qwen3 | ❌ NO | RoPE | -| qwen3_moe | ❌ NO | RoPE | -| mixtral | ❌ NO | RoPE | -| dbrx | ❌ NO | RoPE | -| gpt_oss | ❌ NO | RoPE | -| llama4 | ❌ NO | RoPE | -| mllama | ❌ NO | RoPE (multimodal) | +| Model | Overrides forward()? | Position Encoding Type | +| --------- | -------------------- | ---------------------- | +| llama | ❌ NO | RoPE | +| mistral | ❌ NO | RoPE | +| qwen2 | ❌ NO | RoPE | +| qwen3 | ❌ NO | RoPE | +| qwen3_moe | ❌ NO | RoPE | +| mixtral | ❌ NO | RoPE | +| dbrx | ❌ NO | RoPE | +| gpt_oss | ❌ NO | RoPE | +| llama4 | ❌ NO | RoPE | +| mllama | ❌ NO | RoPE (multimodal) | **Finding**: None of the existing models override `forward()` because they all use RoPE. @@ -120,22 +122,22 @@ All surveyed models in `/NeuronxDistributedInference/src/neuronx_distributed_inf All surveyed models in `/NeuroborosFoundations/src/amzn/neuron/neuroboros/models/`: -| Model | Overrides forward()? | Position Encoding Type | -|-------|---------------------|------------------------| -| gemma3 | ❌ NO | RoPE | -| gpt2 | ❌ NO | RoPE (modified) | -| gptoss | ❌ NO | RoPE | -| phi3 | ❌ NO | RoPE | -| phimoe | ❌ NO | RoPE | -| starcoder2 | ❌ NO | RoPE | +| Model | Overrides forward()? | Position Encoding Type | +| ---------- | -------------------- | ---------------------- | +| gemma3 | ❌ NO | RoPE | +| gpt2 | ❌ NO | RoPE (modified) | +| gptoss | ❌ NO | RoPE | +| phi3 | ❌ NO | RoPE | +| phimoe | ❌ NO | RoPE | +| starcoder2 | ❌ NO | RoPE | **Finding**: None of these models override `forward()` either. Note that `gpt2` in this collection appears to be a modified version using RoPE, not the original GPT-2 with learned positional embeddings. ### Example: Learned Positional Embedding Model -| Model Type | Overrides forward()? | Position Encoding Type | -|-------|---------------------|------------------------| -| Learned Positional Embeddings | ✅ YES | Learned positional embeddings | +| Model Type | Overrides forward()? | Position Encoding Type | +| ----------------------------- | -------------------- | ----------------------------- | +| Learned Positional Embeddings | ✅ YES | Learned positional embeddings | **Models with learned positional embeddings are the first type in the codebase that require forward() override.** @@ -203,6 +205,7 @@ All surveyed models in `/NeuroborosFoundations/src/amzn/neuron/neuroboros/models - **Exception (1% of models)**: MUST override forward() - applies to learned positional embedding models The confusion arises because: + 1. Most existing models use RoPE, so "don't override" is the common case 2. The exception case is documented but not cross-referenced in general guidance 3. Learned positional embedding models are rare in the current codebase @@ -221,7 +224,7 @@ def init_model(self, config: InferenceConfig): dtype=config.neuron_config.torch_dtype, shard_across_embedding=not config.neuron_config.vocab_parallel, ) - + # Positional embeddings (separate, not wrapped) self.embed_positions = ParallelEmbedding( config.max_position_embeddings + offset, # Add offset if model uses it @@ -230,7 +233,7 @@ def init_model(self, config: InferenceConfig): dtype=config.neuron_config.torch_dtype, shard_across_embedding=False, # Don't shard position embeddings ) - + # ... rest of model initialization (layers, norm, lm_head) ``` @@ -243,6 +246,7 @@ def get_input_embeddings(self): ``` **Common Mistake**: Do NOT override with custom signature: + ```python # ❌ WRONG - This breaks the model! def get_input_embeddings(self, input_ids, position_ids): @@ -281,18 +285,18 @@ def forward( ): """ Override forward to add positional embeddings to token embeddings. - + This is required for models with learned positional embeddings (GPT-2, BERT, etc). """ # Only compute embeddings if not already provided if inputs_embeds is None: # Get token embeddings inputs_embeds = self.embed_tokens(input_ids) - + # Add positional embeddings with offset (if applicable) position_embeds = self.embed_positions(position_ids + offset) # offset varies by model inputs_embeds = inputs_embeds + position_embeds - + # Pass ALL parameters to parent forward return super().forward( input_ids=input_ids, @@ -328,20 +332,20 @@ def forward( @staticmethod def convert_hf_to_neuron_state_dict(state_dict: dict, config: InferenceConfig) -> dict: neuron_state_dict = {} - + for key, value in state_dict.items(): new_key = key - + # Remove model-specific prefixes if new_key.startswith('decoder.'): new_key = new_key.replace('decoder.', '', 1) - + # Keys remain flat - no nesting needed: # "embed_tokens.weight" stays as "embed_tokens.weight" # "embed_positions.weight" stays as "embed_positions.weight" - + neuron_state_dict[new_key] = value - + return neuron_state_dict ``` @@ -350,12 +354,14 @@ def convert_hf_to_neuron_state_dict(state_dict: dict, config: InferenceConfig) - ### Mistake 1: Wrong get_input_embeddings() Signature ❌ **WRONG**: + ```python def get_input_embeddings(self, input_ids, position_ids): return self.embed_tokens(input_ids) + self.embed_positions(position_ids) ``` ✅ **CORRECT**: + ```python def get_input_embeddings(self): return self.embed_tokens @@ -366,11 +372,13 @@ def get_input_embeddings(self): ### Mistake 2: Using nn.Embedding Instead of ParallelEmbedding ❌ **WRONG**: + ```python self.embed_positions = nn.Embedding(max_positions, hidden_size) ``` ✅ **CORRECT**: + ```python self.embed_positions = ParallelEmbedding( max_positions, @@ -386,12 +394,14 @@ self.embed_positions = ParallelEmbedding( ### Mistake 3: Incomplete Forward Signature ❌ **WRONG**: + ```python def forward(self, input_ids, attention_mask, position_ids, **kwargs): # Using **kwargs is fragile ``` ✅ **CORRECT**: + ```python def forward( self, @@ -409,6 +419,7 @@ def forward( ### Mistake 4: Not Passing inputs_embeds to Parent ❌ **WRONG**: + ```python def forward(self, input_ids, ...): inputs_embeds = self.embed_tokens(input_ids) + self.embed_positions(position_ids) @@ -417,6 +428,7 @@ def forward(self, input_ids, ...): ``` ✅ **CORRECT**: + ```python def forward(self, input_ids, ...): inputs_embeds = self.embed_tokens(input_ids) + self.embed_positions(position_ids) @@ -432,6 +444,7 @@ def forward(self, input_ids, ...): **Possible Cause**: Learned positional embeddings not being added **Check**: + 1. Does HuggingFace model have `embed_positions`? 2. Is it `nn.Embedding` (learned) or computed (RoPE)? @@ -454,10 +467,12 @@ def forward(self, input_ids, ...): ### When to Override forward() in NeuronBaseModel ✅ **DO Override** if: + - Model uses learned positional embeddings (GPT-2, BERT, RoBERTa, ALBERT) - Model requires custom embedding preprocessing (multimodal) ❌ **DON'T Override** if: + - Model uses RoPE (LLaMA, Mistral, Qwen, Gemma) - 99% of models - Model uses relative position bias (T5, BART) - Model is standard transformer architecture diff --git a/skills/neuron-framework-autoport/references/knowledge_base/PORTING_SLIDING_WINDOW.md b/skills/neuron-framework-autoport/references/knowledge_base/PORTING_SLIDING_WINDOW.md index c060b8e..0ae5f81 100644 --- a/skills/neuron-framework-autoport/references/knowledge_base/PORTING_SLIDING_WINDOW.md +++ b/skills/neuron-framework-autoport/references/knowledge_base/PORTING_SLIDING_WINDOW.md @@ -1,4 +1,5 @@ # AWS Neuron Model Porting: Lessons Learned + ## GenericModel Port - Complete Analysis and Solutions **Date**: 2025-11-14 @@ -13,6 +14,7 @@ Successfully ported GenericModel (3B parameters) from PyTorch/CUDA to AWS Neuron/Trainium. The port encountered four critical issues during compilation and inference, all successfully resolved. The model now generates coherent text at ~15 tokens/second. **Key Statistics:** + - Model Size: 3B parameters - Compilation Time: ~360 seconds (from scratch) - Inference Speed: 14.9-15.0 tokens/second @@ -24,6 +26,7 @@ Successfully ported GenericModel (3B parameters) from PyTorch/CUDA to AWS Neuron ## Model Architecture Characteristics ### Configuration + - **Hidden Size**: 3072 - **Attention Heads**: 24 (query) - **KV Heads**: 2 (grouped query attention, 12:1 ratio) @@ -35,6 +38,7 @@ Successfully ported GenericModel (3B parameters) from PyTorch/CUDA to AWS Neuron - **RoPE Theta**: 999999.44 (extremely high) ### Architecture Differences from LLaMA + 1. **Normalization**: Uses LayerNorm instead of RMSNorm 2. **Activation**: GELU (pytorch_tanh) instead of SwiGLU 3. **Bias**: Uses bias in all linear layers @@ -47,20 +51,24 @@ Successfully ported GenericModel (3B parameters) from PyTorch/CUDA to AWS Neuron ## Issue 1: Missing lm_head.weight during Weight Loading ### Symptoms + ``` RuntimeError: Missing weight tensor with key lm_head.weight ``` ### Root Cause + GenericModel uses **tied embeddings** - the embedding layer (`model.embed_tokens.weight`) and language model head (`lm_head.weight`) share the same weight tensor. The HuggingFace checkpoint only stores one copy as `model.embed_tokens.weight`, but the Neuron model initialization expects both keys in the state dict. The configuration was not properly indicating weight tying: + ```python # WRONG - read from HF config which defaults to False "tie_word_embeddings": hf_config.get("tie_word_embeddings", False) ``` ### Solution + Modified `from_pretrained()` method in the inference config to always set `tie_word_embeddings=True`: ```python @@ -72,10 +80,13 @@ Modified `from_pretrained()` method in the inference config to always set `tie_w This ensures the framework's `update_state_dict_for_tied_weights()` method copies `embed_tokens.weight` to `lm_head.weight` during model loading. ### Key Learning + **Always verify weight tying behavior by inspecting the checkpoint files**, not just the config. Use: + ```bash python -c "import torch; print(torch.load('pytorch_model.bin').keys())" ``` + If `lm_head.weight` is missing but `embed_tokens.weight` is present, embeddings are tied. --- @@ -83,14 +94,17 @@ If `lm_head.weight` is missing but `embed_tokens.weight` is present, embeddings ## Issue 2: Missing Framework-Required Config Attributes ### Symptoms + ``` AttributeError: 'GenericModelInferenceConfig' object has no attribute 'output_attentions' ``` ### Root Cause + NeuronxDistributedInference framework expects certain attributes to exist on the config object during model execution, even if they're not used. These attributes (`output_attentions`, `output_hidden_states`, `use_return_dict`, `use_cache`) are standard in HuggingFace Transformers but weren't defined in our custom config. ### Solution + Added framework-required attributes to `add_derived_config()` method: ```python @@ -113,7 +127,9 @@ def add_derived_config(self): ``` ### Key Learning + **Always check framework base classes for required attributes**. Look at: + - `NeuronBaseModel.__init__()` - `InferenceConfig` parent class - Similar model implementations in the framework @@ -125,12 +141,14 @@ Use defensive `hasattr()` checks to avoid overwriting intentionally set values. ## Issue 3: Wrong Attention Class in Compiled Model ### Symptoms + - Model compiled successfully (exit code 0) - Model loaded without errors - Runtime error during inference: Out-of-bounds memory access - Investigation revealed compiled model was using `NeuronLlamaAttention` instead of `NeuronGenericModelAttention` ### Root Cause + Compilation script used base `NeuronConfig` instead of model-specific `GenericModelNeuronConfig`: ```python @@ -149,13 +167,17 @@ compile_neuron_model( The `NeuronConfig` doesn't specify `attn_cls`, so the framework defaulted to `NeuronLlamaAttention`, which has different behavior than GenericModel's attention mechanism. ### Investigation Process + 1. Checked compiled model's neuron_config.json: + ```bash cat agent_artifacts/data/neff_output/context_encoding_model/_tp0_bk0/neuron_config.json ``` + 2. Found: `"attn_cls": "NeuronLlamaAttention"` instead of expected `"attn_cls": {"__module__": "modeling_genericmodel", "__name__": "NeuronGenericModelAttention"}` ### Solution + Created model-specific `GenericModelNeuronConfig` that sets the correct attention class: ```python @@ -171,6 +193,7 @@ class GenericModelNeuronConfig(NeuronConfig): ``` Updated compilation script: + ```python # File: compile_genericmodel.py, Lines 17-18 from modeling_genericmodel import ( @@ -187,11 +210,14 @@ compile_neuron_model( ``` ### Key Learning + **Always create a model-specific NeuronConfig subclass** that sets: + 1. `attn_cls` - The attention class for your model 2. Any other model-specific compilation parameters **Verify compiled artifacts** after compilation: + ```bash # Check attention class in compiled config cat compiled_path/neuron_config.json | grep -A5 "attn_cls" @@ -202,6 +228,7 @@ cat compiled_path/neuron_config.json | grep -A5 "attn_cls" ## Issue 4: Out-of-Bounds Memory Access in Sliding Window Attention ### Symptoms + ``` RuntimeError: Failed to execute the model status=1006 message=Execution Out-Of-Bounds Memory Access @@ -212,6 +239,7 @@ Received notification generated at runtime: failed to run scatter/gather ``` Error occurred during: + - Context encoding (prefill phase) - First inference call with prompt - Warmup phase showed error but was marked as "safe to ignore" by framework @@ -222,6 +250,7 @@ Error occurred during: The `get_last_kv_window` function in NeuronxDistributedInference assumes K/V tensors are at least `window_size` long. This is violated during context encoding with short prompts when `sliding_window > actual_sequence_length`. **Detailed Breakdown:** + 1. GenericModel configuration: - `sliding_window = 4096` - Compiled with `seq_len = 512` @@ -232,6 +261,7 @@ The `get_last_kv_window` function in NeuronxDistributedInference assumes K/V ten - For a 6-token prompt: `[1, 24, 6, 128]` 3. `get_last_kv_window` execution: + ```python # File: NeuronxDistributedInference/.../attention/utils.py:641 orig_indices = start_idx[:, None] + torch.arange(window_size) @@ -245,6 +275,7 @@ The `get_last_kv_window` function in NeuronxDistributedInference assumes K/V ten 4. Result: `torch.gather` attempts out-of-bounds access → DGE error **Why This Wasn't Caught Earlier:** + - Most models have `sliding_window >= max_position_embeddings` (no sliding window) - Or sliding_window < typical prompt lengths - GenericModel's unusual combination (sliding_window=4096, short prompts) exposed the bug @@ -252,15 +283,18 @@ The `get_last_kv_window` function in NeuronxDistributedInference assumes K/V ten ### Investigation Process 1. **Searched for scatter/gather operations:** + ```bash grep -rn "scatter\|gather" NeuronxDistributedInference/src/.../attention/attention_base.py ``` 2. **Located sliding window function calls:** + - `attention_context_encode_windowed_attention()` (line 1841) - `get_last_kv_window()` (line 2828) 3. **Analyzed `get_last_kv_window` logic:** + ```python # Line 635: Extract actual sequence length batch_size, num_head, actual_seq_len, head_dim = latest_k.shape @@ -299,6 +333,7 @@ def get_last_kv_window(window_size, position_ids, latest_k, latest_v, windowed_c ``` **Why This Works:** + - Pads K/V tensors to `window_size` before gathering - Padding with zeros doesn't affect attention output (will be masked) - After gathering, the KV cache has correct shape for token generation phase @@ -307,12 +342,14 @@ def get_last_kv_window(window_size, position_ids, latest_k, latest_v, windowed_c ### Key Learning **For models with sliding window attention:** + 1. **Test with various prompt lengths** including very short (1-10 tokens) 2. **Check assumptions in framework functions** - don't assume tensors are always full size 3. **Pad tensors defensively** when dealing with variable-length sequences 4. **Sliding window > sequence length is a valid edge case** that must be handled **Debugging scatter/gather errors:** + 1. Check tensor shapes at error site: `print(tensor.shape)` 2. Check gather indices range: `print(index.min(), index.max())` 3. Verify index.max() < tensor.size(gather_dim) @@ -323,6 +360,7 @@ def get_last_kv_window(window_size, position_ids, latest_k, latest_v, windowed_c ## Framework-Specific Learnings ### NeuronBaseModel Pattern + GenericModel followed the NeuronBaseModel pattern which requires: 1. **No custom forward() method** @@ -330,6 +368,7 @@ GenericModel followed the NeuronBaseModel pattern which requires: - Model must implement: `setup_attr_for_model()`, `init_model()`, `convert_hf_to_neuron_state_dict()` 2. **Attention class must inherit from NeuronAttentionBase** + ```python class NeuronGenericModelAttention(NeuronAttentionBase): def __init__(self, config): @@ -368,6 +407,7 @@ Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! ``` **What this means:** + - Framework replicates 2 KV heads → 24 KV heads during weight loading - Each query head gets its own KV head (standard MHA) - This is transparent to the model code @@ -384,6 +424,7 @@ GenericModel uses sliding window attention (4096 tokens). Framework handles this 3. **KV Cache Management:** `get_last_kv_window()` keeps only last window **Critical considerations:** + - Window size can exceed actual sequence length (our bug) - Position IDs wrap around: `position_ids % sliding_window` - KV cache is circular (not linear) @@ -393,6 +434,7 @@ GenericModel uses sliding window attention (4096 tokens). Framework handles this ## Compilation Details ### Compilation Parameters + ```python compile_neuron_model( model_class_path="modeling_genericmodel.NeuronGenericModelForCausalLM", @@ -408,16 +450,19 @@ compile_neuron_model( ``` ### Compilation Output + - **Context Encoding Model**: `context_encoding_model/_tp0_bk0/model.MODULE_*.neff` (~130s) - **Token Generation Model**: `token_generation_model/_tp0_bk0/model.MODULE_*.neff` (~100s) - **Total Time**: ~360 seconds (from scratch) - **Cache Hit**: ~130 seconds (if NEFFs cached) ### Compilation Warnings (Expected) + ``` WARNING:Neuron:TP degree (1) and KV heads (2) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! ``` + This is expected and safe - framework handles GQA → MHA conversion automatically. --- @@ -425,6 +470,7 @@ This is expected and safe - framework handles GQA → MHA conversion automatical ## Testing and Validation ### Test Setup + ```python # Test 1: Code generation prompt = "def fibonacci(n):" @@ -436,17 +482,21 @@ max_tokens = 30 ``` ### Results + **Test 1 - Code Generation:** + - Generated valid Python fibonacci implementation - Output: Complete function with while loop - Tokens/second: 15.0 **Test 2 - General Knowledge:** + - Correctly identified Paris as capital - Output: Multiple choice format (A. Paris, B. Rome, etc.) - Tokens/second: 14.9 ### Performance Metrics + - **Inference Time**: 2-3 seconds for 30-50 tokens - **Throughput**: 14.9-15.0 tokens/second - **Hardware**: trn1.32xlarge (32 Neuron cores) @@ -457,6 +507,7 @@ max_tokens = 30 ## Files Modified/Created ### Created Files + 1. `neuron_port/modeling_genericmodel.py` (489 lines) - Complete Neuron implementation - Classes: NeuronGenericModelAttention, NeuronGenericModelMLP, NeuronGenericModelDecoderLayer, NeuronGenericModelModel, NeuronGenericModelForCausalLM @@ -469,6 +520,7 @@ max_tokens = 30 - Inference test script using run_inference utility ### Modified Files + 1. `NeuronxDistributedInference/src/neuronx_distributed_inference/modules/attention/utils.py:640-646` - Added padding logic in `get_last_kv_window()` function - **This is a framework-level fix that benefits all models with sliding window attention** @@ -478,30 +530,35 @@ max_tokens = 30 ## Best Practices Established ### 1. Configuration Management + - Always verify weight tying by inspecting checkpoint files - Set `tie_word_embeddings` explicitly, don't rely on HF config - Add all framework-required attributes in `add_derived_config()` - Create model-specific NeuronConfig subclass ### 2. Attention Implementation + - Inherit from NeuronAttentionBase - Pass sliding_window parameter to base class - Let framework handle GQA sharding strategies - Don't implement custom attention mechanisms unless necessary ### 3. Compilation Verification + - Check `neuron_config.json` for correct attention class - Verify model compiles (exit code 0) - Check compiled NEFF hashes for cache hits - Test with various prompt lengths ### 4. Debugging Strategy + - Read framework error messages carefully (often give exact line numbers) - Check tensor shapes at error sites - Use grep to find relevant framework code - Test edge cases (very short/long prompts, batch size variations) ### 5. Framework Modifications + - Modify framework code only when necessary - Add defensive checks (if actual_seq_len < window_size) - Document why the change is needed @@ -512,6 +569,7 @@ max_tokens = 30 ## Architecture Patterns ### Successful Pattern: NeuronBaseModel + ```python class NeuronGenericModelForCausalLM(NeuronBaseModel): """ @@ -552,6 +610,7 @@ class NeuronGenericModelForCausalLM(NeuronBaseModel): ``` ### Attention Pattern + ```python class NeuronGenericModelAttention(NeuronAttentionBase): def __init__(self, config): @@ -578,6 +637,7 @@ class NeuronGenericModelAttention(NeuronAttentionBase): ``` ### MLP Pattern + ```python class NeuronGenericModelMLP(nn.Module): def __init__(self, config): @@ -619,22 +679,27 @@ class NeuronGenericModelMLP(nn.Module): ## Common Pitfalls ### 1. Assuming Tensors Are Always Full Size + **Problem:** Framework functions may assume tensors are at least a certain size **Solution:** Add defensive checks and padding when needed ### 2. Using Generic NeuronConfig + **Problem:** Compilation uses wrong attention class (defaults to LLaMA) **Solution:** Always create model-specific NeuronConfig subclass ### 3. Forgetting Weight Tying + **Problem:** Model fails to load due to missing lm_head.weight **Solution:** Check checkpoint files and set tie_word_embeddings explicitly ### 4. Missing Framework Attributes + **Problem:** AttributeError during model execution **Solution:** Add all required attributes in add_derived_config() ### 5. Not Testing Edge Cases + **Problem:** Model works for normal prompts but fails for very short/long ones **Solution:** Test with various lengths: 1, 10, 100, 512 tokens @@ -643,21 +708,25 @@ class NeuronGenericModelMLP(nn.Module): ## Performance Considerations ### Model Size vs. Performance + - **3B parameters**: 14.9-15.0 tokens/second (TP=1) - Larger models will benefit from higher TP degrees - GenericModel's 12:1 GQA ratio → MHA conversion adds memory overhead ### Compilation Time + - **First compile**: ~360 seconds (generates NEFFs) - **Cached compile**: ~130 seconds (reuses NEFFs) - NEFFs are deterministic based on config hash ### Memory Usage + - **Model weights**: ~6GB (3B params × 2 bytes/param for bfloat16) - **KV cache**: Depends on sliding_window (4096 tokens per layer) - **Activation memory**: Depends on batch_size and sequence_length ### Optimization Opportunities + 1. Increase TP degree for larger models (TP=2, 4, 8) 2. Use quantization (INT8) to reduce memory 3. Enable flash attention kernels (faster attention) @@ -668,6 +737,7 @@ class NeuronGenericModelMLP(nn.Module): ## Recommendations for Future Ports ### Pre-Port Checklist + 1. ✅ Identify architecture family (LLaMA-like, GPT-like, etc.) 2. ✅ Check for special features (sliding window, MoE, etc.) 3. ✅ Verify weight tying by inspecting checkpoint @@ -675,6 +745,7 @@ class NeuronGenericModelMLP(nn.Module): 5. ✅ Understand GQA/MHA/MQA configuration ### During Port + 1. ✅ Create model-specific NeuronConfig subclass first 2. ✅ Implement attention class inheriting from NeuronAttentionBase 3. ✅ Implement MLP with ColumnParallel/RowParallelLinear @@ -682,6 +753,7 @@ class NeuronGenericModelMLP(nn.Module): 5. ✅ Test compilation with small model first ### Post-Port Validation + 1. ✅ Verify correct attention class in neuron_config.json 2. ✅ Test with various prompt lengths (1, 10, 100, max_seq_len) 3. ✅ Validate output quality (perplexity, code generation, Q&A) @@ -707,32 +779,39 @@ The model is production-ready and generating coherent text at competitive speeds ## Appendix: Error Messages Reference ### Out-of-Bounds Error Pattern + ``` ERROR TDRV:exec_process_custom_notification: Received notification generated at runtime: failed to run scatter/gather (indirect memory copy via vector DGE), due to out-of-bound access. ``` + **Meaning:** Tensor index operation (gather/scatter) accessed beyond tensor bounds **Common causes:** Gather index range exceeds tensor size on gather dimension **Fix:** Add bounds checking or pad tensors to expected size ### Missing Weight Error Pattern + ``` RuntimeError: Missing weight tensor with key ``` + **Meaning:** State dict doesn't contain expected weight tensor **Common causes:** Weight tying not configured, incorrect weight name mapping **Fix:** Check tie_word_embeddings config, verify convert_hf_to_neuron_state_dict() ### Attribute Error Pattern + ``` AttributeError: '' object has no attribute '' ``` + **Meaning:** Config object missing required attribute **Common causes:** Framework expects standard HuggingFace attributes **Fix:** Add attributes in add_derived_config() method ### Wrong Attention Class Pattern + **Symptoms:** Model compiles but produces wrong outputs or runtime errors **Verification:** Check neuron_config.json for attn_cls value **Fix:** Create model-specific NeuronConfig with correct attn_cls diff --git a/skills/neuron-framework-autoport/references/knowledge_base/ROOT_CAUSE_REPEATED_OUTPUTS.md b/skills/neuron-framework-autoport/references/knowledge_base/ROOT_CAUSE_REPEATED_OUTPUTS.md index 7d1b9f1..b790341 100644 --- a/skills/neuron-framework-autoport/references/knowledge_base/ROOT_CAUSE_REPEATED_OUTPUTS.md +++ b/skills/neuron-framework-autoport/references/knowledge_base/ROOT_CAUSE_REPEATED_OUTPUTS.md @@ -3,10 +3,12 @@ ## Executive Summary **Problem**: Transformer model NeuronX ports can generate repetitive output patterns that fall into two distinct categories: + 1. **Complete Breakdown**: Repeated single tokens (e.g., "is is is is is...") 2. **Context Loss**: Repeated phrases or coherent but looping content (e.g., "in in in at at or in or in...") **Root Causes**: Two different fundamental bugs cause these patterns: + 1. **Tensor Shape Corruption**: Incorrect indexing of parallel layer outputs 2. **Missing Position Context**: Incorrect position embedding implementation @@ -19,12 +21,14 @@ ### Type 1: Complete Breakdown - Single Token Repetition **Symptoms:** + ``` Prompt: What is the capital of France? Output: is is is is is is is is is is is is is is is is is is is is is is is is is is is is is is ``` **Characteristics:** + - Repeats the **same single token** indefinitely - Usually high-frequency tokens like " is", " the", " and" - No coherent language structure @@ -35,12 +39,14 @@ Output: is is is is is is is is is is is is is is is is is is is is is is is is ### Type 2: Context Loss - Phrase/Pattern Repetition **Symptoms:** + ``` Prompt: What is the capital of France? Output: in in in in in in in in in in in at at or in or in or in or in or in or in or in or in or ``` **Characteristics:** + - Repeats **multiple tokens** in patterns - Maintains some grammatical structure - Shows understanding of language but loses context @@ -53,6 +59,7 @@ Output: in in in in in in in in in in in at at or in or in or in or in or in or ## Type 1: Tensor Shape Corruption Bug ### Location & Symptoms + - **File**: MLP forward method in transformer layers - **Symptom**: Single token repetition (e.g., " is is is...") - **Cause**: Incorrect `[0]` indexing on parallel layer outputs @@ -60,6 +67,7 @@ Output: in in in in in in in in in in in at at or in or in or in or in or in or ### The Bug #### What Was Written (WRONG) + ```python class NeuronTransformerMLP: def forward(self, hidden_states): @@ -76,6 +84,7 @@ class NeuronTransformerMLP: ``` #### What Should Be Written (CORRECT) + ```python class NeuronTransformerMLP: def forward(self, hidden_states): @@ -106,6 +115,7 @@ class ColumnParallelLinear: ``` **The Problem:** + - When `skip_bias_add=False` (default), parallel layers return **just a tensor** - `tensor[0]` slices the **first element along dimension 0** - This corrupts the batch dimension: `[batch_size, seq_len, hidden_size]` → `[seq_len, hidden_size]` @@ -113,6 +123,7 @@ class ColumnParallelLinear: - Model can only generate high-frequency tokens ### Cascade Effect + 1. **MLP layer**: Wrong tensor shapes due to `[0]` slicing 2. **Residual connections**: Shape mismatches cause broadcasting errors 3. **All transformer layers**: Receive increasingly corrupted hidden states @@ -124,6 +135,7 @@ class ColumnParallelLinear: ## Type 2: Position Embedding Bug ### Location & Symptoms + - **File**: Main model forward method - **Symptom**: Coherent but repetitive patterns (e.g., "word word word...") - **Cause**: Position embeddings not added to token embeddings @@ -131,55 +143,59 @@ class ColumnParallelLinear: ### The Bug #### What Was Written (WRONG) + ```python class NeuronTransformerModel: def forward(self, input_ids, attention_mask=None, **kwargs): batch_size, seq_len = input_ids.shape - + # Token embeddings inputs_embeds = self.embed_tokens(input_ids) - + # Position embeddings - WRONG IMPLEMENTATION if hasattr(self, 'position_embeddings'): position_ids = torch.arange(seq_len, device=input_ids.device) # ❌ Wrong shape! position_embeds = self.position_embeddings(position_ids) # ❌ Missing batch dimension! # ❌ BUG: Position embeddings not added to token embeddings! - + hidden_states = inputs_embeds # ❌ Missing position information! - + # Rest of forward pass... ``` #### What Should Be Written (CORRECT) + ```python class NeuronTransformerModel: def forward(self, input_ids, attention_mask=None, **kwargs): batch_size, seq_len = input_ids.shape - + # Token embeddings inputs_embeds = self.embed_tokens(input_ids) - + # Position embeddings - CORRECT IMPLEMENTATION if hasattr(self, 'position_embeddings'): position_ids = torch.arange(seq_len, device=input_ids.device).unsqueeze(0) # ✅ Correct shape position_ids = position_ids.expand(batch_size, -1) # ✅ Expand to [batch_size, seq_len] position_embeds = self.position_embeddings(position_ids) # ✅ Correct batch processing inputs_embeds = inputs_embeds + position_embeds # ✅ Add position info to tokens! - + hidden_states = inputs_embeds # ✅ Now includes position information! - + # Rest of forward pass... ``` ### Why This Bug Is Subtle But Critical **Understanding Position Embeddings:** + - Position embeddings provide **sequence position context** to transformers - They must be **added** to token embeddings: `token_embeds + position_embeds` - Without position context, models lose track of sequence structure - Results in coherent language that gets stuck in repetitive patterns ### Cascade Effect + 1. **Missing position info**: Token embeddings lack positional context 2. **Attention confusion**: Self-attention can't properly weight positions 3. **Pattern over-activation**: Model relies on learned patterns without position constraints @@ -193,38 +209,43 @@ class NeuronTransformerModel: ### Quick Diagnosis Checklist **If you see single token repetition (e.g., " is is is..."):** + - ✅ Check MLP forward methods for `[0]` indexing - ✅ Verify parallel layer return types - ✅ Look for tensor shape corruption **If you see coherent but repetitive patterns (e.g., "word word word..."):** + - ✅ Check position embedding implementation - ✅ Verify position embeddings are added to token embeddings - ✅ Ensure position IDs have correct batch dimensions ### Performance Comparison -| Bug Type | Inference Speed | Output Quality | Diagnostic Clue | -|----------|----------------|----------------|-----------------| -| **Tensor Shape Corruption** | Often faster (invalid computation) | **BROKEN - single token loops** | Same token repeated | -| **Position Embedding Missing** | Normal speed | **BROKEN - coherent loops** | Phrases/patterns repeated | -| **Fixed Implementation** | Normal speed | **WORKING - natural text** | Contextually appropriate | +| Bug Type | Inference Speed | Output Quality | Diagnostic Clue | +| ------------------------------ | ---------------------------------- | ------------------------------- | ------------------------- | +| **Tensor Shape Corruption** | Often faster (invalid computation) | **BROKEN - single token loops** | Same token repeated | +| **Position Embedding Missing** | Normal speed | **BROKEN - coherent loops** | Phrases/patterns repeated | +| **Fixed Implementation** | Normal speed | **WORKING - natural text** | Contextually appropriate | --- ## Common Patterns by Model Architecture ### GPT-Style Models (Absolute Position Embeddings) + - **Type 1 Bug**: Check `c_fc` and `c_proj` layers in MLP - **Type 2 Bug**: Check `wte + wpe` embedding addition - **Common tokens**: " is", " the", " and" for Type 1 ### BERT-Style Models (Absolute Position Embeddings) + - **Type 1 Bug**: Check intermediate and output layers in FFN - **Type 2 Bug**: Check token + position + segment embedding addition - **Common tokens**: "[CLS]", "[SEP]" for Type 1 ### T5-Style Models (Relative Position Embeddings) + - **Type 1 Bug**: Check dense layers in FFN - **Type 2 Bug**: Less common (uses relative positions) - **Common tokens**: "", "" for Type 1 @@ -234,18 +255,21 @@ class NeuronTransformerModel: ## Lessons Learned ### For Type 1 (Tensor Shape Corruption) + 1. **Always check return types** of framework layers 2. **Never assume tuple returns** - read the source code 3. **Test tensor shapes** at each layer during debugging 4. **Use assertions** to catch shape mismatches early ### For Type 2 (Position Embedding Missing) + 1. **Position embeddings are mandatory** in transformer models 2. **Always add position to token embeddings** - never skip this step 3. **Verify batch dimensions** in position ID creation 4. **Test generation quality** with diverse prompts ### Universal Debugging Tips + 1. **Strange outputs = fundamental bugs** - don't tweak hyperparameters first 2. **Compare with reference implementations** early and often 3. **Test inference immediately** after compilation @@ -256,6 +280,7 @@ class NeuronTransformerModel: ## Quick Fix Reference ### Type 1 Fix (Remove `[0]` indexing) + ```python # Before (BROKEN): hidden_states = self.linear_layer(hidden_states)[0] @@ -265,6 +290,7 @@ hidden_states = self.linear_layer(hidden_states) ``` ### Type 2 Fix (Add position embeddings) + ```python # Before (BROKEN): inputs_embeds = self.embed_tokens(input_ids) @@ -290,4 +316,4 @@ Both bugs are **completely fixable** with simple code changes, but require diffe **Universal Rule**: In transformer models, both proper tensor shapes AND position embeddings are essential for correct generation. Missing either component will cause repetitive output patterns that make the model unusable. -This applies to all transformer architectures including GPT, BERT, T5, LLaMA, and others when porting to NeuronX hardware. \ No newline at end of file +This applies to all transformer architectures including GPT, BERT, T5, LLaMA, and others when porting to NeuronX hardware. diff --git a/skills/neuron-framework-autoport/references/knowledge_base/TEXT_TO_VIDEO_MODEL_PORTING.md b/skills/neuron-framework-autoport/references/knowledge_base/TEXT_TO_VIDEO_MODEL_PORTING.md index 54afa96..e8f3297 100644 --- a/skills/neuron-framework-autoport/references/knowledge_base/TEXT_TO_VIDEO_MODEL_PORTING.md +++ b/skills/neuron-framework-autoport/references/knowledge_base/TEXT_TO_VIDEO_MODEL_PORTING.md @@ -27,11 +27,11 @@ A complete, self-contained guide for porting text-to-video diffusion models to A This guide has three phases. **Execute them in order. Do not skip ahead.** -| Phase | What | Why | -|-------|------|-----| -| **Discovery** (§3) | Run the original pipeline on CPU. Map every component, extract every constant and shape. | Constants you don't extract here will be guessed wrong later. | -| **Components** (§4) | Implement and validate each component independently on Neuron. | Each component must pass cosine similarity >0.998 vs CPU before moving on. | -| **Assembly** (§5) | Wire components into the final pipeline, matching the original's execution order exactly. | The assembly must reproduce which components run once vs per-step, what inputs they take, and in what order. | +| Phase | What | Why | +| ------------------- | ----------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------ | +| **Discovery** (§3) | Run the original pipeline on CPU. Map every component, extract every constant and shape. | Constants you don't extract here will be guessed wrong later. | +| **Components** (§4) | Implement and validate each component independently on Neuron. | Each component must pass cosine similarity >0.998 vs CPU before moving on. | +| **Assembly** (§5) | Wire components into the final pipeline, matching the original's execution order exactly. | The assembly must reproduce which components run once vs per-step, what inputs they take, and in what order. | **The most common failure mode is skipping Discovery.** An agent that jumps to implementation will guess constants (crop offsets, sequence lengths, conditioning timestep values) and get them wrong. These errors produce plausible but incorrect output with no error messages. @@ -75,17 +75,18 @@ Text Prompt **What varies across models:** -| Aspect | Examples | -|--------|----------| +| Aspect | Examples | +| ----------------------- | ------------------------------------------------------------------------------- | | Number of text encoders | 1 (e.g., T5-XXL), 2 (e.g., multimodal LLM + glyph encoder), 2 (e.g., T5 + CLIP) | -| Conditioning modules | Token refiner, none, pooled projection — varies by model | -| Token preprocessing | Reorder + zero-mask, simple concat, chat-template encoding — varies by model | -| Backbone architecture | Dual-stream MMDiT, single-stream DiT, cross-attention DiT | -| Attention type | Joint self-attention (MMDiT), cross-attention (DiT), both | -| CFG multiplier K | 2 (standard), 1 (no CFG / distilled) | -| VAE type | Non-causal 3D Conv VAE, causal 3D Conv VAE | +| Conditioning modules | Token refiner, none, pooled projection — varies by model | +| Token preprocessing | Reorder + zero-mask, simple concat, chat-template encoding — varies by model | +| Backbone architecture | Dual-stream MMDiT, single-stream DiT, cross-attention DiT | +| Attention type | Joint self-attention (MMDiT), cross-attention (DiT), both | +| CFG multiplier K | 2 (standard), 1 (no CFG / distilled) | +| VAE type | Non-causal 3D Conv VAE, causal 3D Conv VAE | **What is always the same:** + - Text encoders run once, backbone runs N×K times, VAE runs once - The backbone is the optimization target (dominates runtime) - The denoising loop uses a scheduler on CPU @@ -97,28 +98,29 @@ Text Prompt The framework depends on model size: -| Model Size | Framework | TP | When to use | -|---|---|---|---| -| >2B params | NxDI `NeuronApplicationBase` | TP≥2 | Backbone, large text encoders | -| 200M-2B params | Either NxDI or `torch_neuronx.trace` | Optional | Medium text encoders, large conditioning modules | -| <200M params | `torch_neuronx.trace` | No | Small encoders, conditioning modules, utility models | -| Conv3d-heavy (any size) | `torch_neuronx.trace` × N shards | No | VAE decoders (Conv3d exceeds instruction limits) | -| <1ms compute | CPU | N/A | Scheduler, simple pre/post-processing | +| Model Size | Framework | TP | When to use | +| ----------------------- | ------------------------------------ | -------- | ---------------------------------------------------- | +| >2B params | NxDI `NeuronApplicationBase` | TP≥2 | Backbone, large text encoders | +| 200M-2B params | Either NxDI or `torch_neuronx.trace` | Optional | Medium text encoders, large conditioning modules | +| <200M params | `torch_neuronx.trace` | No | Small encoders, conditioning modules, utility models | +| Conv3d-heavy (any size) | `torch_neuronx.trace` × N shards | No | VAE decoders (Conv3d exceeds instruction limits) | +| <1ms compute | CPU | N/A | Scheduler, simple pre/post-processing | ### HBM Budget Rule **Before choosing single-core vs TP, do this math:** + ``` neff_size ≈ (params × 2 bytes) + graph_overhead(0.5-2GB) + scratchpad(1GB) ``` On trn2 with LNC=2: **24GB per logical core**. If `neff_size > 22GB`, you MUST use TP≥2. Do not attempt single-core compilation — it will compile successfully but fail at load time with `Allocation Failure`, wasting the compilation time. -| Params (bf16) | Estimated NEFF | Fits single core? | -|---|---|---| -| <5B | <12GB | ✅ Yes | -| 5-10B | 12-22GB | ⚠️ Maybe — check carefully | -| >10B | >22GB | ❌ No — use TP≥2 | +| Params (bf16) | Estimated NEFF | Fits single core? | +| ------------- | -------------- | -------------------------- | +| <5B | <12GB | ✅ Yes | +| 5-10B | 12-22GB | ⚠️ Maybe — check carefully | +| >10B | >22GB | ❌ No — use TP≥2 | --- @@ -144,17 +146,18 @@ neuronx-cc --version 2>/dev/null Record these values: -| Constant | How to find it | Example (trn2.48xlarge) | -|----------|---------------|------------------------| -| `INSTANCE_TYPE` | Instance metadata or `ec2 describe-instances` | trn2.48xlarge | -| `TOTAL_NEURON_CORES` | `neuron-ls` | 64 (32 physical) | -| `HBM_PER_CORE` | Instance spec (24GB for trn2, 16GB for trn1) | 24GB | -| `LNC_DEFAULT` | `NEURON_RT_VIRTUAL_CORE_SIZE` env var (2 = LNC2, 1 = LNC1) | 2 (LNC2) | -| `EFFECTIVE_HBM` | HBM_PER_CORE / LNC_DEFAULT (with LNC2: 24GB shared by 2 logical cores) | 24GB per logical core pair | -| `INSTRUCTION_LIMIT` | ~5M per NEFF (hard compiler limit) | 5,000,000 | -| `SDK_VERSION` | `pip show torch-neuronx` | Check for known bugs/features | +| Constant | How to find it | Example (trn2.48xlarge) | +| -------------------- | ---------------------------------------------------------------------- | ----------------------------- | +| `INSTANCE_TYPE` | Instance metadata or `ec2 describe-instances` | trn2.48xlarge | +| `TOTAL_NEURON_CORES` | `neuron-ls` | 64 (32 physical) | +| `HBM_PER_CORE` | Instance spec (24GB for trn2, 16GB for trn1) | 24GB | +| `LNC_DEFAULT` | `NEURON_RT_VIRTUAL_CORE_SIZE` env var (2 = LNC2, 1 = LNC1) | 2 (LNC2) | +| `EFFECTIVE_HBM` | HBM_PER_CORE / LNC_DEFAULT (with LNC2: 24GB shared by 2 logical cores) | 24GB per logical core pair | +| `INSTRUCTION_LIMIT` | ~5M per NEFF (hard compiler limit) | 5,000,000 | +| `SDK_VERSION` | `pip show torch-neuronx` | Check for known bugs/features | **Why this matters:** + - A 4.7B-parameter text encoder needs ~9.4GB in bf16. At LNC=2 (24GB per core pair), it fits. At LNC=1 (24GB per single core), it also fits but uses a full physical core. - A 30-block backbone at seq_len=20,280 generates 37M instructions — 7× over the 5M limit. You MUST know this before attempting compilation. - A trn2.48xlarge has 64 cores — you can run backbone on cores 0-7 and VAE on cores 8-15 simultaneously. A trn2.3xlarge has only 16 cores — you may need to run components sequentially. @@ -175,20 +178,21 @@ Save the output frames — these are your ground truth for end-to-end validation Read the pipeline's `__call__` method source code. For every neural network invocation, record: -| Field | What to record | -|-------|---------------| -| **Name** | What the component is (e.g., "T5 text encoder", "token refiner") | -| **Call site** | Where in `__call__` it's invoked — before loop, inside loop, after loop | -| **Frequency** | Once, or per-step | -| **Inputs** | Tensor names, shapes, dtypes | -| **Outputs** | Tensor names, shapes, dtypes | +| Field | What to record | +| ---------------- | ------------------------------------------------------------------------- | +| **Name** | What the component is (e.g., "T5 text encoder", "token refiner") | +| **Call site** | Where in `__call__` it's invoked — before loop, inside loop, after loop | +| **Frequency** | Once, or per-step | +| **Inputs** | Tensor names, shapes, dtypes | +| **Outputs** | Tensor names, shapes, dtypes | | **Conditioning** | Does it take a timestep? If so, is it the loop timestep or a fixed value? | -| **Size** | Parameter count | -| **Dependencies** | Which other components' outputs does it consume? | +| **Size** | Parameter count | +| **Dependencies** | Which other components' outputs does it consume? | **This is the most important step in the entire guide.** The component graph IS the port specification. Every component you find here must be compiled for Neuron. Every constant you extract here prevents a wrong guess later. **Pay special attention to:** + - Components that take a timestep — is it the denoising loop timestep, or a fixed value (e.g., 1000.0)? - Components that run ONCE before the loop vs PER-STEP inside the loop - Any cropping, slicing, or reshaping of intermediate tensors (record exact offsets) @@ -205,12 +209,12 @@ For each text encoder, trace the tokenization and embedding extraction: Record: -| Constant | How to find it | -|----------|---------------| -| `SEQ_LEN` | `max_length` in the tokenizer call | -| `CROP_START` | Offset where pipeline slices hidden states (0 if no crop) | -| `CROP_LEN` | Number of tokens extracted | -| `EXTRACT_LAYER` | Which hidden state layer (e.g., -1, -2, -3) | +| Constant | How to find it | +| ---------------- | ----------------------------------------------------------------------- | +| `SEQ_LEN` | `max_length` in the tokenizer call | +| `CROP_START` | Offset where pipeline slices hidden states (0 if no crop) | +| `CROP_LEN` | Number of tokens extracted | +| `EXTRACT_LAYER` | Which hidden state layer (e.g., -1, -2, -3) | | `SYSTEM_MESSAGE` | Exact system prompt text, if any (affects tokenization and crop offset) | **Do NOT guess these values.** Read the pipeline source code. A wrong `CROP_START` silently shifts the token window and degrades output. @@ -220,15 +224,18 @@ Record: These traps were discovered during prior T2V ports. Each one silently produces wrong output. **Trap 1: System message whitespace.** Extract the exact system message using `repr()`: + ```python import inspect sig = inspect.signature(Pipeline._get_mllm_prompt_embeds) system_message = sig.parameters['system_message'].default print(repr(system_message)) # Shows \n, \t, multiple spaces ``` + A system message with `\n 1.` (newline + 8 spaces) vs `1.` (single space) changes tokenization by 5 tokens. This 40-character whitespace difference was the difference between generating a cat and generating a horse. **Trap 2: Tokenizer class mismatch.** The pipeline may use a fast tokenizer while `AutoTokenizer` loads the slow version. They can produce different token counts. Always verify: + ```python assert my_tokenizer_output['attention_mask'].sum() == pipe_tokenizer_output['attention_mask'].sum() ``` @@ -236,12 +243,15 @@ assert my_tokenizer_output['attention_mask'].sum() == pipe_tokenizer_output['att **Trap 3: `apply_chat_template` input format.** `format_text_input()` returns `[[{system}, {user}]]` (list of conversations). The pipeline passes `formatted[0]` (single conversation) or `formatted` (batch). Passing the wrong nesting level changes tokenization silently. **Trap 4: Decoder models used as encoders.** Some text encoders are decoder-architecture models with `is_decoder=True`. They use **causal** (lower-triangular) attention masking even when used for encoding. Check: + ```python print(model.config.is_decoder) # True = MUST use causal mask ``` + Using bidirectional masking on a causal model produces cosine ~0.60 vs the reference — close enough to look like a numerical issue, far enough to produce wrong images. **Trap 5: Hidden layer extraction index.** `output_hidden_states=True` returns `[embedding, layer_0, ..., layer_N]` — that's `N+1` entries. `hidden_states[-3]` is layer `N-2`, not layer `N-3`. Verify the index: + ```python out = model(input_ids=ids, attention_mask=mask, output_hidden_states=True) print(f"Total hidden states: {len(out.hidden_states)}") # N+1 @@ -252,12 +262,12 @@ print(f"hidden_states[-3] is layer index: {len(out.hidden_states) - 3}") For each conditioning module (refiner, projector, etc.): -| Constant | How to find it | -|----------|---------------| -| `RUNS_ONCE` | Is it called before the loop (True) or inside the loop (False)? | +| Constant | How to find it | +| ---------------- | ------------------------------------------------------------------------- | +| `RUNS_ONCE` | Is it called before the loop (True) or inside the loop (False)? | | `FIXED_TIMESTEP` | If it takes a timestep, is it fixed (e.g., 1000.0) or from the scheduler? | -| `INPUT_SHAPES` | Exact shapes of all inputs | -| `OUTPUT_SHAPES` | Exact shapes of all outputs | +| `INPUT_SHAPES` | Exact shapes of all inputs | +| `OUTPUT_SHAPES` | Exact shapes of all outputs | ### 3.5 Extract VAE Constants @@ -267,12 +277,12 @@ with open("path/to/vae/config.json") as f: cfg = json.load(f) ``` -| Constant | How to find it | -|----------|---------------| -| `SCALING_FACTOR` | `cfg["scaling_factor"]` | -| `SPATIAL_FACTOR` | `cfg["spatial_compression_ratio"]` or `cfg["ffactor_spatial"]` | +| Constant | How to find it | +| ----------------- | ---------------------------------------------------------------- | +| `SCALING_FACTOR` | `cfg["scaling_factor"]` | +| `SPATIAL_FACTOR` | `cfg["spatial_compression_ratio"]` or `cfg["ffactor_spatial"]` | | `TEMPORAL_FACTOR` | `cfg["temporal_compression_ratio"]` or `cfg["ffactor_temporal"]` | -| `LATENT_CHANNELS` | `cfg["latent_channels"]` | +| `LATENT_CHANNELS` | `cfg["latent_channels"]` | ### 3.6 Compute Latent Dimensions @@ -299,6 +309,7 @@ with torch.no_grad(): ### 3.8 Estimate HBM Memory Per Core Each NeuronCore has 24GB HBM (trn2). Estimate: + ``` weights_per_core = total_params × 2 bytes / tp_degree attention_per_core = heads_per_core × seq_len² × 2 bytes @@ -368,6 +379,7 @@ For each component in your component graph (§3.2), implement and validate it in **Applies to:** Large text encoders (e.g., T5-XXL, multimodal LLMs), transformer backbones. **Pattern:** + 1. Define a `nn.Module` subclass with TP layers (see §6) 2. Define an `InferenceConfig` subclass with model constants 3. Define a `ModelWrapper` subclass with `input_generator()` and `get_model_instance()` @@ -375,6 +387,7 @@ For each component in your component graph (§3.2), implement and validate it in 5. Compile, load, validate **TP layer replacement:** + - Attention Q/K/V: `ColumnParallelLinear(gather_output=False)` - Attention O: `RowParallelLinear(input_is_parallel=True)` - FFN up: `ColumnParallelLinear(gather_output=False)` → activation → FFN down: `RowParallelLinear(input_is_parallel=True)` @@ -383,6 +396,7 @@ For each component in your component graph (§3.2), implement and validate it in - QK norm (if present): `CustomRMSNorm` from NxDI (fused hardware call) **What goes inside the NEFF vs outside:** + - Inside: everything that runs on every forward call — embeddings, projections, attention, FFN, norms, output projection - Inside: time/timestep embedding (sinusoidal → linear → activation → linear) — this is cheap and avoids a CPU→device transfer per step - Inside: RoPE as `register_buffer` (pre-computed in `__init__`, NOT in `forward()`) @@ -404,6 +418,7 @@ def _init_rope(self, config): **Attention mask:** Use `-1e9` for masked positions, NOT `float('-inf')` (causes NaN in softmax). Pre-compute on CPU if the mask depends on variable-length inputs. **Weight conversion:** Write a `convert_hf_to_neuron_state_dict` that maps HF weight keys to your Neuron model's keys. Always verify parameter count matches: + ```python orig = sum(v.numel() for v in hf_sd.values() if key_is_relevant(v)) neuron = sum(v.numel() for v in neuron_sd.values()) @@ -419,13 +434,16 @@ assert orig == neuron, f"Weight count mismatch: {orig} vs {neuron}" **Critical differences from standard encoder implementation:** 1. **Causal attention mask required.** Pre-compute as `register_buffer`: + ```python causal = torch.tril(torch.ones(SEQ_LEN, SEQ_LEN)) self.register_buffer("causal_mask", causal) ``` + In forward: `mask = self.causal_mask[:S, :S] * attention_mask[:, None, None, :]` 2. **M-RoPE (multimodal RoPE).** Some multimodal models use section-based RoPE with `mrope_section` config. For text-only input, all position dimensions are identical, but the cos/sin interleaving must match exactly: + ```python mrope_section_2 = [s * 2 for s in config.rope_scaling['mrope_section']] cos_chunks = cos_full.split(mrope_section_2, dim=-1) @@ -433,12 +451,14 @@ assert orig == neuron, f"Weight count mismatch: {orig} vs {neuron}" ``` 3. **GQA (Grouped Query Attention).** With `num_kv_heads < num_attention_heads`, use `repeat_interleave` to expand KV heads before attention: + ```python k = k.repeat_interleave(num_heads // num_kv_heads, dim=1) v = v.repeat_interleave(num_heads // num_kv_heads, dim=1) ``` 4. **Hidden state extraction.** The pipeline extracts a specific intermediate layer, not the final output. Your model must return the correct layer: + ```python self.extract_idx = config.num_hidden_layers + EXTRACT_LAYER # e.g., 28 + (-3) = 25 for i, layer in enumerate(self.layers): @@ -457,6 +477,7 @@ assert orig == neuron, f"Weight count mismatch: {orig} vs {neuron}" **Applies to:** Small text encoders (e.g., CLIP, ByT5), token refiners, conditioning modules, utility models. **Pattern:** + 1. Wrap the model in a thin `nn.Module` with explicit `forward` signature 2. Trace with `torch_neuronx.trace` 3. Save with `torch.jit.save` @@ -487,6 +508,7 @@ torch.jit.save(traced, save_path) **⚠️ Causal vs Non-Causal:** If the VAE uses `CausalConv3d`, see §12.1 for the cache-aware tracing procedure. The pattern below applies to non-causal VAEs. For causal VAEs, each block needs multiple variants for different cache states. **Pattern:** + 1. Use the shapes from Discovery (§3.7) — no guessing 2. Split into individual blocks (conv_in, mid_resnets, mid_attn, up_resnets, up_upsamples, norm_conv_out) 3. Trace each block independently @@ -646,28 +668,29 @@ def convert_hf_to_neuron_state_dict(state_dict, config): ### Operations That DON'T Trace (inside NxDI NEFFs) -| Operation | Fix | -|---|---| -| `torch.arange` in forward() | Pre-compute as `register_buffer` | -| `torch.meshgrid` in forward() | Pre-compute as `register_buffer` | +| Operation | Fix | +| ------------------------------------ | -------------------------------------------------- | +| `torch.arange` in forward() | Pre-compute as `register_buffer` | +| `torch.meshgrid` in forward() | Pre-compute as `register_buffer` | | `torch.where` (dynamic output shape) | Use `torch.index_select` with pre-computed indices | -| `float('-inf')` in masks | Use `-1e9` | -| `F.pad(mode='replicate')` on 5D | Use `mode='constant'` | -| 8D tensor reshapes | Decompose into sequential ≤6D operations | -| `torch.sort` (trn2 only) | Use `torch.topk` | -| Unused input tensors | Remove — XLA ignores them, then crashes | +| `float('-inf')` in masks | Use `-1e9` | +| `F.pad(mode='replicate')` on 5D | Use `mode='constant'` | +| 8D tensor reshapes | Decompose into sequential ≤6D operations | +| `torch.sort` (trn2 only) | Use `torch.topk` | +| Unused input tensors | Remove — XLA ignores them, then crashes | ### Operations That Trace Fine -| Operation | Notes | -|---|---| -| `F.scaled_dot_product_attention` | Preferred for attention | -| `F.gelu(approximate="tanh")` / `nn.GELU(approximate="tanh")` | Works | -| `torch.cat`, `torch.index_select` | Works | -| `nn.Conv3d` | Works but generates large instruction graphs | -| `torch.topk` | Works on both trn1 and trn2 | +| Operation | Notes | +| ------------------------------------------------------------ | -------------------------------------------- | +| `F.scaled_dot_product_attention` | Preferred for attention | +| `F.gelu(approximate="tanh")` / `nn.GELU(approximate="tanh")` | Works | +| `torch.cat`, `torch.index_select` | Works | +| `nn.Conv3d` | Works but generates large instruction graphs | +| `torch.topk` | Works on both trn1 and trn2 | ### Input Constraints + - All inputs must be consumed by at least one operation - All input dtypes must match the compiled dtype (typically bfloat16, except indices which are long) - All input shapes are fixed at compile time @@ -677,7 +700,9 @@ def convert_hf_to_neuron_state_dict(state_dict, config): ## 8. Validation ### Per-Component Validation + After each component compiles, compare Neuron output vs CPU reference: + ```python cosine = F.cosine_similarity(cpu_out.flatten().float(), neuron_out.flatten().float(), dim=0) assert cosine > 0.998, f"Component failed: cosine={cosine}" @@ -686,7 +711,9 @@ assert cosine > 0.998, f"Component failed: cosine={cosine}" For VAE decoder (due to padding workaround): cosine >0.97 is acceptable. ### End-to-End Validation + Generate video with the same prompt used in Discovery (§3.1). Compare against saved CPU reference frames: + 1. Subject is present and recognizable 2. Frame std > 60 (not flat/gray) 3. Spatial structure is correct @@ -697,25 +724,27 @@ Generate video with the same prompt used in Discovery (§3.1). Compare against s **Numerical metrics can pass while the image is garbage.** A frame with std=108 and "correct" color statistics can be random colorful noise. A frame with cosine 0.99 against a reference can show a horse instead of a cat. **Rule:** After every pipeline change, visually inspect the output image. Compare side-by-side with the CPU reference. If you cannot view images directly, compute per-region statistics: + ```python arr = np.array(img) top, bot = arr[:H//2].mean(), arr[H//2:].mean() print(f"top={top:.0f} bot={bot:.0f} R={arr[:,:,0].mean():.0f} G={arr[:,:,1].mean():.0f} B={arr[:,:,2].mean():.0f}") ``` + A beach sunset should have: bright top (sky) > dark bottom (ground), R > G > B (warm tones). ### Failure Diagnosis -| Symptom | Likely Cause | -|---|---| -| Correct scene but missing/wrong subject | Text encoder `CROP_START` or `SEQ_LEN` is wrong | -| Flat gray output (std < 20) | Conditioning module called with wrong timestep, or broken CFG | -| NaN in output | `float('-inf')` in attention mask | -| Wrong output resolution | Wrong VAE spatial/temporal factor | -| Blurry/wrong spatial detail | Wrong RoPE convention or axis ordering | -| Garbled output, correct shapes | Token concatenation order wrong | -| Degraded quality, subtle | Dropped weight layer in conditioning module | -| Good quality but wrong content | Text encoder `SEQ_LEN` too short — prompt truncated | +| Symptom | Likely Cause | +| --------------------------------------- | ------------------------------------------------------------- | +| Correct scene but missing/wrong subject | Text encoder `CROP_START` or `SEQ_LEN` is wrong | +| Flat gray output (std < 20) | Conditioning module called with wrong timestep, or broken CFG | +| NaN in output | `float('-inf')` in attention mask | +| Wrong output resolution | Wrong VAE spatial/temporal factor | +| Blurry/wrong spatial detail | Wrong RoPE convention or axis ordering | +| Garbled output, correct shapes | Token concatenation order wrong | +| Degraded quality, subtle | Dropped weight layer in conditioning module | +| Good quality but wrong content | Text encoder `SEQ_LEN` too short — prompt truncated | --- @@ -724,48 +753,63 @@ A beach sunset should have: bright top (sky) > dark bottom (ground), R > G > B ( Each pitfall below was discovered from failed ports. They produce plausible but wrong output with no error message. ### P1: Conditioning module called per-step instead of once (or vice versa) + Read the original pipeline's `__call__` method. If a module is called before the loop, call it before the loop. If it takes a fixed timestep (not the scheduler's), use that fixed value. Getting this wrong changes the conditioning signal. ### P2: Wrong text encoder crop offset or sequence length + These constants must be extracted from the pipeline source code (§3.3), not guessed. A wrong crop offset shifts the entire token window. A too-short sequence length truncates the prompt. Both produce output that looks "almost right" but is wrong. ### P3: Skipping token preprocessing + If the original pipeline reorders, masks, or projects tokens before feeding them to the backbone, your port must do the same. Simply zeroing padding tokens is NOT equivalent to reordering them — padding tokens still occupy attention positions. ### P4: Dropping weight layers to "simplify" + If the original has `linear_1 → activation → linear_2`, your port must have both. Every `nn.Linear` in the original must exist in the port. Verify with parameter count matching. ### P5: Wrong RoPE convention + Multiple conventions exist (real vs complex, interleaved vs split-half, rotate_half vs unbind/stack). Verify your implementation matches the original by comparing Q/K values after RoPE application on CPU. ### P6: Building dynamic tensors inside the NEFF + `torch.arange`, `torch.meshgrid`, dynamic mask construction — all must be pre-computed as `register_buffer` or passed as inputs. They create CPU tensors during XLA tracing. ### P7: Using CPU fallbacks "temporarily" + If a component runs on CPU "to unblock testing," it will stay on CPU forever. Use the completion checklist (§10). ### P8: Inconsistent world_size across NxDI components + All NxDI components must share the same `world_size` (set to the maximum TP degree needed). ### P9: Wrong VAE scaling factor or spatial factor + Read from config, don't hardcode. A wrong spatial factor doubles or halves all dimensions. ### P10: Token concatenation order + If the backbone expects encoder tokens in a specific order (e.g., `[text_1, text_2, image]`), reversing or rearranging silently breaks attention patterns. ### 🆕 P11: System message whitespace + If the pipeline uses a system message for text encoding (e.g., chat-template-based encoders), the exact whitespace matters. A system message with `\n 1.` vs ` 1.` changes tokenization by multiple tokens, shifting the entire embedding window. Extract the system message with `repr()` to see hidden whitespace. ### 🆕 P12: Causal vs bidirectional attention in text encoders + Decoder-architecture models used as text encoders require causal (lower-triangular) attention masking. Using bidirectional masking produces cosine ~0.60 — close enough to look like a numerical issue, far enough to produce completely wrong conditioning. Always check `model.config.is_decoder`. ### 🆕 P13: Cosine similarity comparison target + When debugging a Neuron-compiled model, always compare the **same model** on Neuron vs CPU first (`cosine(neuron_out, cpu_out_same_model)`). Only then compare against the HF reference model. If the first is 1.0 and the second is <0.99, the bug is in your model architecture (attention mask, RoPE, layer extraction), not in Neuron or NxDI weight loading. ### 🆕 P14: Encoder output is all-zeros for certain prompts + Some text encoders (e.g., glyph-specific encoders) only activate for specific input patterns. For normal prompts, the output is all-zeros. This is correct behavior — do not "fix" it by running the encoder on the raw prompt. ### P15: Compiling partial models + Compiling a subset of layers (e.g., 8 of 54) to "test faster" hides critical issues: HBM limits that only appear at full model size, numerical accumulation errors, and semantically wrong output that passes numerical checks. Always compile the full component. --- @@ -775,12 +819,14 @@ Compiling a subset of layers (e.g., 8 of 54) to "test faster" hides critical iss **The port is not done until every box is checked.** This checklist is generated from YOUR component graph (§3.2), not from a fixed list. ### For EACH component in your component graph: + - [ ] Compiled for Neuron (NxDI or traced, per framework selection) - [ ] Validated: cosine >0.998 vs CPU reference (>0.97 for VAE) - [ ] Weight conversion verified: parameter count matches original - [ ] Runs on Neuron in the final pipeline (not a CPU fallback) ### Pipeline assembly: + - [ ] Execution order matches original pipeline's `__call__` method - [ ] Components that run once DO run once (not per-step) - [ ] Components that run per-step DO run per-step (not once) @@ -789,16 +835,17 @@ Compiling a subset of layers (e.g., 8 of 54) to "test faster" hides critical iss - [ ] Scheduler step uses float32 for numerical stability ### End-to-end: + - [ ] Generated video matches CPU reference (subject present, correct scene) - [ ] Frame std > 60 (not flat/gray) - [ ] No NaN in any intermediate tensor ### Artifacts: + - [ ] `pipeline_constants.py` — all constants from Discovery in one file - [ ] `compile_all.py` — single script that compiles every component - [ ] `run_inference.py` — single script that runs the full pipeline - --- ## 11. 🆕 Field-Tested Debugging Lessons @@ -834,12 +881,12 @@ The NxDI Flux example (`generate_flux.py`) demonstrates this pattern — it's a Dynamic tensor creation in `forward()` creates CPU tensors during XLA tracing. Pre-compute in `__init__` as `register_buffer`: -| Must pre-compute | Why | -|---|---| +| Must pre-compute | Why | +| -------------------------------------------- | -------------------------------------------- | | RoPE cos/sin (including M-RoPE interleaving) | `torch.arange` in forward creates CPU tensor | -| Causal attention mask | `torch.tril` in forward creates CPU tensor | -| Position indices | Same | -| Any grid or meshgrid | Same | +| Causal attention mask | `torch.tril` in forward creates CPU tensor | +| Position indices | Same | +| Any grid or meshgrid | Same | ### 11.5 VAE Conv3d Instruction Limits @@ -847,26 +894,25 @@ The VAE decoder (even at 1.26B params) generates 7.9M instructions from Conv3d o ### 11.6 Visual Validation Is the Only Ground Truth -| Metric | What it tells you | What it DOESN'T tell you | -|--------|-------------------|--------------------------| -| `std > 60` | Not flat/gray | Could be random colorful noise | -| `cosine > 0.99` | Numerically close | Could show a horse instead of a cat | -| `R > G > B` | Warm tones | Could be any warm-toned scene | -| **Visual inspection** | **Everything** | — | +| Metric | What it tells you | What it DOESN'T tell you | +| --------------------- | ----------------- | ----------------------------------- | +| `std > 60` | Not flat/gray | Could be random colorful noise | +| `cosine > 0.99` | Numerically close | Could show a horse instead of a cat | +| `R > G > B` | Warm tones | Could be any warm-toned scene | +| **Visual inspection** | **Everything** | — | Always look at the actual image. Compare side-by-side with the CPU reference. ### 11.7 Summary of Root Causes from Non-Causal VAE Port -| Symptom | Root Cause | How to Avoid | -|---------|-----------|--------------| -| Colorful noise output | Raw encoder output without pipeline preprocessing | Run Discovery phase (§3) completely | +| Symptom | Root Cause | How to Avoid | +| ------------------------------------------ | ------------------------------------------------- | --------------------------------------- | +| Colorful noise output | Raw encoder output without pipeline preprocessing | Run Discovery phase (§3) completely | | Wrong subject (e.g., horse instead of cat) | Bidirectional attention on a causal encoder model | Check `model.config.is_decoder` (§3.3a) | -| Slightly wrong subject | System message whitespace difference | Extract with `repr()` (§3.3a) | -| Cosine 0.60 "framework broken" | Comparing different models, not same model | Use comparison protocol (§11.2) | -| NEFF won't load (OOM) | Model too large for single core | Check HBM budget first (§2) | -| Cosine 0.59 "RoPE wrong" | Actually correct — wrong comparison target | Use comparison protocol (§11.2) | - +| Slightly wrong subject | System message whitespace difference | Extract with `repr()` (§3.3a) | +| Cosine 0.60 "framework broken" | Comparing different models, not same model | Use comparison protocol (§11.2) | +| NEFF won't load (OOM) | Model too large for single core | Check HBM budget first (§2) | +| Cosine 0.59 "RoPE wrong" | Actually correct — wrong comparison target | Use comparison protocol (§11.2) | --- @@ -878,14 +924,14 @@ These lessons come from porting a 1.3B-parameter T2V model with a causal 3D VAE The §4 Type C pattern (split into blocks, trace each, chain at runtime) was written for a **non-causal** VAE. Some models (e.g., those with `CausalConv3d`) use a **causal** VAE. The difference: -| | Non-Causal VAE | Causal VAE | -|---|---|---| -| Temporal padding | Symmetric (sees past + future) | Asymmetric (sees past only) | -| Processing | All frames at once | One frame at a time with cache | -| Cache | None | 33-entry stateful cache | -| Cache lifecycle | N/A | None → "Rep" string → 1-frame tensor → 2-frame tensor | -| Block variants needed | 1 per block | 3 per block (_0f, _1f, regular) + upsampler variants | -| Total traced blocks | ~18 | ~55 | +| | Non-Causal VAE | Causal VAE | +| --------------------- | ------------------------------ | ------------------------------------------------------ | +| Temporal padding | Symmetric (sees past + future) | Asymmetric (sees past only) | +| Processing | All frames at once | One frame at a time with cache | +| Cache | None | 33-entry stateful cache | +| Cache lifecycle | N/A | None → "Rep" string → 1-frame tensor → 2-frame tensor | +| Block variants needed | 1 per block | 3 per block (\_0f, \_1f, regular) + upsampler variants | +| Total traced blocks | ~18 | ~55 | **For causal VAEs, follow this procedure:** @@ -902,11 +948,13 @@ The §4 Type C pattern (split into blocks, trace each, chain at runtime) was wri If a `forward()` method takes a tensor input but only uses `tensor.shape` (not the values), XLA tracing drops the input entirely. The NEFF runs without error but produces wrong output. **Known affected patterns:** + - RoPE computed from `hidden_states.shape` (e.g., a `RotaryPosEmbed` module that only reads `.shape`) - Dynamic mask construction from `input.shape[1]` - Any positional encoding derived from input dimensions **Detection:** Look for this XLA warning during tracing: + ``` UserWarning: Received an input tensor that was unused or used in a non-static way ``` @@ -919,21 +967,21 @@ UserWarning: Received an input tensor that was unused or used in a non-static wa Two operations commonly found in causal 3D VAEs crash XLA: -| Operation | Error | Fix | -|---|---|---| -| `x[:, :, -N:, :, :]` where dim < N | "Value out of range" | `x[:, :, max(0, x.shape[2]-N):, :, :]` | -| `F.interpolate(mode='nearest-exact')` | "Unknown custom-call API version" | `mode='nearest'` | +| Operation | Error | Fix | +| ------------------------------------- | --------------------------------- | -------------------------------------- | +| `x[:, :, -N:, :, :]` where dim < N | "Value out of range" | `x[:, :, max(0, x.shape[2]-N):, :, :]` | +| `F.interpolate(mode='nearest-exact')` | "Unknown custom-call API version" | `mode='nearest'` | **Apply fixes via monkey-patching** (replace forward methods at runtime), NOT by modifying the installed diffusers source. This keeps the port self-contained. ### 12.4 NxDI ModelBuilder vs torch_neuronx.trace — Decision Framework -| Situation | Use | Why | -|---|---|---| -| Model fits on 1 core, no TP needed | `torch_neuronx.trace` | Simpler, no framework overhead | -| Model needs TP≥2 | NxDI ModelBuilder | Handles weight sharding, parallel layers, multi-core compilation | -| NxDI constant-folds a dynamic input | `torch_neuronx.trace` | NxDI's tracing can bake input values as constants | -| Need context parallelism | `torch_neuronx.trace` with NxD communication primitives | CP splits sequence, not model — use TP infrastructure for all-gathers | +| Situation | Use | Why | +| ----------------------------------- | ------------------------------------------------------- | --------------------------------------------------------------------- | +| Model fits on 1 core, no TP needed | `torch_neuronx.trace` | Simpler, no framework overhead | +| Model needs TP≥2 | NxDI ModelBuilder | Handles weight sharding, parallel layers, multi-core compilation | +| NxDI constant-folds a dynamic input | `torch_neuronx.trace` | NxDI's tracing can bake input values as constants | +| Need context parallelism | `torch_neuronx.trace` with NxD communication primitives | CP splits sequence, not model — use TP infrastructure for all-gathers | **Always try NxDI first for TP.** Only fall back to `torch_neuronx.trace` if NxDI constant-folds critical dynamic inputs (test by changing an input at inference and checking if output changes). @@ -942,21 +990,23 @@ Two operations commonly found in causal 3D VAEs crash XLA: The Neuron compiler has a hard 5M instruction limit per NEFF. This is the most common blocker for video models. **What generates instructions:** + - Attention: O(seq_len²) — dominates for long sequences - Conv3d: O(channels × spatial × temporal × kernel) — dominates for VAE - All-gathers: ~1K instructions each, adds up with many layers **Instruction count scaling (example: 30-block DiT backbone):** -| seq_len | TP=1 | CP=4 | CP=8 | -|---------|------|------|------| -| 6,240 (13 frames) | ~5M ✅ | N/A | N/A | -| 7,800 (17 frames) | 6.9M ❌ | N/A | N/A | -| 20,280 (49 frames) | 37M ❌ | 9.5M ❌ | 5.1M ❌* | +| seq_len | TP=1 | CP=4 | CP=8 | +| ------------------ | ------- | ------- | --------- | +| 6,240 (13 frames) | ~5M ✅ | N/A | N/A | +| 7,800 (17 frames) | 6.9M ❌ | N/A | N/A | +| 20,280 (49 frames) | 37M ❌ | 9.5M ❌ | 5.1M ❌\* | -*CP=8 at 5.1M is 2.5% over — solved by splitting into 2 NEFFs of 15 blocks each (~2.5M each). +\*CP=8 at 5.1M is 2.5% over — solved by splitting into 2 NEFFs of 15 blocks each (~2.5M each). **Solutions (in order of preference):** + 1. **Context parallelism** — splits sequence across ranks, reduces per-rank compute AND instructions 2. **Split into multiple NEFFs** — trace groups of blocks separately, chain at runtime (1-2ms overhead per NEFF boundary) 3. **Block-by-block tracing** — trace each block individually (higher overhead, last resort) @@ -1036,17 +1086,17 @@ A recurring mistake in this port was jumping into the first viable approach with ### 12.13 Summary of Causal-VAE Port Bugs -| Bug | Symptom | Detection | Fix | -|-----|---------|-----------|-----| -| XLA drops shape-only inputs (RoPE) | Cosine 0.994, muted output | XLA warning in trace log | Pre-compute on CPU, pass as input | -| NxDI constant-folds attention mask | T5 ignores padding | Same output for different masks | Use `torch_neuronx.trace` | -| Padding not zeroed after T5 | Colorful noise | Cosine 0.05 vs pipeline | Zero positions beyond seq_len | -| XLA negative indexing | Compile error | Error message | `max(0, x.shape[2] - N)` | -| XLA nearest-exact | Compile error | Error message | `mode='nearest'` | -| CausalConv3d segfaults XLA | Segfault | Crash | Block-by-block tracing | -| Stale VAE cache | Crosshatch on Neuron-decoded frames | Visual inspection | Return updated cache from blocks | -| Low frame count crosshatch | Crosshatch on ALL platforms | CPU fp32 reference | Generate ≥49 frames | -| 5M instruction limit | Compile error (exit 70) | Compiler message | CP + split NEFFs | -| LNC mismatch | Can't load T5 + backbone | Runtime error | Separate subprocesses | -| Old NEFFs cached in memory | Shape mismatch after retrace | Runtime error | Reload blocks after retracing | -| VAE block traced before patches | Wrong output (cos 0.004) | Per-block cosine check | Always apply patches first, trace ALL blocks fresh | +| Bug | Symptom | Detection | Fix | +| ---------------------------------- | ----------------------------------- | ------------------------------- | -------------------------------------------------- | +| XLA drops shape-only inputs (RoPE) | Cosine 0.994, muted output | XLA warning in trace log | Pre-compute on CPU, pass as input | +| NxDI constant-folds attention mask | T5 ignores padding | Same output for different masks | Use `torch_neuronx.trace` | +| Padding not zeroed after T5 | Colorful noise | Cosine 0.05 vs pipeline | Zero positions beyond seq_len | +| XLA negative indexing | Compile error | Error message | `max(0, x.shape[2] - N)` | +| XLA nearest-exact | Compile error | Error message | `mode='nearest'` | +| CausalConv3d segfaults XLA | Segfault | Crash | Block-by-block tracing | +| Stale VAE cache | Crosshatch on Neuron-decoded frames | Visual inspection | Return updated cache from blocks | +| Low frame count crosshatch | Crosshatch on ALL platforms | CPU fp32 reference | Generate ≥49 frames | +| 5M instruction limit | Compile error (exit 70) | Compiler message | CP + split NEFFs | +| LNC mismatch | Can't load T5 + backbone | Runtime error | Separate subprocesses | +| Old NEFFs cached in memory | Shape mismatch after retrace | Runtime error | Reload blocks after retracing | +| VAE block traced before patches | Wrong output (cos 0.004) | Per-block cosine check | Always apply patches first, trace ALL blocks fresh | diff --git a/skills/neuron-framework-autoport/references/knowledge_base/TRACE_PORT.md b/skills/neuron-framework-autoport/references/knowledge_base/TRACE_PORT.md index 207696a..dc143ed 100644 --- a/skills/neuron-framework-autoport/references/knowledge_base/TRACE_PORT.md +++ b/skills/neuron-framework-autoport/references/knowledge_base/TRACE_PORT.md @@ -14,12 +14,14 @@ Successfully ported GenericModel from CUDA to AWS Neuron hardware. The model com ## Session Context ### Starting Point + - Model implementation already existed (`modeling_genericmodel.py`, 622 lines) - Previous session had completed initial port but encountered runtime issues - Compiled artifacts existed but inference was failing - Test infrastructure partially set up ### Environment + - **Hardware:** trn1.32xlarge (32 NeuronCores, 16GB per core) - **Framework Versions:** - NeuronxDistributed @@ -28,6 +30,7 @@ Successfully ported GenericModel from CUDA to AWS Neuron hardware. The model com - Compiler: neuronx-cc ### Model Architecture + - **Type:** Decoder-only transformer for code generation - **Key Features:** - Grouped Query Attention (24 query heads, 2 KV heads) @@ -45,6 +48,7 @@ Successfully ported GenericModel from CUDA to AWS Neuron hardware. The model com **Timestamp:** Initial inference test **Error Message:** + ``` RuntimeError: Missing weight tensor with key lm_head.bias ``` @@ -53,6 +57,7 @@ RuntimeError: Missing weight tensor with key lm_head.bias The lm_head layer was initialized with `bias=config.use_bias` (True), but GenericModel uses weight tying where lm_head shares weights with embed_tokens and has no bias parameter. **Investigation Steps:** + 1. Examined checkpoint contents to verify no lm_head weights exist 2. Checked HuggingFace config.json - confirmed `use_bias: true` applies to most layers 3. Reviewed weight tying pattern from original HuggingFace implementation @@ -92,6 +97,7 @@ if "embed_tokens.weight" in neuron_state_dict and "lm_head.weight" not in neuron ``` **Outcome:** + - Cleared compiler caches (`/tmp/neuron-compile-cache`, Python `__pycache__`) - Recompiled model successfully - Token generation model: 123.65 seconds (PASS) @@ -103,12 +109,14 @@ if "embed_tokens.weight" in neuron_state_dict and "lm_head.weight" not in neuron **Timestamp:** First inference test after recompilation **Error Message:** + ``` AttributeError: 'GenericModelInferenceConfig' object has no attribute 'output_attentions' ``` **Root Cause:** The NeuronBaseModel framework's `_setup_func_config()` method (model_base.py:3407) expects config attributes: + - `output_attentions` - `output_hidden_states` - `use_return_dict` @@ -116,6 +124,7 @@ The NeuronBaseModel framework's `_setup_func_config()` method (model_base.py:340 These are standard HuggingFace config attributes used to control optional outputs during inference. **Investigation Steps:** + 1. Traced error to `model_base.py:3407`: `self.text_config.output_attentions` 2. Examined InferenceConfig base class - no default values provided 3. Checked original HuggingFace config.json - confirmed `use_cache: true` exists @@ -143,6 +152,7 @@ def add_derived_config(self): These are runtime configuration attributes, so no recompilation was needed. The fix could be applied and tested immediately. **Outcome:** + - Inference test passed immediately after config update - No recompilation required - Generated correct code output @@ -152,6 +162,7 @@ These are runtime configuration attributes, so no recompilation was needed. The ## Compilation Process ### Configuration Used + ```python NeuronConfig: - tp_degree: 1 @@ -163,6 +174,7 @@ NeuronConfig: ``` ### Compilation Timeline + 1. **HLO Generation:** 13.9 seconds - Context encoding model: 6.57 seconds - Token generation model: 6.91 seconds @@ -180,10 +192,12 @@ NeuronConfig: 4. **Total Compilation Time:** ~140 seconds ### Compiler Warnings (Expected/Ignorable) + ``` WARNING: TP degree (1) and KV heads (2) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! ``` + - Appears 60 times (30 layers × 2 models) - This is expected behavior for TP=1 with GQA - Framework automatically converts to MHA for single-device execution @@ -193,8 +207,10 @@ WARNING: TP degree (1) and KV heads (2) are not divisible. ## Inference Testing ### Test 1: Multiple Choice Question + **Prompt:** "What is the capital of France?" **Generated Output:** + ``` B @@ -210,11 +226,13 @@ Ber ``` **Analysis:** + - Model generated multiple-choice format (typical of code/test generation models) - Correctly included "Paris" as an option - Shows model is functioning but responds in structured format **Metrics:** + - Inference time: 0.62 seconds - Generated tokens: 20 - Throughput: 32.1 tokens/second @@ -222,8 +240,10 @@ Ber --- ### Test 2: Code Generation (Fibonacci) + **Prompt:** "def fibonacci(n):" **Generated Output:** + ```python if n <= 1: return n @@ -237,6 +257,7 @@ def fibonacci_ ``` **Analysis:** + - ✅ Correct recursive Fibonacci implementation - ✅ Proper base case handling - ✅ Syntactically valid Python @@ -244,6 +265,7 @@ def fibonacci_ - ✅ Multi-lingual capability (Chinese comment) **Metrics:** + - Inference time: 1.53 seconds - Generated tokens: 50 - Throughput: 32.6 tokens/second @@ -253,12 +275,15 @@ def fibonacci_ ## Key Learnings and Patterns ### 1. Weight Tying Considerations + **Pattern:** Models with weight tying require special handling: + - lm_head typically has `bias=False` even if other layers use bias - Weight copying logic needed in `convert_hf_to_neuron_state_dict()` - Check original HuggingFace implementation for tying behavior **Detection Method:** + ```python # Check if lm_head weights exist in checkpoint checkpoint_keys = state_dict.keys() @@ -269,9 +294,11 @@ lm_head_keys = [k for k in checkpoint_keys if 'lm_head' in k] --- ### 2. Config Attributes for Framework Compatibility + **Pattern:** NeuronxDistributed framework expects standard HuggingFace config attributes: **Required Attributes:** + ```python # In add_derived_config(): self.output_attentions = False # Control attention weights output @@ -280,6 +307,7 @@ self.use_return_dict = True # Use dictionary return format ``` **Best Practice:** + - Always implement `add_derived_config()` in custom InferenceConfig classes - Set sensible defaults for inference (False for optional outputs) - Use `hasattr()` checks to allow override from config.json @@ -287,7 +315,9 @@ self.use_return_dict = True # Use dictionary return format --- ### 3. Framework-Specific Requirements + **NeuronBaseModel Pattern:** + ```python # ❌ Don't override __init__() # ❌ Don't define custom forward() @@ -312,9 +342,11 @@ def init_model(self, config): --- ### 4. Sliding Window Attention Compatibility + **Issue:** Sliding window size (4096) > sequence length (128) causes errors **Solution:** Disable for initial port: + ```python super().__init__( config=config, @@ -324,13 +356,16 @@ super().__init__( ``` **Note for Production:** Re-enable sliding window with appropriate sequence length: + - Ensure `seq_len >= sliding_window` - Or implement dynamic handling in attention mechanism --- ### 5. Debugging Strategy + **When Compilation Fails:** + 1. Check compiler logs: `agent_artifacts/data/neff_output/*/log-neuron-cc.txt` 2. Clear all caches: ```bash @@ -340,6 +375,7 @@ super().__init__( 3. Review PYTHONPATH setup - ensure all framework paths included **When Inference Fails:** + 1. Check if it's a weight loading issue (missing tensors) 2. Verify config attributes are complete 3. Test with simple prompt first before complex scenarios @@ -349,12 +385,14 @@ super().__init__( ## Files Created/Modified ### Core Implementation + - ✅ `/home/ec2-user/agents/gaal/neuron_port/modeling_genericmodel.py` (622 lines) - Fixed lm_head bias configuration - Added weight tying support - Added missing config attributes ### Test Infrastructure + - ✅ `/home/ec2-user/agents/gaal/test_genericmodel_inference.py` (2.0KB) - Main test script using NeuroborosFoundations utilities - Follows phi3 test pattern @@ -366,11 +404,13 @@ super().__init__( - Complete usage documentation ### Compiled Artifacts + - ✅ `agent_artifacts/data/genericmodel_compiled/model.pt` (12MB) - ✅ `agent_artifacts/data/genericmodel_compiled/neuron_config.json` - ✅ `agent_artifacts/data/genericmodel_compiled/weights/` (model weights) ### Temporary Files (in agent_artifacts/tmp/) + - `compile_genericmodel.py` - Compilation wrapper - `test_genericmodel_simple.py` - Original test script - Various compilation/inference logs @@ -380,6 +420,7 @@ super().__init__( ## Environment Setup ### Required PYTHONPATH + ```bash export PYTHONPATH="/home/ec2-user/agents/gaal/neuron_port:\ /home/ec2-user/agents/gaal/NeuroborosFoundations/src:\ @@ -388,6 +429,7 @@ export PYTHONPATH="/home/ec2-user/agents/gaal/neuron_port:\ ``` ### Directory Structure + ``` /home/ec2-user/agents/gaal/ ├── neuron_port/ @@ -408,18 +450,21 @@ export PYTHONPATH="/home/ec2-user/agents/gaal/neuron_port:\ ## Performance Characteristics ### Model Size + - **Parameters:** ~1.3B (estimated from hidden_size=3072, 30 layers) - **Original checkpoint:** 12GB (safetensors) - **Compiled model:** 12MB (NEFF + metadata) - **Weights:** Stored separately, sharded ### Inference Performance + - **Throughput:** 32.6 tokens/second - **Latency:** ~30ms per token - **Configuration:** TP=1, batch=1, seq_len=128, bfloat16 - **Memory:** Single NeuronCore utilized ### Compilation Performance + - **Total time:** ~2.5 minutes - **Caching:** Token generation NEFF reused for context encoding - **Optimization levels:** -O2 (token gen), -O1 (context enc) @@ -442,6 +487,7 @@ export PYTHONPATH="/home/ec2-user/agents/gaal/neuron_port:\ ## Recommendations for Future Ports ### Pre-Implementation Checklist + 1. ☑️ Review HuggingFace config.json for architecture details 2. ☑️ Check for weight tying (missing lm_head in checkpoint) 3. ☑️ Identify attention mechanism (MHA, MQA, GQA) @@ -450,6 +496,7 @@ export PYTHONPATH="/home/ec2-user/agents/gaal/neuron_port:\ 6. ☑️ Check for sliding window or other special features ### Implementation Pattern + 1. Create InferenceConfig with all required attributes 2. Implement attention class inheriting from NeuronAttentionBase 3. Implement MLP class with appropriate parallelism @@ -459,6 +506,7 @@ export PYTHONPATH="/home/ec2-user/agents/gaal/neuron_port:\ 7. Implement convert_hf_to_neuron_state_dict() for weight mapping ### Testing Strategy + 1. Start with small seq_len (128) for faster iteration 2. Test compilation first (catch architecture issues early) 3. Test weight loading (catch mapping issues) @@ -467,6 +515,7 @@ export PYTHONPATH="/home/ec2-user/agents/gaal/neuron_port:\ 6. Increase seq_len/batch_size as needed ### Common Pitfalls to Avoid + - ❌ Putting lm_head in ForCausalLM instead of base Model - ❌ Using config.use_bias for lm_head when weight tying exists - ❌ Forgetting to add output_attentions/output_hidden_states @@ -478,6 +527,7 @@ export PYTHONPATH="/home/ec2-user/agents/gaal/neuron_port:\ ## Appendix: Command Reference ### Compilation + ```bash cd /home/ec2-user/agents/gaal/agent_artifacts/tmp export PYTHONPATH="..." @@ -490,12 +540,14 @@ python3 compile_genericmodel.py \ ``` ### Inference Testing + ```bash cd /home/ec2-user/agents/gaal ./run_genericmodel_test.sh ``` ### Cache Clearing (if needed) + ```bash rm -rf /tmp/neuron-compile-cache find . -type d -name __pycache__ -exec rm -rf {} + @@ -503,6 +555,7 @@ rm -rf agent_artifacts/data/genericmodel_compiled/* ``` ### Checking Hardware + ```bash neuron-ls # Verify NeuronCore availability ``` @@ -517,6 +570,7 @@ Successfully ported GenericModel to AWS Neuron with working inference at 32.6 to 2. **Missing config attributes** - Resolved by adding framework-expected attributes in add_derived_config() The port is production-ready for single-core inference with seq_len=128. Future work could include: + - Increasing sequence length (512, 2048, 4096) - Enabling tensor parallelism (TP > 1) - Re-enabling sliding window attention diff --git a/skills/neuron-framework-autoport/references/knowledge_base/TROUBLESHOOTING.md b/skills/neuron-framework-autoport/references/knowledge_base/TROUBLESHOOTING.md index 2538e24..c07a3ea 100644 --- a/skills/neuron-framework-autoport/references/knowledge_base/TROUBLESHOOTING.md +++ b/skills/neuron-framework-autoport/references/knowledge_base/TROUBLESHOOTING.md @@ -7,14 +7,16 @@ This document provides a comprehensive breakdown of all errors encountered durin ### 1. Missing Model Path in Base Class **Error Message:** + ``` TypeError: NeuronBaseModel.__init__() missing 1 required positional argument: 'model_path' ``` -**Root Cause:** +**Root Cause:** The `NeuronLlama3Model` class wasn't properly calling the parent class constructor with required arguments. **Solution:** + ```python def __init__(self, config): # Properly initialize parent class with model_path @@ -28,14 +30,15 @@ def __init__(self, config): ### 2. Incorrect Compilation Method Name **Error Message:** + ``` AttributeError: 'NeuronLlama3ForCausalLM' object has no attribute 'compile_model' ``` -**Root Cause:** +**Root Cause:** Used wrong method name for model compilation. -**Solution:** +**Solution:** Changed from `model.compile_model()` to `model.compile()` **Key Insight:** NeuronX uses `compile()` method, not `compile_model()`. @@ -44,29 +47,30 @@ Changed from `model.compile_model()` to `model.compile()` ### 3. Missing Checkpoint Files in Compiled Directory -**Error:** +**Error:** Model compilation succeeded but inference failed due to missing weight files. -**Root Cause:** +**Root Cause:** The compilation process didn't copy necessary checkpoint files to the output directory. **Solution:** + ```python def copy_necessary_files(checkpoint_path, output_dir): """Copy necessary files for inference""" print("Copying necessary files for inference...") - + # Copy main checkpoint if os.path.exists(checkpoint_path): shutil.copy2(checkpoint_path, os.path.join(output_dir, "pytorch_model.bin")) print("Copied checkpoint as pytorch_model.bin") - + # Copy safetensors files safetensors_files = glob.glob(os.path.join(os.path.dirname(checkpoint_path), "*.safetensors")) for file in safetensors_files: shutil.copy2(file, output_dir) print(f"Copied {os.path.basename(file)}") - + # Copy tokenizer files tokenizer_files = ["tokenizer.model", "tokenizer_config.json"] for file in tokenizer_files: @@ -85,28 +89,32 @@ def copy_necessary_files(checkpoint_path, output_dir): ### 4. Critical Intermediate Size Calculation Error **Error Message:** + ``` RuntimeError: expected shape torch.Size([5632, 2048]) for layers.0.mlp.gate_proj.weight but found torch.Size([8192, 2048]) ``` -**Root Cause:** +**Root Cause:** Incorrect calculation of `intermediate_size` in the Llama3 MLP layers. The original formula was wrong: -- **Wrong:** `hidden_dim = int(2 * self.hidden_size / 3)` + +- **Wrong:** `hidden_dim = int(2 * self.hidden_size / 3)` - **Correct:** `hidden_dim = 4 * dim; hidden_dim = int(2 * hidden_dim / 3)` **Debugging Process:** + 1. Checked actual weight shapes in checkpoint: `8192` 2. Verified our calculation produced: `5632` 3. Traced back to original Llama3 implementation 4. Found the correct formula in the original codebase **Solution:** + ```python def calculate_intermediate_size(params): """ Calculate intermediate_size from ffn_dim_multiplier like original Llama3 Based on FeedForward.__init__ in original Llama3 implementation - + Original logic: hidden_dim = 4 * dim # Start with 4x the hidden dimension hidden_dim = int(2 * hidden_dim / 3) # Apply 2/3 factor @@ -117,22 +125,23 @@ def calculate_intermediate_size(params): dim = params['dim'] multiple_of = params.get('multiple_of', 256) ffn_dim_multiplier = params.get('ffn_dim_multiplier') - + # Base calculation: 4 * dim, then 2/3 of that hidden_dim = 4 * dim hidden_dim = int(2 * hidden_dim / 3) - + # Apply multiplier if specified if ffn_dim_multiplier is not None: hidden_dim = int(ffn_dim_multiplier * hidden_dim) - + # Round to nearest multiple hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of) - + return hidden_dim ``` **Verification:** + ```python # For Llama3.2-1B with dim=2048, ffn_dim_multiplier=1.5, multiple_of=256 # hidden_dim = 4 * 2048 = 8192 @@ -150,14 +159,16 @@ def calculate_intermediate_size(params): ### 5. Missing Configuration Attributes **Error Message:** + ``` AttributeError: 'Llama3InferenceConfig' object has no attribute 'output_attentions' ``` -**Root Cause:** +**Root Cause:** The NeuronX framework expected certain configuration attributes that weren't defined in our custom config class. **Solution:** + ```python config_dict = { 'hidden_size': params['dim'], @@ -187,19 +198,21 @@ config_dict = { ### 6. Configuration Not Persisting After Compilation -**Error:** +**Error:** Even after fixing the code, the compiled model still used old configuration values. -**Root Cause:** +**Root Cause:** The model was loading configuration from the saved `neuron_config.json` file rather than using the updated code. **Debugging Process:** + 1. Fixed code but error persisted 2. Checked compiled `neuron_config.json` file 3. Found old values still present 4. Realized compilation caches configuration **Solution:** + ```bash # Always delete compiled model after code changes rm -rf llama3_compiled @@ -220,20 +233,23 @@ grep "output_attentions" llama3_compiled/neuron_config.json ### 7. Neuron Runtime Initialization Failure **Error Message:** + ``` -RuntimeError: The PyTorch Neuron Runtime could not be initialized. +RuntimeError: The PyTorch Neuron Runtime could not be initialized. Logical Neuron Core(s) not available - Requested:32 Available:0 ``` -**Root Cause:** +**Root Cause:** Temporary unavailability of Trainium accelerator cores. **Debugging Process:** + 1. Initially thought it was a code issue 2. Checked hardware status 3. Confirmed it was a resource availability issue **Solution:** + - **Wait for hardware availability** (this was a temporary resource issue) - **Retry the inference** once cores became available - No code changes needed - this was an infrastructure issue @@ -246,13 +262,14 @@ Temporary unavailability of Trainium accelerator cores. ### 8. Incorrect Development Workflow -**Error:** +**Error:** Attempting to run inference before proper compilation, leading to various cascading errors. -**Root Cause:** +**Root Cause:** Not following the proper NeuronX development workflow. **Solution - Correct Workflow:** + ```bash # 1. Always compile first python compile_llama3.py @@ -272,41 +289,42 @@ python run_inference.py ### 9. Weight Format Conversion Issues -**Error:** +**Error:** Various shape mismatches during weight loading. -**Root Cause:** +**Root Cause:** Inconsistent handling of original Llama3 weight format vs. NeuronX expected format. **Solution:** + ```python def convert_state_dict_to_neuronx_format(original_state_dict, config): """Convert original Llama3 weights to NeuronX format""" converted_state_dict = {} - + # Handle embedding layers if 'tok_embeddings.weight' in original_state_dict: converted_state_dict['embed_tokens.weight'] = original_state_dict['tok_embeddings.weight'] - + # Handle output layer if 'output.weight' in original_state_dict: converted_state_dict['lm_head.weight'] = original_state_dict['output.weight'] - + # Handle transformer layers for layer_idx in range(config.num_hidden_layers): layer_prefix = f'layers.{layer_idx}' - + # Attention weights if f'{layer_prefix}.attention.wq.weight' in original_state_dict: converted_state_dict[f'{layer_prefix}.self_attn.q_proj.weight'] = \ original_state_dict[f'{layer_prefix}.attention.wq.weight'] - + if f'{layer_prefix}.attention.wk.weight' in original_state_dict: converted_state_dict[f'{layer_prefix}.self_attn.k_proj.weight'] = \ original_state_dict[f'{layer_prefix}.attention.wk.weight'] - + # ... continue for all weight mappings - + return converted_state_dict ``` @@ -334,6 +352,7 @@ def convert_state_dict_to_neuronx_format(original_state_dict, config): ### Development Best Practices Learned: #### ✅ **Do's:** + - Always compile before testing inference - Verify mathematical formulas against original implementations - Check saved configuration files, not just source code @@ -342,6 +361,7 @@ def convert_state_dict_to_neuronx_format(original_state_dict, config): - Delete compiled models after significant code changes #### ❌ **Don'ts:** + - Don't assume parent class constructors work the same way - Don't skip weight format conversion - Don't ignore hardware resource availability @@ -356,6 +376,7 @@ def convert_state_dict_to_neuronx_format(original_state_dict, config): 5. **Test incrementally** - compile and test after each major change ### Total Development Time: + - **Initial Implementation:** ~2 hours - **Error Resolution:** ~1.5 hours - **Testing & Validation:** ~30 minutes @@ -366,6 +387,7 @@ def convert_state_dict_to_neuronx_format(original_state_dict, config): ## 🎯 Success Metrics **Final Results:** + - ✅ Model compiles successfully - ✅ All weight shapes match perfectly - ✅ Configuration includes all required attributes @@ -374,9 +396,10 @@ def convert_state_dict_to_neuronx_format(original_state_dict, config): - ✅ Complete end-to-end functionality achieved **Generated Output Example:** + ``` Input: "Hello, how are you?" Output: "Hello, how are you? I am I am I am I am I am" ``` -While the output shows some repetition (normal for a small 1B model), the core functionality works perfectly, demonstrating successful model porting from CUDA to NeuronX. \ No newline at end of file +While the output shows some repetition (normal for a small 1B model), the core functionality works perfectly, demonstrating successful model porting from CUDA to NeuronX. diff --git a/skills/neuron-framework-autoport/references/knowledge_base/WEIGHT_SHARDING_FIXES_SUMMARY.md b/skills/neuron-framework-autoport/references/knowledge_base/WEIGHT_SHARDING_FIXES_SUMMARY.md index 3dbd8c0..209f59d 100644 --- a/skills/neuron-framework-autoport/references/knowledge_base/WEIGHT_SHARDING_FIXES_SUMMARY.md +++ b/skills/neuron-framework-autoport/references/knowledge_base/WEIGHT_SHARDING_FIXES_SUMMARY.md @@ -18,6 +18,7 @@ model.compile(output_path) # This triggers automatic sharding ``` **Key Changes**: + - Removed manual model loading and weight transfer - Removed manual state dict saving - Added proper `model.compile()` call that triggers the framework's compilation and sharding process @@ -37,6 +38,7 @@ model.compile(output_path) # This triggers automatic sharding 4. **`from_pretrained()`**: Loads compiled models from directories #### Removed Duplicates: + - Removed duplicate `NeuronGptOssForCausalLM` class definitions - Removed duplicate `NeuronGptOssModel` class definitions - Moved `lm_head` to the model class where it belongs @@ -50,7 +52,7 @@ def convert_hf_to_neuron_state_dict(state_dict: dict, config: InferenceConfig) - # Handles complex MoE parameter structure # Splits combined QKV projections # Adds rank utilities for tensor parallel support - + def checkpoint_loader_fn(self, mmap: bool = False): """Redirect to original checkpoint directory for weight loading.""" # Similar to phi3 implementation @@ -97,14 +99,14 @@ gpt_oss_compiled_neuron/ ### 5. Key Differences: Before vs After -| Aspect | Before (Broken) | After (Fixed) | -|--------|----------------|---------------| -| Compilation | Manual `torch.save()` | Framework `model.compile()` | -| Sharding | Not performed | Automatic during compilation | -| Checkpoint Loading | Basic loading | Proper redirection and conversion | -| Framework Integration | Bypassed | Full integration | -| Parameter Conversion | Missing | Complete HF→NeuronX mapping | -| Weight Files | Single `pytorch_model.bin` | Multiple sharded `.safetensors` | +| Aspect | Before (Broken) | After (Fixed) | +| --------------------- | -------------------------- | --------------------------------- | +| Compilation | Manual `torch.save()` | Framework `model.compile()` | +| Sharding | Not performed | Automatic during compilation | +| Checkpoint Loading | Basic loading | Proper redirection and conversion | +| Framework Integration | Bypassed | Full integration | +| Parameter Conversion | Missing | Complete HF→NeuronX mapping | +| Weight Files | Single `pytorch_model.bin` | Multiple sharded `.safetensors` | ### 6. Memory Issue Resolution Options @@ -133,4 +135,4 @@ cat neuronx_gpt_oss/gpt_oss_compiled_neuron/neuron_config.json | grep save_shard The GPT-OSS implementation now properly integrates with the NeuronX distributed inference framework and will perform automatic weight sharding during compilation, matching the working phi3 implementation pattern. -The memory issue encountered is a separate concern related to model size and hardware constraints, not the sharding implementation itself. \ No newline at end of file +The memory issue encountered is a separate concern related to model size and hardware constraints, not the sharding implementation itself. diff --git a/skills/neuron-framework-autoport/references/knowledge_base/compilation_errors_and_fixes.md b/skills/neuron-framework-autoport/references/knowledge_base/compilation_errors_and_fixes.md index 2287537..1fe556c 100644 --- a/skills/neuron-framework-autoport/references/knowledge_base/compilation_errors_and_fixes.md +++ b/skills/neuron-framework-autoport/references/knowledge_base/compilation_errors_and_fixes.md @@ -3,6 +3,7 @@ This document provides a comprehensive guide to all compilation errors encountered and their solutions when implementing Llama3 for the NeuronxDistributed framework. ## Table of Contents + 1. [Base Class Integration Errors](#1-base-class-integration-errors) 2. [Constructor Signature Issues](#2-constructor-signature-issues) 3. [Configuration Loading Problems](#3-configuration-loading-problems) @@ -18,6 +19,7 @@ This document provides a comprehensive guide to all compilation errors encounter ## 1. Base Class Integration Errors ### Error: Missing Required Methods + ``` AttributeError: 'NeuronLlama3Model' object has no attribute 'setup_attr_for_model' ``` @@ -51,6 +53,7 @@ class NeuronLlama3Model(NeuronBaseModel): ## 2. Constructor Signature Issues ### Error: Incompatible Constructor Parameters + ``` TypeError: __init__() got unexpected keyword arguments ``` @@ -63,7 +66,7 @@ TypeError: __init__() got unexpected keyword arguments def __init__(self, model_path: str = None, config: Llama3InferenceConfig = None, **kwargs): if model_path is None: model_path = "" # Provide empty string as default - + if config is not None: super().__init__(model_path, config=config, **kwargs) else: @@ -77,6 +80,7 @@ def __init__(self, model_path: str = None, config: Llama3InferenceConfig = None, ## 3. Configuration Loading Problems ### Error: Multiple Configuration Format Support + ``` FileNotFoundError: No configuration file found in /path/to/model ``` @@ -90,7 +94,7 @@ FileNotFoundError: No configuration file found in /path/to/model def from_pretrained(cls, model_path): params_file = os.path.join(model_path, "params.json") config_file = os.path.join(model_path, "config.json") - + if os.path.exists(params_file): # Load original Llama3 format with open(params_file, 'r') as f: @@ -112,6 +116,7 @@ def from_pretrained(cls, model_path): ## 4. Parameter Name Mapping ### Error: Inconsistent Parameter Names + ``` KeyError: 'hidden_size' not found in configuration ``` @@ -145,6 +150,7 @@ def from_original_params(cls, params): ## 5. Checkpoint Conversion Issues ### Error: Parameter Count Mismatch + ``` Original checkpoint: 147 parameters Converted checkpoint: 164 parameters @@ -153,8 +159,9 @@ Converted checkpoint: 164 parameters **Root Cause**: Framework adds metadata and rank utilities for distributed training. **Solution**: This is expected behavior. The increase is due to: + - Framework splitting combined weight matrices -- Adding metadata and configuration parameters +- Adding metadata and configuration parameters - Adding rank utilities for tensor parallel support - Tensor reshaping for distributed training compatibility @@ -162,20 +169,20 @@ Converted checkpoint: 164 parameters @staticmethod def convert_hf_to_neuron_state_dict(state_dict: dict, config: InferenceConfig) -> dict: neuron_config = config.neuron_config - + # Add rank utilities for tensor parallel support if neuron_config.vocab_parallel: state_dict["embed_tokens.rank_util.rank"] = torch.arange( 0, neuron_config.local_ranks_size ) - + num_layers = config.num_hidden_layers tp_degree = neuron_config.tp_degree for i in range(num_layers): state_dict[f"layers.{i}.self_attn.rank_util.rank"] = torch.arange( 0, tp_degree, dtype=torch.int32 ) - + return state_dict ``` @@ -186,6 +193,7 @@ def convert_hf_to_neuron_state_dict(state_dict: dict, config: InferenceConfig) - ## 6. Forward Method Signature Conflicts ### Error: Framework Forward Method Conflict + ``` ERROR: You cannot specify both input_ids and inputs_embeds at the same time ``` @@ -206,11 +214,11 @@ class NeuronLlama3Model(NeuronBaseModel): def setup_attr_for_model(self, config): # Setup attributes pass - + def init_model(self, config): # Initialize components pass - + # No forward method - base class handles this ✅ ``` @@ -221,6 +229,7 @@ class NeuronLlama3Model(NeuronBaseModel): ## 7. Layer Return Format Mismatches ### Error: Tuple Unpacking Mismatch + ``` ERROR: too many values to unpack (expected 3) ``` @@ -255,6 +264,7 @@ return outputs ## 8. Import and Module Structure ### Error: Module Import Problems + ``` ImportError: cannot import name 'NeuronLlama3Model' from 'neuronx_llama3' ``` @@ -273,7 +283,7 @@ from .modeling_llama3 import ( __all__ = [ "NeuronLlama3Model", - "Llama3InferenceConfig", + "Llama3InferenceConfig", "NeuronLlama3ForCausalLM" ] ``` @@ -285,6 +295,7 @@ __all__ = [ ## 9. Framework Pattern Compliance ### Error: Not Following Framework Patterns + Multiple compilation issues due to not following the established framework patterns. **Solution**: Study existing models (Qwen3, Mistral) and follow their patterns: @@ -308,7 +319,7 @@ class NeuronLlama3Model(NeuronBaseModel): # Framework Pattern for CausalLM Classes class NeuronLlama3ForCausalLM(NeuronBaseForCausalLM): _model_cls = NeuronLlama3Model - + @staticmethod def convert_hf_to_neuron_state_dict(state_dict, config): """Required: Convert state dict format""" @@ -340,8 +351,8 @@ When implementing a new model for NeuronxDistributed, ensure: The key insight for successful compilation is understanding that NeuronxDistributedInference uses a different pattern than typical PyTorch models: 1. **Base class handles forward pass** - Don't implement custom forward methods -2. **Consistent return formats** - All layers must return expected tuple formats +2. **Consistent return formats** - All layers must return expected tuple formats 3. **Framework compliance** - Follow established patterns from existing models 4. **Proper initialization** - Use `init_model` for component setup, not `__init__` -Following these patterns ensures smooth compilation and integration with the NeuronX hardware optimization pipeline. \ No newline at end of file +Following these patterns ensures smooth compilation and integration with the NeuronX hardware optimization pipeline. diff --git a/skills/neuron-framework-autoport/references/knowledge_base/genericmoe_port_session.md b/skills/neuron-framework-autoport/references/knowledge_base/genericmoe_port_session.md index ab0e54d..a709170 100644 --- a/skills/neuron-framework-autoport/references/knowledge_base/genericmoe_port_session.md +++ b/skills/neuron-framework-autoport/references/knowledge_base/genericmoe_port_session.md @@ -20,6 +20,7 @@ Successfully ported a Mixture-of-Experts model from HuggingFace Transformers to ## Task Breakdown ### 1. Initial Analysis Phase + - Reviewed knowledge base for MoE patterns (Generic MoE 29B model) - Analyzed NeuronxDistributed and NeuronxDistributedInference codebases - Studied reference implementations (Qwen3-MoE, Mixtral) @@ -28,12 +29,14 @@ Successfully ported a Mixture-of-Experts model from HuggingFace Transformers to ### 2. Implementation Phase **Created Files**: + - `neuron_port/modeling_genericmoe_neuron.py` - Complete model implementation - `neuron_port/README_GENERICMOE_PORT.md` - Documentation - `agent_artifacts/tmp/compile_genericmoe.py` - Compilation script - `agent_artifacts/tmp/test_genericmoe_inference.py` - Inference test script **Key Components Implemented**: + ```python # Configuration class GenericMoEInferenceConfig(InferenceConfig): @@ -70,6 +73,7 @@ def convert_genericmoe_hf_to_neuron_state_dict(): ### 3. Compilation Phase **First Attempt - FAILED** ❌ + ```python config = CompilationConfig( model_class=NeuronGenericMoEForCausalLM, @@ -85,22 +89,26 @@ config = CompilationConfig( ``` **Issue**: `use_fp16=False` resulted in torch.float32, causing: + - 2x memory usage - Hundreds of casting warnings: `bfloat16 → float32` - Slower processing **Root Cause**: Didn't read model_compiler.py source first + ```python # From model_compiler.py:102 dtype = torch.bfloat16 if self.config.use_fp16 else torch.float32 ``` **Solution**: Kill compilation, fix parameter, restart + ```python use_fp16=True, # ✅ CORRECT: True = bfloat16, False = float32 ``` **Second Attempt - SUCCESS** ✅ + - Cleared cache: `rm -rf /var/tmp/neuron-compile-cache` - Restarted with correct dtype - Compilation completed in ~13 minutes @@ -109,26 +117,34 @@ use_fp16=True, # ✅ CORRECT: True = bfloat16, False = float32 ### 4. Inference Phase **Attempt 1 - NameError** ❌ + ```python model_class=NeuronGenericMoeForCausalLM, # ❌ Wrong: Moe vs MoE ``` + **Fix**: Corrected class name typo **Attempt 2 - AttributeError** ❌ + ```python # Error: 'NeuronConfig' object has no attribute 'router_config' if self.neuron_config.router_config is not None: # ❌ No hasattr check ``` + **Fix**: Added conditional check + ```python if hasattr(self.neuron_config, 'router_config') and self.neuron_config.router_config is not None: ``` **Attempt 3 - AttributeError** ❌ + ```python # Error: 'GenericMoEInferenceConfig' object has no attribute 'output_attentions' ``` + **Fix**: Added standard HuggingFace config attributes + ```python # Standard HF attributes (needed for inference) if not hasattr(self, 'output_attentions'): @@ -140,6 +156,7 @@ if not hasattr(self, 'return_dict'): ``` **Attempt 4 - SUCCESS** ✅ + ``` Prompt: What is the capital of France? Response: The capital of France is Paris. Paris is not only the capital but @@ -159,6 +176,7 @@ Performance: 4.9 tokens/second, 10.3 seconds total (50 tokens) **Decision**: Use `nn.LayerNorm` instead of RMSNorm for ALL normalization layers **Rationale** (from Generic MoE knowledge base): + - AWS Neuron's custom RMSNorm kernel has subtle numerical differences - Lack of mean-centering causes activation drift across deep layers - MoE router is extremely sensitive to input distribution @@ -167,6 +185,7 @@ Performance: 4.9 tokens/second, 10.3 seconds total (50 tokens) - **Result**: Using RMSNorm produces gibberish/repetitive output **Implementation**: + ```python # In NeuronGenericMoEDecoderLayer self.input_layernorm = nn.LayerNorm( @@ -186,12 +205,14 @@ self.post_attention_layernorm = nn.LayerNorm( **Decision**: bfloat16 (native model dtype) **Why**: + - Native dtype from HuggingFace model - Half memory vs float32 - Better performance - No accuracy loss for this model **Configuration**: + ```python use_fp16=True # Confusing naming, but this gives bfloat16 # Results in: torch.bfloat16 @@ -202,6 +223,7 @@ use_fp16=True # Confusing naming, but this gives bfloat16 **Decision**: TP=16, EP=1 (tensor parallelism only, no expert parallelism) **Rationale**: + - Expert parallelism (EP>1) not supported for token generation - With TP=16, expert weights sharded across dimensions - Memory per rank: ~5.5GB (manageable on 16GB HBM) @@ -213,6 +235,7 @@ use_fp16=True # Confusing naming, but this gives bfloat16 **Decision**: FP32 router with softmax activation **Implementation**: + ```python if hasattr(self.neuron_config, 'router_config') and self.neuron_config.router_config is not None: self.neuron_config.router_config.dtype = torch.float32 @@ -220,6 +243,7 @@ if hasattr(self.neuron_config, 'router_config') and self.neuron_config.router_co ``` **Why**: + - MoE routing extremely sensitive to numerical precision - FP32 prevents routing weight quantization errors - Softmax matches model's routing implementation @@ -231,6 +255,7 @@ if hasattr(self.neuron_config, 'router_config') and self.neuron_config.router_co **Flag**: `--internal-hlo2tensorizer-options='--verify-hlo=false'` **Rationale**: + - HLO verifier fails with "Expert routing patterns not recognized" - Dynamic expert routing creates conditional computation graphs - Disabled verifier + comprehensive post-compilation validation @@ -273,15 +298,18 @@ if hasattr(self.neuron_config, 'router_config') and self.neuron_config.router_co ## Improvement Strategy for Next Time ### 1. Read APIs First (Most Critical) + ```python # BEFORE writing compile script: Read(model_compiler.py) # Understand parameters Read(inference_config_base.py) # See required attributes Read(moe_neuron_config.py) # MoE-specific config ``` + **Impact**: Would save entire recompilation cycle (13 minutes) ### 2. Study Similar Successful Ports + ```bash # Find most similar model grep -r "MoENeuronConfig" NeuronxDistributedInference/ @@ -289,10 +317,13 @@ grep -r "MoENeuronConfig" NeuronxDistributedInference/ Read(qwen3_moe/modeling_qwen3_moe.py) # Copy pattern, don't reinvent ``` + **Impact**: Would catch all config attributes upfront (save 3 restarts) ### 3. Defensive Config Class Template + Start with this pattern: + ```python def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) @@ -310,6 +341,7 @@ def __init__(self, *args, **kwargs): ``` ### 4. Pre-Compilation Checklist + - [ ] Read compiler API source - [ ] Study 1-2 similar model implementations - [ ] Copy standard config attributes @@ -319,6 +351,7 @@ def __init__(self, *args, **kwargs): - [ ] Validate weight shapes ### 5. Time Savings Estimate + **Actual**: ~45 minutes (compilation + retries) **With improvements**: ~15 minutes (one compilation + one inference) **Savings**: ~30 minutes (67% reduction) @@ -328,6 +361,7 @@ def __init__(self, *args, **kwargs): ## Architecture Details ### Model Specifications + - **Layers**: 32 decoder layers - **Experts**: 16 experts per layer - **Active Experts**: 2 per token (top-2 routing) @@ -339,6 +373,7 @@ def __init__(self, *args, **kwargs): - **Vocab Size**: 32,064 ### Weight Transformation + ```python # HuggingFace format: # w1: gate_proj [intermediate_size, hidden_size] @@ -360,6 +395,7 @@ down_proj = w2.T ## Performance Metrics ### Compilation + - **HLO Generation**: ~13 seconds - **Token Generation Model**: ~105 seconds (PASS) - **Context Encoding Model**: ~20 seconds (PASS) @@ -367,6 +403,7 @@ down_proj = w2.T - **Total**: ~13 minutes ### Inference + - **Weight Loading**: ~11 seconds - **Warmup**: ~0.8 seconds - **Generation**: 10.3 seconds for 50 tokens @@ -374,6 +411,7 @@ down_proj = w2.T - **Latency**: ~200ms per token ### Memory + - **Per NeuronCore**: ~5.5GB HBM - **Total Model**: ~88GB (16 cores × 5.5GB) - **Host RAM During Compilation**: ~188GB peak @@ -413,18 +451,21 @@ performance: ## Critical Files **Model Implementation**: + ``` neuron_port/modeling_genericmoe_neuron.py # 523 lines neuron_port/README_GENERICMOE_PORT.md # Documentation ``` **Scripts**: + ``` agent_artifacts/tmp/compile_genericmoe.py # Compilation agent_artifacts/tmp/test_genericmoe_inference.py # Inference test ``` **Compiled Artifacts**: + ``` agent_artifacts/data/genericmoe_compiled/ ├── config.json @@ -444,6 +485,7 @@ agent_artifacts/data/genericmoe_compiled/ **Test Prompt**: "What is the capital of France?" **Generated Response**: + ``` The capital of France is Paris. Paris is not only the capital but also the largest city in France, known for its history, culture, and landmarks such @@ -451,6 +493,7 @@ as the Eiffel Tower, Notre-Dame Cathedral, and the Louvre... ``` **Validation**: ✅ **PASSED** + - Factually correct - Coherent sentence structure - No repetition or gibberish @@ -461,16 +504,19 @@ as the Eiffel Tower, Notre-Dame Cathedral, and the Louvre... ## Key Takeaways ### What Made This Port Successful + 1. ✅ **Knowledge Base Usage**: LayerNorm decision from Generic MoE example 2. ✅ **Reference Implementation**: Qwen3-MoE weight conversion pattern 3. ✅ **Systematic Debugging**: Clear error messages → targeted fixes ### What Could Be Improved + 1. ❌ **API Documentation Reading**: Should have read compiler source first 2. ❌ **Config Attribute Checklist**: Should have copied all standard attrs upfront 3. ❌ **Dry Run Testing**: Should have tested config instantiation before compilation ### Bottom Line + **"Read reference implementations and API docs BEFORE writing code"** - fastest debugging is the bug you never write. --- diff --git a/skills/neuron-framework-autoport/references/knowledge_base/genericmoe_v14_analysis.md b/skills/neuron-framework-autoport/references/knowledge_base/genericmoe_v14_analysis.md index d44d39f..ed14b9b 100644 --- a/skills/neuron-framework-autoport/references/knowledge_base/genericmoe_v14_analysis.md +++ b/skills/neuron-framework-autoport/references/knowledge_base/genericmoe_v14_analysis.md @@ -1,15 +1,19 @@ # GenericMoE v14 Analysis - Simplified RoPE Approach ## Date + October 27, 2025 ## Summary + v14 successfully compiled with simplified standard RoPE (removing LongRoPE scaling parameters). Compilation was faster due to cached NEFF files. ## Changes in v14 + Removed `use_scaled_rope` parameter from `NeuronGenericMoEAttention.__init__()` in both modeling files: ### Before (v13): + ```python super().__init__( config=config, @@ -25,6 +29,7 @@ super().__init__( ``` ### After (v14): + ```python super().__init__( config=config, @@ -42,6 +47,7 @@ super().__init__( ## Compilation Results ### v14 Compilation (Successful) + - **Time**: 12.8 minutes (using cached NEFFs from v13) - **HLO Generation**: 11.6 seconds - **Token Generation Model**: 0.21 seconds (cached) @@ -50,12 +56,14 @@ super().__init__( - **Status**: ✅ SUCCESS ### Comparison with v13: + - v13: 142 seconds total compilation time - v14: Used cached NEFF files, faster overall ## Expected Inference Behavior Based on v12 and v13 results (which also had vocab masking), we expect v14 to show **SIMILAR gibberish output**: + - Test 1: Empty or very short response - Test 2: "Trans Trans" or similar repetition - Test 3: "Writing Writing..." repetitive gibberish @@ -63,6 +71,7 @@ Based on v12 and v13 results (which also had vocab masking), we expect v14 to sh ### Why Removing LongRoPE Won't Fix Gibberish The investigation document (`longrope_investigation_v13.md`) established that: + 1. `use_scaled_rope=True` was already ineffective in v13 (only affects different code path) 2. LongRoPE mscale multiplier (1.243x) is a relatively small adjustment (~24%) 3. The gibberish persisted despite the flag being present @@ -74,14 +83,17 @@ Therefore, **removing the ineffective flag should not change behavior**. Since RoPE configuration (v9, v13, v14) hasn't resolved gibberish, need to investigate: ### 1. Normalization Type (HIGH PRIORITY) + **Lesson from knowledge_base**: Generic MoE had identical gibberish symptoms due to LayerNorm vs RMSNorm mismatch. **Action Required**: + - Verify GenericMoE uses RMSNorm (not LayerNorm) - Check `get_rmsnorm_cls()` in modeling_genericmoe.py:165-169 - Compare with HuggingFace GenericMoE implementation **Current Implementation**: + ```python def get_rmsnorm_cls(): # If infer on NXD -> CustomRMSNorm @@ -92,24 +104,30 @@ def get_rmsnorm_cls(): **Question**: Is `CustomRMSNorm` compatible with GenericMoE's normalization requirements? ### 2. Weight Loading Verification + - Ensure 0 missing keys in weight loading - Verify state dict conversion is correctly mapping all weights - Check if redundant keys being removed are actually irrelevant ### 3. MoE Router Configuration -From GenericmoeInferenceConfig.__init__() (lines 188-193): + +From GenericmoeInferenceConfig.**init**() (lines 188-193): + ```python self.neuron_config.router_config.dtype = torch.float32 self.neuron_config.router_config.act_fn = "softmax" ``` **Questions**: + - Is router correctly routing to experts? - Are expert weights being applied correctly? - Check if `glu_mlp=True` configuration is working ### 4. Expert Weight Application + From convert_genericmoe_hf_to_neuron_state_dict(): + - Gate projection (w1) and up projection (w3) concatenation - Down projection (w2) mapping - Weight transposition (.T operations) @@ -117,10 +135,12 @@ From convert_genericmoe_hf_to_neuron_state_dict(): **Potential Issue**: Weight dimension ordering mismatch? ## Files Modified + - `/home/ec2-user/agents/hariseldon/NeuroborosFoundations/src/amzn/neuron/neuroboros/models/genericmoe/modeling_genericmoe.py` (line 325) - `/home/ec2-user/agents/hariseldon/neuron_port/modeling_genericmoe.py` (line 325) ## Compilation Artifacts + - Compiled model: `/home/ec2-user/agents/hariseldon/agent_artifacts/data/genericmoe_compiled` - Context encoding model: Compiled (cached) - Token generation model: Compiled (cached) @@ -133,12 +153,14 @@ From convert_genericmoe_hf_to_neuron_state_dict(): **NEXT ACTION**: Investigate normalization type following the lesson from Generic MoE success story in knowledge_base. Focus investigation on: + 1. RMSNorm implementation compatibility 2. Weight loading/mapping correctness 3. MoE router functionality 4. Expert weight application ## Version History + - v9: Fixed rope_theta (1M → 10K) - v10: Added attention_bias and lm_head_bias - v11: Attempted weight truncation (reverted per user feedback) diff --git a/skills/neuron-framework-autoport/references/knowledge_base/genericmoe_v15_layernorm_fix.md b/skills/neuron-framework-autoport/references/knowledge_base/genericmoe_v15_layernorm_fix.md index c613c9e..11413ff 100644 --- a/skills/neuron-framework-autoport/references/knowledge_base/genericmoe_v15_layernorm_fix.md +++ b/skills/neuron-framework-autoport/references/knowledge_base/genericmoe_v15_layernorm_fix.md @@ -13,12 +13,14 @@ Identified and fixed the **root cause** of GenericMoE gibberish output by invest ### The Issue **HuggingFace GenericMoE Implementation**: + ```python # transformers/src/transformers/models/genericmoe/modeling_genericmoe.py:657 self.norm = nn.LayerNorm(config.hidden_size, eps=config.rms_norm_eps, elementwise_affine=True) ``` **Previous Neuron Implementation (v14 and earlier)**: + ```python # modeling_genericmoe.py:461 self.norm = get_rmsnorm_cls()(self.hidden_size, eps=config.rms_norm_eps) @@ -29,6 +31,7 @@ self.norm = get_rmsnorm_cls()(self.hidden_size, eps=config.rms_norm_eps) The final model normalization layer used **RMSNorm** in the Neuron implementation, but HuggingFace GenericMoE uses **LayerNorm**. This normalization type mismatch caused the model to produce gibberish output. **Key Observations**: + - Decoder layers correctly use RMSNorm: `input_layernorm` and `post_attention_layernorm` - Only the final model normalization (`self.norm`) uses LayerNorm - The parameter name `rms_norm_eps` in config.json (value: 1e-05) is misleading - it's used for both RMSNorm AND LayerNorm epsilon values @@ -36,6 +39,7 @@ The final model normalization layer used **RMSNorm** in the Neuron implementatio ### Knowledge Base Pattern Match This is **identical** to the Generic MoE debugging case documented in the knowledge_base: + - Generic MoE had gibberish/repetitive output - Root cause: LayerNorm vs RMSNorm mismatch - After fixing normalization type: 100% accuracy achieved @@ -45,6 +49,7 @@ This is **identical** to the Generic MoE debugging case documented in the knowle ### File: `NeuroborosFoundations/src/amzn/neuron/neuroboros/models/genericmoe/modeling_genericmoe.py` **Lines 460-463** (Changed from v14): + ```python # Final normalization layer # CRITICAL FIX v15: HuggingFace GenericMoE uses LayerNorm (not RMSNorm) for final model normalization @@ -53,6 +58,7 @@ self.norm = nn.LayerNorm(self.hidden_size, eps=config.rms_norm_eps, elementwise_ ``` **Previous v14 code**: + ```python # Final normalization layer self.norm = get_rmsnorm_cls()(self.hidden_size, eps=config.rms_norm_eps) @@ -61,6 +67,7 @@ self.norm = get_rmsnorm_cls()(self.hidden_size, eps=config.rms_norm_eps) ### File: `neuron_port/modeling_genericmoe.py` **Lines 476-479** (Changed from v14): + ```python # Final normalization layer # CRITICAL FIX v15: HuggingFace GenericMoE uses LayerNorm (not RMSNorm) for final model normalization @@ -71,21 +78,25 @@ self.norm = nn.LayerNorm(self.hidden_size, eps=config.rms_norm_eps, elementwise_ ## Investigation Process ### 1. Read HuggingFace GenericMoE Source + - Located GenericMoE implementation in local transformers directory - Found `genericmoeRMSNorm` class at line 567-584 - Discovered final model normalization uses `nn.LayerNorm` at line 657 ### 2. Examined Neuron Implementation + - Reviewed `get_rmsnorm_cls()` function (lines 165-169) - Checked `CustomRMSNorm` in NeuronxDistributedInference - Found all decoder layer normalizations correctly use RMSNorm ### 3. Identified Mismatch + - Final model normalization: **LayerNorm** (HF) vs **RMSNorm** (Neuron) - Config parameter: `rms_norm_eps = 1e-05` - This matches Generic MoE debugging pattern exactly ### 4. Applied Fix + - Changed `self.norm` from `get_rmsnorm_cls()` to `nn.LayerNorm` - Maintained `elementwise_affine=True` to match HuggingFace - Updated both main and backup modeling files @@ -93,6 +104,7 @@ self.norm = nn.LayerNorm(self.hidden_size, eps=config.rms_norm_eps, elementwise_ ## Compilation Details ### Configuration + - **Model Path**: `agent_artifacts/data/generic-moe-model` - **Output Path**: `agent_artifacts/data/genericmoe_compiled` - **TP Degree**: 16 @@ -101,17 +113,20 @@ self.norm = nn.LayerNorm(self.hidden_size, eps=config.rms_norm_eps, elementwise_ - **Precision**: bfloat16 ### Compilation Started + - **Start Time**: 2025-10-27 03:56:12 UTC - **Expected Duration**: 30-60 minutes - **Status**: Running in background (shell ID: f26f15) ### Cache Handling + - Cleared `/tmp/neuron-compile-cache/` before compilation - Cleared `/var/tmp/neuron-compile-cache/` after initial failure due to cached failed NEFF ## Expected Outcome Based on the Generic MoE pattern from knowledge_base: + - **Before fix**: Gibberish/repetitive output (identical to v9-v14) - **After fix**: Should produce coherent, accurate responses - **Generic MoE result**: 100% accuracy after fixing normalization type @@ -119,6 +134,7 @@ Based on the Generic MoE pattern from knowledge_base: ## Version History Context ### Eliminated Causes (v9-v14) + - ✅ v9: Fixed `rope_theta` (1M → 10K) - ✅ v10: Added `attention_bias` and `lm_head_bias` - ✅ v11: Attempted weight truncation (reverted per user feedback) @@ -129,6 +145,7 @@ Based on the Generic MoE pattern from knowledge_base: **All v9-v14 versions produced identical gibberish output**, confirming those fixes did not address the root cause. ### v15 Fix (THIS VERSION) + - **Target**: Normalization type mismatch in final model layer - **Change**: RMSNorm → LayerNorm for `self.norm` - **Confidence**: HIGH (matches proven Generic MoE fix pattern) diff --git a/skills/neuron-framework-autoport/references/knowledge_base/genericmoe_v16_complete_layernorm_fix.md b/skills/neuron-framework-autoport/references/knowledge_base/genericmoe_v16_complete_layernorm_fix.md index 157927d..b8e8aad 100644 --- a/skills/neuron-framework-autoport/references/knowledge_base/genericmoe_v16_complete_layernorm_fix.md +++ b/skills/neuron-framework-autoport/references/knowledge_base/genericmoe_v16_complete_layernorm_fix.md @@ -11,13 +11,16 @@ Successfully identified the **complete root cause** by comparing with successful ## Critical Discovery: v15 Was Incomplete ### v15 Fix (Partial - Still Produced Gibberish) + **Only changed final model normalization**: + ```python # modeling_genericmoe.py:463 (v15) self.norm = nn.LayerNorm(self.hidden_size, eps=config.rms_norm_eps, elementwise_affine=True) ``` **Left decoder layers as RMSNorm**: + ```python # modeling_genericmoe.py:356-363 (v15 - INCOMPLETE) self.input_layernorm = get_rmsnorm_cls()(config.hidden_size, eps=config.rms_norm_eps) @@ -48,6 +51,7 @@ self.post_attention_layernorm = nn.LayerNorm(config.hidden_size, eps=config.rms_ **File 1**: `NeuroborosFoundations/src/amzn/neuron/neuroboros/models/genericmoe/modeling_genericmoe.py` **Lines 355-358** (Decoder layer normalizations): + ```python # CRITICAL FIX v16: Use LayerNorm for ALL normalization layers to match successful port # The successful port uses LayerNorm for decoder layers, not RMSNorm @@ -56,6 +60,7 @@ self.post_attention_layernorm = nn.LayerNorm(config.hidden_size, eps=config.rms_ ``` **Lines 460-463** (Final model normalization - from v15): + ```python # CRITICAL FIX v15: HuggingFace GenericMoE uses LayerNorm (not RMSNorm) for final model normalization # This matches the pattern from Generic MoE where LayerNorm vs RMSNorm mismatch caused gibberish output @@ -66,15 +71,16 @@ self.norm = nn.LayerNorm(self.hidden_size, eps=config.rms_norm_eps, elementwise_ ### Summary of ALL v16 Normalization Changes -| Layer | v9-v14 | v15 | v16 (This Version) | Successful Port | -|-------|--------|-----|-------------------|-----------------| -| `input_layernorm` | RMSNorm | RMSNorm | **LayerNorm** ✅ | LayerNorm | -| `post_attention_layernorm` | RMSNorm | RMSNorm | **LayerNorm** ✅ | LayerNorm | -| `self.norm` (final) | RMSNorm | LayerNorm | **LayerNorm** ✅ | LayerNorm | +| Layer | v9-v14 | v15 | v16 (This Version) | Successful Port | +| -------------------------- | ------- | --------- | ------------------ | --------------- | +| `input_layernorm` | RMSNorm | RMSNorm | **LayerNorm** ✅ | LayerNorm | +| `post_attention_layernorm` | RMSNorm | RMSNorm | **LayerNorm** ✅ | LayerNorm | +| `self.norm` (final) | RMSNorm | LayerNorm | **LayerNorm** ✅ | LayerNorm | ## Version History ### v9-v14: Incremental Fixes (All Produced Gibberish) + - v9: Fixed `rope_theta` (1M → 10K) - v10: Added `attention_bias` and `lm_head_bias` - v11: Weight truncation attempt (reverted) @@ -83,11 +89,13 @@ self.norm = nn.LayerNorm(self.hidden_size, eps=config.rms_norm_eps, elementwise_ - v14: Simplified RoPE (removed `use_scaled_rope`) ### v15: Partial Fix (Still Gibberish) + - Changed **only final** `self.norm` to LayerNorm - Left decoder layers as RMSNorm - Result: Still produced gibberish output identical to v14 ### v16: Complete Fix (This Version) + - Changed **ALL** normalization layers to LayerNorm - Matches successful port backup exactly - Expected: Should produce coherent output @@ -95,6 +103,7 @@ self.norm = nn.LayerNorm(self.hidden_size, eps=config.rms_norm_eps, elementwise_ ## Compilation Details ### Configuration + - **Model Path**: `agent_artifacts/data/generic-moe-model` - **Output Path**: `agent_artifacts/data/genericmoe_compiled` - **TP Degree**: 16 @@ -103,6 +112,7 @@ self.norm = nn.LayerNorm(self.hidden_size, eps=config.rms_norm_eps, elementwise_ - **Precision**: bfloat16 ### Compilation Progress + - **Start Time**: 2025-10-27 19:11:55 UTC - **HLO Generation**: 13.6 seconds - **Expected Duration**: 30-60 minutes @@ -119,6 +129,7 @@ self.norm = nn.LayerNorm(self.hidden_size, eps=config.rms_norm_eps, elementwise_ ### LayerNorm vs RMSNorm Differences **LayerNorm**: + ```python mean = x.mean(dim=-1, keepdim=True) var = x.var(dim=-1, keepdim=True, unbiased=False) @@ -127,6 +138,7 @@ output = weight * x_normalized + bias # if elementwise_affine=True ``` **RMSNorm** (simpler, no mean subtraction): + ```python rms = sqrt(x.pow(2).mean(dim=-1, keepdim=True) + eps) x_normalized = x / rms @@ -138,15 +150,18 @@ The successful port found that **LayerNorm works better** for GenericMoE on AWS ## Files Modified ### Code Changes + 1. `NeuroborosFoundations/src/amzn/neuron/neuroboros/models/genericmoe/modeling_genericmoe.py` (lines 355-358, 460-463) 2. `neuron_port/modeling_genericmoe.py` (lines 371-374, 476-479) ### New Files Created + 1. `agent_artifacts/tmp/compile_genericmoe_v16_all_layernorm.py` - Compilation script 2. `agent_artifacts/traces/compile_genericmoe_v16_all_layernorm.log` - Compilation log 3. `agent_artifacts/traces/genericmoe_v16_complete_layernorm_fix.md` - This document ### Downloaded for Comparison + 1. `successful_port/modeling_genericmoe_working.py` - Working implementation from S3 backup ## Next Steps @@ -160,11 +175,13 @@ The successful port found that **LayerNorm works better** for GenericMoE on AWS ## Expected Test Results ### v15 Results (Gibberish - For Comparison) + - Test 1: "The capital of France is Paris is correct. The capital is capital is capital..." - Test 2: Empty output - Test 3: "The fibbyline is fibbyline..." ### v16 Expected Results + - Test 1: "The capital of France is Paris." - Test 2: "A mixture of experts model..." - Test 3: "def fibonacci(n):..." diff --git a/skills/neuron-framework-autoport/references/knowledge_base/genericmoe_v16_final_success_summary.md b/skills/neuron-framework-autoport/references/knowledge_base/genericmoe_v16_final_success_summary.md index 33bf8a6..80422f9 100644 --- a/skills/neuron-framework-autoport/references/knowledge_base/genericmoe_v16_final_success_summary.md +++ b/skills/neuron-framework-autoport/references/knowledge_base/genericmoe_v16_final_success_summary.md @@ -13,24 +13,27 @@ Successfully ported **GenericMoE (generic-moe-model)** to AWS Neuron hardware af ## Critical Fix: Complete LayerNorm Migration ### The Problem + - v9-v14: All produced gibberish/repetitive output - v15: Changed only final `self.norm` to LayerNorm → **Still gibberish** - Root cause: Incomplete normalization layer fix ### The Solution (v16) + Changed **ALL THREE** normalization layers from RMSNorm to LayerNorm: -| Layer | Location | v9-v15 | v16 (Working) | -|-------|----------|--------|---------------| -| `input_layernorm` | Decoder layer pre-attention | RMSNorm | **LayerNorm** ✅ | -| `post_attention_layernorm` | Decoder layer pre-MoE | RMSNorm | **LayerNorm** ✅ | -| `self.norm` | Final model normalization | RMSNorm (v9-v14) / LayerNorm (v15) | **LayerNorm** ✅ | +| Layer | Location | v9-v15 | v16 (Working) | +| -------------------------- | --------------------------- | ---------------------------------- | ---------------- | +| `input_layernorm` | Decoder layer pre-attention | RMSNorm | **LayerNorm** ✅ | +| `post_attention_layernorm` | Decoder layer pre-MoE | RMSNorm | **LayerNorm** ✅ | +| `self.norm` | Final model normalization | RMSNorm (v9-v14) / LayerNorm (v15) | **LayerNorm** ✅ | ### Code Changes **File 1**: `NeuroborosFoundations/src/amzn/neuron/neuroboros/models/genericmoe/modeling_genericmoe.py` **Lines 355-358** (Decoder layer): + ```python # CRITICAL FIX v16: Use LayerNorm for ALL normalization layers to match successful port # The successful port uses LayerNorm for decoder layers, not RMSNorm @@ -39,6 +42,7 @@ self.post_attention_layernorm = nn.LayerNorm(config.hidden_size, eps=config.rms_ ``` **Lines 460-463** (Final normalization): + ```python # CRITICAL FIX v15: HuggingFace GenericMoE uses LayerNorm (not RMSNorm) for final model normalization # This matches the pattern from Generic MoE where LayerNorm vs RMSNorm mismatch caused gibberish output @@ -52,6 +56,7 @@ self.norm = nn.LayerNorm(self.hidden_size, eps=config.rms_norm_eps, elementwise_ ## Compilation Details ### Configuration + - **Model**: generic-moe-model (41B parameters, 16 experts) - **Hardware**: AWS Trainium (trn1.32xlarge, 32 NeuronCores) - **Tensor Parallelism**: 16-way (TP=16) @@ -60,16 +65,18 @@ self.norm = nn.LayerNorm(self.hidden_size, eps=config.rms_norm_eps, elementwise_ - **Precision**: bfloat16 ### Compilation Time Breakdown -| Phase | Duration | -|-------|----------| -| HLO Generation | 13.6 seconds | -| Token Generation Model Compilation | 100.8 seconds (~1.7 min) | -| Context Encoding Model Compilation | 15.9 seconds | -| Model Building | 152.8 seconds (~2.5 min) | -| **Weight Sharding** | **722.7 seconds (~12 min)** | -| **Total** | **~14.7 minutes** | + +| Phase | Duration | +| ---------------------------------- | --------------------------- | +| HLO Generation | 13.6 seconds | +| Token Generation Model Compilation | 100.8 seconds (~1.7 min) | +| Context Encoding Model Compilation | 15.9 seconds | +| Model Building | 152.8 seconds (~2.5 min) | +| **Weight Sharding** | **722.7 seconds (~12 min)** | +| **Total** | **~14.7 minutes** | ### Output Artifacts + - **Location**: `agent_artifacts/data/genericmoe_compiled/` - **Weight Shards**: 16 files (one per TP rank) - **Config Files**: `neuron_config.json`, `config.json` @@ -82,21 +89,25 @@ self.norm = nn.LayerNorm(self.hidden_size, eps=config.rms_norm_eps, elementwise_ ### Test Results (v16 - All PASSED ✅) **Test 1: Factual Question** + - **Prompt**: "What is the capital of France?" - **Output**: "The capital of France is Paris. It is not only the largest city in France but also serves as the country's political, cultural, and economic center. Paris is renowned for its history, art, architecture, and fashion..." - **Status**: ✅ Coherent, accurate response **Test 2: Technical Explanation** + - **Prompt**: "Explain what a mixture of experts model is in one sentence." - **Output**: "A mixture of experts model is an ensemble learning approach that combines the outputs of multiple specialized models to improve overall prediction accuracy." - **Status**: ✅ Clear, concise explanation **Test 3: Code Generation** + - **Prompt**: "Write a Python function to calculate fibonacci numbers." - **Output**: "Certainly! Below is a Python function that calculates Fibonacci numbers using both iterative and recursive approaches. I'll start with the iterative approach, which is more efficient in terms of time and space complexity.\n\n```python\ndef fibonacci_iterative(n):\n..." - **Status**: ✅ Working code with explanation ### Performance Metrics + - **Throughput**: 5.5 tokens/second - **Success Rate**: 3/3 tests (100%) - **Inference Latency**: @@ -109,11 +120,13 @@ self.norm = nn.LayerNorm(self.hidden_size, eps=config.rms_norm_eps, elementwise_ ## Comparison: v15 vs v16 ### v15 Results (Partial Fix - GIBBERISH) + - **Test 1**: "The capital of France is Paris is correct. The capital is capital is capital is capital..." - **Test 2**: Empty output - **Test 3**: "The fibbyline is fibbyline is fibbyline..." ### v16 Results (Complete Fix - WORKING) + - **Test 1**: ✅ Coherent factual answer about Paris - **Test 2**: ✅ Clear MoE explanation - **Test 3**: ✅ Working Python code @@ -125,6 +138,7 @@ self.norm = nn.LayerNorm(self.hidden_size, eps=config.rms_norm_eps, elementwise_ ## Version History ### v9-v14: Progressive Fixes (All Gibberish) + - **v9**: Fixed `rope_theta` (1M → 10K) - **v10**: Added `attention_bias` and `lm_head_bias` - **v11**: Weight truncation attempt (reverted) @@ -133,11 +147,13 @@ self.norm = nn.LayerNorm(self.hidden_size, eps=config.rms_norm_eps, elementwise_ - **v14**: Simplified RoPE (removed ineffective `use_scaled_rope`) ### v15: Partial LayerNorm Fix (Still Gibberish) + - Changed only final `self.norm` from RMSNorm to LayerNorm - Result: Identical gibberish to v14 - **Learning**: Incomplete fix - decoder layers still used RMSNorm ### v16: Complete LayerNorm Fix (SUCCESS) + - Changed ALL three normalization layers to LayerNorm - Downloaded successful port backup from S3 for comparison - Discovered successful port uses LayerNorm for **ALL** layers @@ -150,6 +166,7 @@ self.norm = nn.LayerNorm(self.hidden_size, eps=config.rms_norm_eps, elementwise_ ### LayerNorm vs RMSNorm on Neuron Hardware **LayerNorm** (what works): + ```python mean = x.mean(dim=-1, keepdim=True) var = x.var(dim=-1, keepdim=True, unbiased=False) @@ -158,6 +175,7 @@ output = weight * x_normalized + bias # if elementwise_affine=True ``` **RMSNorm** (caused gibberish): + ```python rms = sqrt(x.pow(2).mean(dim=-1, keepdim=True) + eps) x_normalized = x / rms @@ -169,6 +187,7 @@ output = weight * x_normalized # no bias, no mean subtraction ### HuggingFace Source Code Discrepancy **HuggingFace Implementation**: + ```python # transformers/src/transformers/models/genericmoe/modeling_genericmoe.py class GenericmoeDecoderLayer(nn.Module): @@ -178,6 +197,7 @@ class GenericmoeDecoderLayer(nn.Module): ``` **Successful Neuron Port**: + ```python # Must use LayerNorm for all layers self.input_layernorm = nn.LayerNorm(config.hidden_size, eps=config.rms_norm_eps, elementwise_affine=True) @@ -191,6 +211,7 @@ This discrepancy highlights hardware-specific requirements that may not match re ## Files Modified ### Core Model Files + 1. `NeuroborosFoundations/src/amzn/neuron/neuroboros/models/genericmoe/modeling_genericmoe.py` - Lines 355-358: Decoder layer normalization - Lines 460-463: Final model normalization @@ -200,16 +221,19 @@ This discrepancy highlights hardware-specific requirements that may not match re - Lines 476-479: Final model normalization ### Compilation & Testing Scripts + 1. `agent_artifacts/tmp/compile_genericmoe_v16_all_layernorm.py` - v16 compilation script 2. `agent_artifacts/tmp/test_genericmoe_v14_inference.py` - Inference test script (reused for all versions) ### Trace Files + 1. `agent_artifacts/traces/compile_genericmoe_v16_all_layernorm.log` - Compilation log 2. `agent_artifacts/traces/inference_test_v16_complete_layernorm.log` - Inference test results 3. `agent_artifacts/traces/genericmoe_v16_complete_layernorm_fix.md` - Technical analysis 4. `agent_artifacts/traces/genericmoe_v16_final_success_summary.md` - This document ### Reference Files + 1. `successful_port/modeling_genericmoe_working.py` - Working implementation from S3 backup --- @@ -217,6 +241,7 @@ This discrepancy highlights hardware-specific requirements that may not match re ## Deployment Instructions ### Prerequisites + ```bash # Ensure you're on AWS Trainium instance neuron-ls # Should show 32 NeuronCores @@ -226,6 +251,7 @@ source /opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/bin/activate ``` ### Compilation + ```bash cd /home/ec2-user/agents/hariseldon @@ -235,6 +261,7 @@ PYTHONPATH="./NeuroborosFoundations/src:$PYTHONPATH" \ ``` ### Inference + ```python from amzn.neuron.neuroboros.utils.run_inference import run_inference_with_classes from amzn.neuron.neuroboros.models.genericmoe.modeling_genericmoe import ( diff --git a/skills/neuron-framework-autoport/references/knowledge_base/issues_analysis_and_improved_prompts.md b/skills/neuron-framework-autoport/references/knowledge_base/issues_analysis_and_improved_prompts.md index d3c65cd..b26f155 100644 --- a/skills/neuron-framework-autoport/references/knowledge_base/issues_analysis_and_improved_prompts.md +++ b/skills/neuron-framework-autoport/references/knowledge_base/issues_analysis_and_improved_prompts.md @@ -3,18 +3,22 @@ ## Summary of Issues Encountered ### 1. **Initial Analysis and Architecture Understanding** + - **Issue**: Needed comprehensive understanding of both NeuronxDistributed and NeuronxDistributedInference frameworks - **Resolution**: Thorough analysis of model architectures, attention mechanisms, and framework patterns ### 2. **Model Implementation Challenges** + - **Issue**: Creating a complete Llama3 port from CUDA implementation to NeuronxDistributed - **Resolution**: Built comprehensive implementation with proper base class inheritance and framework compliance ### 3. **Configuration and Weight Handling** + - **Issue**: Multiple checkpoint formats (original Llama3, HuggingFace converted) - **Resolution**: Created multi-format conversion system with auto-detection ### 4. **Package Installation and Dependencies** + - **Issue**: Proper package structure and dependency management - **Resolution**: Created proper setup.py and package structure @@ -23,8 +27,9 @@ Here are your original prompts with augmentations to prevent future issues: ### Original Prompt 1 (Augmented): + ``` -Analyze both of these projects, and particularly the src and test directories as well as the docs directories. The project contains the NeuronSDK source covering architecture and model definitions for a set of models and some implementation guides for models. +Analyze both of these projects, and particularly the src and test directories as well as the docs directories. The project contains the NeuronSDK source covering architecture and model definitions for a set of models and some implementation guides for models. AUGMENTED VERSION: Analyze both NeuronxDistributed and NeuronxDistributedInference projects comprehensively: @@ -56,13 +61,14 @@ Provide architectural descriptions with code examples and highlight common trait ``` ### Original Prompt 2 (Augmented): + ``` Based on your understanding, including all of the existing available components in the Neuron SDK you have analyzed, please now analyze a CUDA specific implementation of Llama3 and create a version that works on the neuronx-distributed framework. AUGMENTED VERSION: Create a comprehensive Llama3 implementation for NeuronxDistributed framework with these requirements: -1. **Source Analysis**: +1. **Source Analysis**: - Analyze CUDA implementation in /home/ec2-user/NeuronxSDK/source/llama3 - Study model architecture, generation logic, and configuration handling - Document all architectural components and their relationships @@ -103,7 +109,8 @@ Ensure the implementation is production-ready with proper error handling, valida ``` ### Original Prompt 3 (Augmented): -``` + +```` Please proceed in the steps mentioned above to convert weights, compile and test model and run inference. AUGMENTED VERSION: @@ -116,55 +123,59 @@ Execute the complete Llama3 implementation pipeline with comprehensive validatio - Ensure sufficient system resources 2. **Step-by-step Execution with Validation**: - + **Step 1: Package Installation** ```bash cd neuronx_llama3 pip install -e . # Validate: Check package imports work correctly python -c "import neuronx_llama3; print('Package installed successfully')" - ``` +```` - **Step 2: Weight Conversion with Multi-format Support** - ```bash - python convert_checkpoint.py \ - --input_path /home/ec2-user/.llama/checkpoints/Llama3.2-1B \ - --output_path ./converted \ - --validate_conversion \ - --verbose - # Validate: Check converted weights match original architecture - ``` - - **Step 3: Model Testing (Pre-compilation)** - ```bash - python test_model.py \ - --checkpoint_path ./converted \ - --test_levels config,model,forward \ - --verbose - # Validate: Ensure model loads and forward pass works - ``` - - **Step 4: Model Compilation** - ```bash - python compile_model.py \ - --checkpoint_path ./converted \ - --output_path ./compiled \ - --batch_size 1 \ - --sequence_length 128 \ - --verbose - # Validate: Check compilation succeeds and artifacts are created - ``` - - **Step 5: Inference Testing** - ```bash - python run_inference.py \ - --model_path ./compiled \ - --prompt "The meaning of life is" \ - --max_tokens 50 \ - --temperature 0.7 \ - --verbose - # Validate: Check inference produces coherent output - ``` +**Step 2: Weight Conversion with Multi-format Support** + +```bash +python convert_checkpoint.py \ + --input_path /home/ec2-user/.llama/checkpoints/Llama3.2-1B \ + --output_path ./converted \ + --validate_conversion \ + --verbose +# Validate: Check converted weights match original architecture +``` + +**Step 3: Model Testing (Pre-compilation)** + +```bash +python test_model.py \ + --checkpoint_path ./converted \ + --test_levels config,model,forward \ + --verbose +# Validate: Ensure model loads and forward pass works +``` + +**Step 4: Model Compilation** + +```bash +python compile_model.py \ + --checkpoint_path ./converted \ + --output_path ./compiled \ + --batch_size 1 \ + --sequence_length 128 \ + --verbose +# Validate: Check compilation succeeds and artifacts are created +``` + +**Step 5: Inference Testing** + +```bash +python run_inference.py \ + --model_path ./compiled \ + --prompt "The meaning of life is" \ + --max_tokens 50 \ + --temperature 0.7 \ + --verbose +# Validate: Check inference produces coherent output +``` 3. **Error Handling and Recovery**: - At each step, check for errors and provide clear diagnostics @@ -185,6 +196,7 @@ Execute the complete Llama3 implementation pipeline with comprehensive validatio - Create summary report of all operations Execute with verbose logging and validate each step before proceeding to the next. + ``` ## Key Improvements in Augmented Prompts @@ -235,4 +247,5 @@ Execute with verbose logging and validate each step before proceeding to the nex 9. **Create fallback strategies for common failure modes** 10. **Validate each step before proceeding to the next** -These augmented prompts should help prevent the issues you encountered and provide a more robust implementation process for future model ports. \ No newline at end of file +These augmented prompts should help prevent the issues you encountered and provide a more robust implementation process for future model ports. +``` diff --git a/skills/neuron-framework-autoport/references/knowledge_base/keyerror_fix_detailed_explanation.md b/skills/neuron-framework-autoport/references/knowledge_base/keyerror_fix_detailed_explanation.md index 962a27e..6cec12a 100644 --- a/skills/neuron-framework-autoport/references/knowledge_base/keyerror_fix_detailed_explanation.md +++ b/skills/neuron-framework-autoport/references/knowledge_base/keyerror_fix_detailed_explanation.md @@ -3,6 +3,7 @@ ## Problem Summary During the Phi3 model compilation for NeuronX, we encountered a critical KeyError: + ``` KeyError: 'layers.0.self_attn.qkv_proj.q_proj.weight' ``` @@ -14,10 +15,12 @@ This error occurred because the NeuronX framework expected weights in a specific ### 1. Weight Structure Mismatch **HuggingFace Phi3 Format:** + - Uses fused QKV projections: `layers.{i}.self_attn.qkv_proj.weight` - Uses fused gate/up projections: `layers.{i}.mlp.gate_up_proj.weight` **NeuronX Expected Format:** + - Expects separate projections: `layers.{i}.self_attn.qkv_proj.q_proj.weight`, `k_proj.weight`, `v_proj.weight` - Expects separate MLP projections: `layers.{i}.mlp.gate_proj.weight`, `up_proj.weight` @@ -36,51 +39,51 @@ Added `convert_hf_to_neuron_state_dict` method to split fused weights: def convert_hf_to_neuron_state_dict(hf_state_dict, config): """Convert HuggingFace state dict to NeuronX format""" neuron_state_dict = {} - + # Calculate dimensions hidden_size = config.hidden_size num_attention_heads = config.num_attention_heads num_key_value_heads = config.num_key_value_heads head_dim = hidden_size // num_attention_heads - + q_hidden_size = num_attention_heads * head_dim kv_hidden_size = num_key_value_heads * head_dim - + for key, tensor in hf_state_dict.items(): if key.startswith('model.'): key = key[6:] # Remove 'model.' prefix - + if 'self_attn.qkv_proj.weight' in key: # Split QKV fused weight into separate Q, K, V weights layer_idx = key.split('.')[1] - + # Split the fused QKV weight q_weight = tensor[:q_hidden_size, :] k_weight = tensor[q_hidden_size:q_hidden_size + kv_hidden_size, :] v_weight = tensor[q_hidden_size + kv_hidden_size:, :] - + # Create separate weight keys base_key = f"layers.{layer_idx}.self_attn.qkv_proj" neuron_state_dict[f"{base_key}.q_proj.weight"] = q_weight neuron_state_dict[f"{base_key}.k_proj.weight"] = k_weight neuron_state_dict[f"{base_key}.v_proj.weight"] = v_weight - + elif 'mlp.gate_up_proj.weight' in key: # Split gate_up fused weight into separate gate and up weights layer_idx = key.split('.')[1] intermediate_size = tensor.shape[0] // 2 - + gate_weight = tensor[:intermediate_size, :] up_weight = tensor[intermediate_size:, :] - + base_key = f"layers.{layer_idx}.mlp" neuron_state_dict[f"{base_key}.gate_proj.weight"] = gate_weight neuron_state_dict[f"{base_key}.up_proj.weight"] = up_weight - + else: # Copy other weights as-is neuron_state_dict[key] = tensor - + return neuron_state_dict ``` @@ -98,7 +101,7 @@ def load_state_dict(self, state_dict, strict=True): print("🔧 Converting HuggingFace weights to NeuronX format...") state_dict = self.convert_hf_to_neuron_state_dict(state_dict, self.config) print(f"✅ Weight conversion completed. Total keys: {len(state_dict)}") - + return super().load_state_dict(state_dict, strict) def _is_hf_state_dict(self, state_dict): @@ -139,6 +142,7 @@ else: ### 2. Compilation Test The fix was validated by successful compilation: + - **Before**: KeyError during compilation - **After**: Successful compilation with weight conversion message: ``` @@ -163,6 +167,7 @@ The weight splitting preserves the original tensor dimensions and functionality ## Impact This fix enables: + - ✅ Successful Phi3 model compilation for NeuronX - ✅ Automatic weight format conversion - ✅ Compatibility with HuggingFace model checkpoints @@ -187,4 +192,4 @@ This fix enables: 3. **Thorough Testing**: Creating specific debug scripts to verify the fix ensures the solution works correctly. -4. **Framework Integration**: Understanding how the target framework loads and expects weights is crucial for successful model porting. \ No newline at end of file +4. **Framework Integration**: Understanding how the target framework loads and expects weights is crucial for successful model porting. diff --git a/skills/neuron-framework-autoport/references/knowledge_base/layernorm_vs_rmsnorm_analysis.md b/skills/neuron-framework-autoport/references/knowledge_base/layernorm_vs_rmsnorm_analysis.md index 4458c88..655e871 100644 --- a/skills/neuron-framework-autoport/references/knowledge_base/layernorm_vs_rmsnorm_analysis.md +++ b/skills/neuron-framework-autoport/references/knowledge_base/layernorm_vs_rmsnorm_analysis.md @@ -10,6 +10,7 @@ **The Question**: Why does HuggingFace GenericMoE work correctly with RMSNorm for decoder layers while AWS Neuron hardware requires LayerNorm for ALL normalization layers? **The Answer**: The issue stems from a combination of: + 1. **Hardware-specific CustomCall implementation** in Neuron's RMSNorm 2. **Numerical precision differences** in bfloat16 execution 3. **Residual connection interaction** with normalization instability @@ -24,6 +25,7 @@ **File**: `transformers/src/transformers/models/genericmoe/modeling_genericmoe.py` **Decoder Layer Normalization** (lines 595-596): + ```python class GenericmoeDecoderLayer(GradientCheckpointingLayer): def __init__(self, config: GenericmoeConfig, layer_idx: int): @@ -33,6 +35,7 @@ class GenericmoeDecoderLayer(GradientCheckpointingLayer): ``` **Final Model Normalization** (line 657): + ```python class GenericmoeModel(GenericmoePreTrainedModel): def __init__(self, config: GenericmoeConfig): @@ -41,6 +44,7 @@ class GenericmoeModel(GenericmoePreTrainedModel): ``` **GenericmoeRMSNorm Implementation** (lines 567-581): + ```python class GenericmoeRMSNorm(nn.Module): def __init__(self, hidden_size, eps=1e-6): @@ -57,6 +61,7 @@ class GenericmoeRMSNorm(nn.Module): ``` **Key Features**: + - Upcasts to FP32 for normalization computation - Uses pure PyTorch operations (pow, mean, rsqrt) - Runs on mature CUDA kernels (GPU) or optimized CPU ops @@ -69,6 +74,7 @@ class GenericmoeRMSNorm(nn.Module): **File**: `neuronx_distributed_inference/modules/custom_calls.py` **CustomRMSNorm Implementation** (lines 11-38): + ```python class CustomRMSNorm(nn.Module): def __init__(self, hidden_size=None, eps=1e-6): @@ -91,6 +97,7 @@ class CustomRMSNorm(nn.Module): ``` **RmsNorm Hardware Call** (`torch_neuronx/xla_impl/ops.py` lines 1471-1482): + ```python class RmsNorm(torch.autograd.Function): @xla_hlo_call @@ -107,6 +114,7 @@ class RmsNorm(torch.autograd.Function): ``` **Key Features**: + - Calls `AwsNeuronRmsNorm` custom hardware kernel - Implementation is **opaque** (compiled into hardware) - Optimized for Neuron tensor cores @@ -119,6 +127,7 @@ class RmsNorm(torch.autograd.Function): ### 1. Hardware CustomCall Behavioral Differences **Problem**: `AwsNeuronRmsNorm` is a custom hardware kernel with implementation details hidden from the PyTorch layer. While functionally equivalent in theory, subtle differences in: + - Numerical precision handling - Rounding modes - Intermediate computation order @@ -133,6 +142,7 @@ Can cause **activation distribution drift** when compounded across 32 decoder la ### 2. Mean-Centered vs Non-Mean-Centered Normalization **RMSNorm** (no mean subtraction): + ```python rms = sqrt(x.pow(2).mean(-1, keepdim=True) + eps) x_normalized = x / rms @@ -140,6 +150,7 @@ output = weight * x_normalized ``` **LayerNorm** (mean-centered): + ```python mean = x.mean(dim=-1, keepdim=True) var = x.var(dim=-1, keepdim=True, unbiased=False) @@ -148,6 +159,7 @@ output = weight * x_normalized + bias ``` **Key Difference**: LayerNorm **subtracts the mean** before normalizing, which: + - Centers activations around zero - Prevents drift in positive/negative directions - More robust to outliers @@ -160,6 +172,7 @@ output = weight * x_normalized + bias GenericMoE has 16 experts with a router that selects top-2 experts per token. The router uses **softmax** over logits to compute expert probabilities. **RMSNorm Impact**: + ```python # Without mean subtraction, activations can drift hidden_states_norm = hidden_states / rms # Could be biased positive or negative @@ -170,6 +183,7 @@ expert_probs = softmax(router_logits) # Softmax is sensitive to input scale ``` **LayerNorm Impact**: + ```python # Mean subtraction centers activations hidden_states_norm = (hidden_states - mean) / std # Centered around 0 @@ -180,6 +194,7 @@ expert_probs = softmax(router_logits) # Better expert selection ``` **Why This Matters on Neuron**: The `AwsNeuronRmsNorm` custom kernel may have subtle numerical differences that accumulate through: + 1. 32 decoder layers 2. Each with 2 normalization calls (pre-attention, pre-MoE) 3. 64 total normalization operations @@ -192,6 +207,7 @@ Small biases compound → Router makes poor expert choices → Gibberish output ### 4. Residual Connection Instability **Residual Connection Pattern**: + ```python # Pre-attention residual = hidden_states @@ -207,12 +223,14 @@ hidden_states = residual + hidden_states # Residual add ``` **RMSNorm Problem**: + - No mean centering → Activations can have non-zero mean - Residual connections **accumulate bias** across layers - After 32 layers: `hidden_states = initial + Σ(biased_updates)` - Result: Activations grow unbounded or collapse **LayerNorm Solution**: + - Mean subtraction keeps activations centered - Residual updates are zero-mean - Stable across all 32 layers @@ -224,11 +242,13 @@ hidden_states = residual + hidden_states # Residual add **GenericMoE Configuration**: Uses `torch.bfloat16` for computation **bfloat16 Characteristics**: + - 8-bit exponent (same as FP32) → Good dynamic range - 7-bit mantissa (vs 23-bit in FP32) → **Low precision** - Rounding errors accumulate quickly **RMSNorm in bfloat16**: + ```python variance = hidden_states.pow(2).mean(-1, keepdim=True) # Squared values → Large numbers rms = torch.rsqrt(variance + eps) # rsqrt of large numbers → Small numbers @@ -236,6 +256,7 @@ output = hidden_states * rms # Multiplication may lose precision ``` **LayerNorm in bfloat16**: + ```python mean = x.mean(dim=-1, keepdim=True) # Mean is moderate var = x.var(dim=-1, keepdim=True) # Variance more stable than squared mean @@ -243,6 +264,7 @@ output = (x - mean) / sqrt(var + eps) # Better conditioned ``` **Why LayerNorm Works Better**: + - Mean subtraction reduces magnitude before division - Better numerical conditioning in low-precision arithmetic - PyTorch's LayerNorm has optimized bfloat16 kernels @@ -253,19 +275,21 @@ output = (x - mean) / sqrt(var + eps) # Better conditioned **Interesting Observation**: HuggingFace uses **different norms for different layers**: -| Layer | Normalization Type | -|-------|-------------------| -| Decoder `input_layernorm` | GenericmoeRMSNorm | -| Decoder `post_attention_layernorm` | GenericmoeRMSNorm | -| Final `model.norm` | **LayerNorm** ← Different! | +| Layer | Normalization Type | +| ---------------------------------- | -------------------------- | +| Decoder `input_layernorm` | GenericmoeRMSNorm | +| Decoder `post_attention_layernorm` | GenericmoeRMSNorm | +| Final `model.norm` | **LayerNorm** ← Different! | **Why This Works on GPU**: + 1. **Mature CUDA kernels**: RMSNorm has been extensively optimized for GPU 2. **Higher precision**: GPUs often use TF32 (19-bit mantissa) for intermediate computations 3. **Better compiler**: CUDA compiler has years of optimization for transformer ops 4. **Final LayerNorm saves the day**: The final LayerNorm re-centers activations before the LM head, correcting any accumulated bias from the decoder RMSNorms **Why This Fails on Neuron**: + 1. **Custom kernel**: `AwsNeuronRmsNorm` is newer, less mature 2. **Strict bfloat16**: No TF32 fallback on Neuron cores 3. **Compiler limitations**: XLA-to-Neuron compilation may not optimize as aggressively @@ -278,13 +302,16 @@ output = (x - mean) / sqrt(var + eps) # Better conditioned ### v15 Implementation (Failed - Still Gibberish) **Changed**: + - ✅ Final `self.norm`: RMSNorm → LayerNorm **Unchanged**: + - ❌ Decoder `input_layernorm`: Still RMSNorm - ❌ Decoder `post_attention_layernorm`: Still RMSNorm **Result**: + ``` Test 1: "The capital of France is Paris is correct. The capital is capital is capital..." Test 2: Empty output @@ -292,6 +319,7 @@ Test 3: "The fibbyline is fibbyline..." ``` **Analysis**: The final LayerNorm tried to fix the activations, but by that point: + - 32 layers had accumulated bias from RMSNorm - Router had made poor expert selections - Token representations were corrupted @@ -302,11 +330,13 @@ Test 3: "The fibbyline is fibbyline..." ### v16 Implementation (Success - Perfect Output) **Changed**: + - ✅ Decoder `input_layernorm`: RMSNorm → LayerNorm - ✅ Decoder `post_attention_layernorm`: RMSNorm → LayerNorm - ✅ Final `self.norm`: RMSNorm → LayerNorm **Result**: + ``` Test 1: "The capital of France is Paris. It is not only the largest city in France..." Test 2: "A mixture of experts model is an ensemble learning approach..." @@ -314,6 +344,7 @@ Test 3: "Certainly! Below is a Python function that calculates Fibonacci numbers ``` **Analysis**: LayerNorm everywhere: + - Keeps activations centered at every layer - Router receives well-distributed inputs - Expert selection is accurate @@ -325,15 +356,18 @@ Test 3: "Certainly! Below is a Python function that calculates Fibonacci numbers ## Why Does HuggingFace Use RMSNorm? **Historical Context**: RMSNorm was introduced as a simplification of LayerNorm: + - **Fewer operations**: No mean subtraction, no bias - **Faster on GPU**: Fewer memory accesses - **Similar accuracy**: On mature hardware with good kernels **Research Papers**: + - "Root Mean Square Layer Normalization" (Zhang & Sennrich, 2019) - Shows RMSNorm achieves similar results to LayerNorm on GPU **Why It Works for HuggingFace**: + 1. Trained on GPU with mature CUDA kernels 2. Inference on GPU with same kernels 3. Final LayerNorm provides safety net @@ -344,16 +378,19 @@ Test 3: "Certainly! Below is a Python function that calculates Fibonacci numbers ## Why LayerNorm Is Required on Neuron **Hardware Constraints**: + - Newer custom kernel (`AwsNeuronRmsNorm`) with different behavior - Strict bfloat16 arithmetic (no TF32) - XLA-to-Neuron compilation pipeline **Architectural Sensitivity**: + - 32 layers × 2 norms/layer = 64 normalization operations - MoE router sensitive to input distribution - Residual connections accumulate bias **Numerical Stability**: + - Mean subtraction in LayerNorm centers activations - More robust in low-precision arithmetic - Prevents drift across many layers @@ -363,37 +400,45 @@ Test 3: "Certainly! Below is a Python function that calculates Fibonacci numbers ## Key Takeaways ### 1. Hardware Matters + **Same weights, same architecture, different hardware → different behavior** The `AwsNeuronRmsNorm` custom call is not a drop-in replacement for PyTorch's RMSNorm due to: + - Implementation differences - Precision handling - Rounding modes - Optimization trade-offs ### 2. Normalization Is Critical for MoE + **Router stability depends on input distribution** RMSNorm without mean subtraction can cause: + - Activation drift - Biased router logits - Poor expert selection - Catastrophic output degradation ### 3. Residual Connections Amplify Problems + **Bias accumulates across layers** Without mean centering: + - Each layer adds biased updates - Residual connections compound the bias - After 32 layers, activations are corrupted ### 4. Don't Trust Reference Implementations Blindly + **HuggingFace works on GPU ≠ HuggingFace works everywhere** Hardware-specific optimizations require hardware-specific fixes. The successful port knew this and used LayerNorm everywhere. ### 5. Debugging Deep Learning Is Hard + **Small numerical differences → catastrophic failure** It took 16 versions to identify that ALL normalization layers needed the fix, not just the final one. The v15 → v16 jump was the critical insight. diff --git a/skills/neuron-framework-autoport/references/knowledge_base/llama3_neuronx_implementation_summary.md b/skills/neuron-framework-autoport/references/knowledge_base/llama3_neuronx_implementation_summary.md index 7a3ac13..792bee5 100644 --- a/skills/neuron-framework-autoport/references/knowledge_base/llama3_neuronx_implementation_summary.md +++ b/skills/neuron-framework-autoport/references/knowledge_base/llama3_neuronx_implementation_summary.md @@ -7,9 +7,11 @@ This document summarizes the comprehensive analysis and implementation of Meta's ## Project Overview ### Objective + Port Meta's Llama3 model from its original CUDA implementation to the NeuronxDistributed framework for efficient inference on AWS Neuron hardware. ### Scope + - Analyze original Llama3 architecture from `/home/ec2-user/source/llama3` - Study existing NeuronxDistributed model implementations - Create a complete Llama3 implementation following framework patterns @@ -24,6 +26,7 @@ Port Meta's Llama3 model from its original CUDA implementation to the NeuronxDis The original Llama3 implementation revealed the following key architectural components: #### Core Components + - **Model Parameters**: Llama3.2-1B configuration - Hidden size: 2048 - Attention heads: 32 @@ -35,13 +38,15 @@ The original Llama3 implementation revealed the following key architectural comp - Uses scaled RoPE: true #### Key Architectural Features + 1. **Grouped-Query Attention (GQA)**: Multiple query heads share key-value heads for efficiency 2. **Rotary Position Embeddings (RoPE)**: With scaling support for extended context -3. **SwiGLU Activation**: In the MLP layers (gate_proj * silu(up_proj)) +3. **SwiGLU Activation**: In the MLP layers (gate_proj \* silu(up_proj)) 4. **RMSNorm**: For layer normalization 5. **KV Caching**: For efficient autoregressive generation #### Original Implementation Structure + ```python class Transformer(nn.Module): - tok_embeddings: VocabParallelEmbedding @@ -59,6 +64,7 @@ class TransformerBlock(nn.Module): ### NeuronxDistributed Framework Architecture #### Framework Structure + The NeuronxDistributed framework follows a specific pattern for model implementations: ``` @@ -68,17 +74,20 @@ src/neuronx_distributed_inference/models/{model_name}/ ``` #### Base Classes + 1. **NeuronApplicationBase**: Root application class requiring `model_path` parameter 2. **NeuronBaseForCausalLM**: Base class for causal language models 3. **NeuronBaseModel**: Base class for the actual model implementation 4. **NeuronAttentionBase**: Base class for attention mechanisms #### Configuration System + - **InferenceConfig**: Model-specific configuration - **NeuronConfig**: Neuron hardware-specific configuration - **OnDeviceSamplingConfig**: For on-device text generation #### Parallelization Support + - **Tensor Parallelism (TP)**: Distribute parameters across devices - **Sequence Parallelism (SP)**: Distribute sequence processing - **Context Parallelism (CP)**: For long sequence attention @@ -91,6 +100,7 @@ src/neuronx_distributed_inference/models/{model_name}/ Created a comprehensive implementation in `neuronx_llama3/` directory: #### Core Model Classes + 1. **Llama3InferenceConfig**: Configuration class with HF and original format support 2. **NeuronLlama3Attention**: GQA implementation using NeuronAttentionBase 3. **NeuronLlama3MLP**: SwiGLU MLP implementation @@ -99,6 +109,7 @@ Created a comprehensive implementation in `neuronx_llama3/` directory: 6. **NeuronLlama3ForCausalLM**: Causal LM wrapper #### Key Implementation Features + - **Dual Format Support**: Handles both original Llama3 format (`params.json`) and HuggingFace format (`config.json`) - **Proper Parallelization**: Integrates with NeuronxDistributed parallelization strategies - **Optimized Components**: Uses CustomRMSNorm and other Neuron-optimized components @@ -107,9 +118,11 @@ Created a comprehensive implementation in `neuronx_llama3/` directory: ### Checkpoint Conversion System #### Multi-Format Support + The conversion system handles three checkpoint formats: 1. **Original Llama3 Format** (`consolidated.00.pth`): + ``` tok_embeddings.weight → model.embed_tokens.weight layers.{i}.attention.wq.weight → model.layers.{i}.self_attn.q_proj.weight @@ -117,6 +130,7 @@ The conversion system handles three checkpoint formats: ``` 2. **HuggingFace SafeTensors Format**: + ``` model.embed_tokens.weight → model.embed_tokens.weight model.layers.{i}.self_attn.q_proj.weight → model.layers.{i}.self_attn.qkv_proj.q_proj.weight @@ -125,6 +139,7 @@ The conversion system handles three checkpoint formats: 3. **Neuron Format**: Final format expected by the NeuronxDistributed model #### Weight Mapping Challenges + - **State Dict Key Mismatch**: Different frameworks expect different parameter naming conventions - **Tensor Parallelism Metadata**: Need to add rank information for distributed execution - **QKV Projection Structure**: NeuronxDistributed uses grouped QKV projections @@ -132,15 +147,17 @@ The conversion system handles three checkpoint formats: ### Utility Scripts and Tools #### Complete Toolchain + 1. **convert_checkpoint.py**: Multi-format checkpoint conversion 2. **compile_model.py**: Model compilation for Neuron hardware 3. **run_inference.py**: Inference execution 4. **example_chat.py**: Interactive chat interface 5. **test_model.py**: Model testing without compilation 6. **run_pipeline.sh**: Complete pipeline automation -7. **debug_*.py**: Various debugging utilities +7. **debug\_\*.py**: Various debugging utilities #### Configuration Management + - **Minimal Stable Settings**: Following best practices from MODEL_IMPLEMENTATION_GUIDE.md - **Progressive Optimization**: Start simple, add optimizations incrementally - **Comprehensive Error Handling**: Detailed logging and error reporting @@ -148,11 +165,13 @@ The conversion system handles three checkpoint formats: ## Technical Challenges and Solutions ### Challenge 1: State Dictionary Key Mapping + **Problem**: The converted weights had keys that didn't match the model's expected parameter names. **Root Cause**: Different frameworks use different naming conventions for parameters. -**Investigation**: +**Investigation**: + - Original Llama3 uses: `tok_embeddings.weight`, `layers.{i}.attention.wq.weight` - HuggingFace uses: `model.embed_tokens.weight`, `model.layers.{i}.self_attn.q_proj.weight` - NeuronxDistributed expects: `model.layers.{i}.self_attn.qkv_proj.q_proj.weight` @@ -160,6 +179,7 @@ The conversion system handles three checkpoint formats: **Solution Approach**: Created a multi-stage conversion pipeline that handles all three formats. ### Challenge 2: Model Initialization Pattern + **Problem**: The NeuronApplicationBase constructor requires specific parameters that weren't initially understood. **Root Cause**: The base class expects a `model_path` parameter for loading compiled models. @@ -167,15 +187,18 @@ The conversion system handles three checkpoint formats: **Solution**: Updated the `from_config` method to properly initialize with required parameters. ### Challenge 3: Distributed Environment Setup + **Problem**: Models require proper distributed environment initialization before creation. **Solution**: Implemented proper initialization sequence: + ```python dist.init_process_group(backend="xla", init_method="pjrt://", world_size=1, rank=0) nxd.parallel_layers.parallel_state.initialize_model_parallel(tensor_model_parallel_size=1) ``` ### Challenge 4: Missing Dependencies + **Problem**: Many existing models depend on a missing `llama` module that was intentionally removed. **Finding**: This revealed the interconnected nature of the model implementations and the need for careful dependency management. @@ -185,17 +208,20 @@ nxd.parallel_layers.parallel_state.initialize_model_parallel(tensor_model_parall ### NeuronxDistributed Architecture Patterns #### Model Implementation Pattern + 1. **Configuration Class**: Extends `InferenceConfig` with model-specific parameters 2. **Attention Class**: Extends `NeuronAttentionBase` with model-specific attention 3. **Model Class**: Extends `NeuronBaseModel` with layer definitions 4. **CausalLM Class**: Extends `NeuronBaseForCausalLM` as the main interface #### Parallelization Strategy + - **Automatic Sharding**: Framework handles parameter distribution - **Rank Metadata**: Models need rank information for proper distributed execution - **Process Groups**: Different parallelization strategies use different process groups #### Optimization Features + - **Flash Attention**: Optimized attention computation - **Custom Kernels**: Neuron-specific optimized operations - **KV Cache Management**: Efficient autoregressive generation @@ -204,16 +230,19 @@ nxd.parallel_layers.parallel_state.initialize_model_parallel(tensor_model_parall ### Best Practices Identified #### Configuration Management + 1. **Dual Format Support**: Handle both original and HuggingFace formats 2. **Minimal Initial Settings**: Start with simple, stable configurations 3. **Progressive Enhancement**: Add optimizations after basic functionality works #### Weight Conversion + 1. **Multi-Stage Pipeline**: Handle different source formats systematically 2. **Metadata Addition**: Add required rank and parallelization metadata 3. **Validation**: Verify weight shapes and key mappings #### Error Handling + 1. **Comprehensive Logging**: Track all stages of model creation and compilation 2. **Graceful Degradation**: Handle missing optimizations gracefully 3. **Debug Utilities**: Create tools to inspect model state and parameters @@ -223,53 +252,63 @@ nxd.parallel_layers.parallel_state.initialize_model_parallel(tensor_model_parall ### Model Categories in NeuronxDistributed #### Decoder-Only Models (Causal LM) + - **Mistral**: Uses sliding window attention, GQA, SwiGLU - **Qwen3**: Uses Q-K normalization, GQA, SwiGLU - **DeepSeek**: Transformer-based architecture #### Mixture-of-Experts Models + - **Mixtral**: 8 experts, top-k=2 routing, based on Mistral - **Qwen3-MoE**: Multiple experts with normalized routing probabilities #### Encoder-Decoder Models + - **T5**: Text-to-text transformer, independent of other models #### Multi-Modal Models + - **CLIP**: Vision-text dual encoder - **Pixtral**: Multi-modal capabilities ### Common Architectural Patterns #### Attention Mechanisms + 1. **Multi-Head Attention**: Standard transformer attention 2. **Grouped-Query Attention**: Shared key-value heads for efficiency 3. **Sliding Window Attention**: Limited context for efficiency (Mistral) #### Normalization + - **RMSNorm**: Preferred over LayerNorm for performance - **CustomRMSNorm**: Neuron-optimized implementation - **Q-K Normalization**: Applied to query and key vectors (Qwen3) #### Activation Functions -- **SwiGLU**: Standard for modern LLMs (gate_proj * silu(up_proj)) + +- **SwiGLU**: Standard for modern LLMs (gate_proj \* silu(up_proj)) - **GELU**: Traditional activation function - **SiLU**: Sigmoid Linear Unit ## Lessons Learned ### Framework Understanding + 1. **Base Class Requirements**: Understanding constructor parameters is crucial 2. **Distributed Initialization**: Proper sequence is essential for model creation 3. **State Dict Conventions**: Each framework has specific naming expectations 4. **Parallelization Integration**: Models must be designed with distribution in mind ### Implementation Strategy + 1. **Start Simple**: Begin with minimal configurations and build up 2. **Follow Patterns**: Existing models provide excellent templates 3. **Incremental Development**: Add features one at a time 4. **Comprehensive Testing**: Create debug utilities early in the process ### Common Pitfalls + 1. **State Dict Mismatches**: Most common source of loading errors 2. **Missing Initialization**: Distributed environment must be set up first 3. **Parameter Naming**: Inconsistent naming across frameworks causes issues @@ -278,6 +317,7 @@ nxd.parallel_layers.parallel_state.initialize_model_parallel(tensor_model_parall ## Current Status and Next Steps ### Implementation Status + - ✅ **Complete Architecture**: All model components implemented - ✅ **Checkpoint Conversion**: Multi-format support working - ✅ **Configuration System**: Flexible configuration management @@ -286,12 +326,14 @@ nxd.parallel_layers.parallel_state.initialize_model_parallel(tensor_model_parall - ❌ **Successful Compilation**: Blocked by state dict issues ### Immediate Next Steps + 1. **Resolve State Dict Mapping**: Debug exact parameter structure expected 2. **Test Compilation**: Attempt model compilation once weights load correctly 3. **Validate Inference**: Ensure generated text quality matches original 4. **Performance Optimization**: Add Neuron-specific optimizations ### Future Enhancements + 1. **Tensor Parallelism**: Test with multiple devices 2. **Context Parallelism**: Support for very long sequences 3. **On-Device Sampling**: Implement for faster generation @@ -313,11 +355,13 @@ The comprehensive toolchain, documentation, and debugging utilities created duri ## Files Created ### Core Implementation + - `neuronx_llama3/src/neuronx_llama3/modeling_llama3.py`: Complete model implementation - `neuronx_llama3/src/neuronx_llama3/__init__.py`: Module exports - `neuronx_llama3/setup.py`: Package configuration ### Utility Scripts + - `neuronx_llama3/convert_checkpoint.py`: Multi-format checkpoint conversion - `neuronx_llama3/compile_model.py`: Model compilation for Neuron - `neuronx_llama3/run_inference.py`: Inference execution @@ -326,15 +370,17 @@ The comprehensive toolchain, documentation, and debugging utilities created duri - `neuronx_llama3/run_pipeline.sh`: Complete pipeline automation ### Documentation + - `neuronx_llama3/README.md`: Usage instructions and examples - `neuronx_llama3/IMPLEMENTATION_DETAILS.md`: Detailed implementation guide - `docs/model_architectures.md`: Comprehensive model architecture analysis - `docs/llama3_neuronx_implementation_summary.md`: This summary document ### Debug Utilities + - `neuronx_llama3/debug_keys.py`: Parameter structure debugging - `neuronx_llama3/debug_mistral_keys.py`: Mistral model comparison - `neuronx_llama3/debug_qwen_keys.py`: Qwen model comparison - `neuronx_llama3/debug_t5_keys.py`: T5 model comparison -This comprehensive implementation demonstrates the complexity and depth required for successfully porting models to specialized hardware frameworks while maintaining performance and functionality. \ No newline at end of file +This comprehensive implementation demonstrates the complexity and depth required for successfully porting models to specialized hardware frameworks while maintaining performance and functionality. diff --git a/skills/neuron-framework-autoport/references/knowledge_base/longrope_investigation_v13.md b/skills/neuron-framework-autoport/references/knowledge_base/longrope_investigation_v13.md index f8ca2ce..e98b797 100644 --- a/skills/neuron-framework-autoport/references/knowledge_base/longrope_investigation_v13.md +++ b/skills/neuron-framework-autoport/references/knowledge_base/longrope_investigation_v13.md @@ -1,10 +1,13 @@ # LongRoPE Investigation - GenericMoE v13 Analysis ## Date + October 27, 2025 ## Problem Statement + Despite adding `use_scaled_rope=True` in v13, GenericMoE model continues producing gibberish output: + - Test 1: Empty response - Test 2: "Trans Trans" repetition - Test 3: "Writing Writing..." gibberish @@ -12,11 +15,14 @@ Despite adding `use_scaled_rope=True` in v13, GenericMoE model continues produci ## Investigation Process ### 1. Knowledge Base Review + Reviewed `/NeuroborosFoundations/knowledge_base/` for similar issues: + - **MoE_Port_Master_Summary.md**: Generic MoE achieved 100% accuracy through systematic debugging - **Category3_Accuracy_Debugging_Analysis.md**: Key lesson - normalization type mismatch (LayerNorm vs RMSNorm) caused similar gibberish output ### 2. HuggingFace GenericMoE RoPE Implementation + Discovered GenericMoE uses custom `GenericmoeRotaryEmbedding`: ```python @@ -49,12 +55,14 @@ class GenericmoeRotaryEmbedding(nn.Module): ``` **Key Points:** + - Applies `mscale` multiplier to cos/sin embeddings - Uses `short_mscale` for seq_len <= 4096, `long_mscale` for longer sequences - Both mscale values = 1.243163121016122 - Has 64 `short_factor` and 64 `long_factor` arrays for frequency scaling ### 3. NeuronX Framework RoPE Implementation + Examined `NeuronAttentionBase` in `attention_base.py`: ```python @@ -74,12 +82,15 @@ elif use_polar_compatible_rope: ``` **Critical Finding:** + - `use_scaled_rope` flag **ONLY** affects `use_polar_compatible_rope` code path - Default path uses `RotaryEmbedding` module which doesn't support LongRoPE mscale - The simple `RotaryEmbedding(dim, max_position_embeddings, base)` has no scaling logic ### 4. Checked Llama4 Implementation + Verified Llama4 uses identical pattern: + ```python # Line 311 in modeling_llama4_text.py use_scaled_rope=getattr(config, "rope_scaling", None) is not None, @@ -99,6 +110,7 @@ But this works for Llama4 because it may use `use_polar_compatible_rope` mode or ## Evidence ### GenericMoE config.json rope_scaling: + ```json { "rope_scaling": { @@ -113,6 +125,7 @@ But this works for Llama4 because it may use `use_polar_compatible_rope` mode or ``` ### v13 Implementation (INCORRECT): + ```python # NeuronGenericmoeAttention.__init__ rotary_emb = RotaryEmbedding( @@ -131,6 +144,7 @@ super().__init__( ## Conclusion The LongRoPE scaling in GenericMoE requires custom implementation that: + 1. Reads `short_mscale` and `long_mscale` from config 2. Applies appropriate mscale multiplier based on sequence length 3. Uses `short_factor`/`long_factor` arrays for frequency adjustments @@ -157,22 +171,26 @@ super().__init__( ``` **Rationale:** + - Standard RoPE with `rope_theta=10000` and `max_position_embeddings=131072` may work for inference - LongRoPE's mscale=1.243 is a ~24% adjustment - may not be critical for basic functionality - Allows us to test if other issues (normalization, weight loading, etc.) are causing gibberish - Can revisit LongRoPE implementation if basic RoPE works ## Next Steps (v14) + 1. Remove `use_scaled_rope` from both modeling files 2. Recompile model 3. Test inference 4. Compare output quality with/without LongRoPE ## Files Modified in v13 (to be reverted in v14) + - `/home/ec2-user/agents/hariseldon/NeuroborosFoundations/src/amzn/neuron/neuroboros/models/genericmoe/modeling_genericmoe.py` (line 341) - `/home/ec2-user/agents/hariseldon/neuron_port/modeling_genericmoe.py` (line 325) ## Compilation Status + - v13 compiled successfully in 142 seconds - No compilation errors - Issue is runtime accuracy, not compilation diff --git a/skills/neuron-framework-autoport/references/knowledge_base/model_architectures.md b/skills/neuron-framework-autoport/references/knowledge_base/model_architectures.md index 71c8c8e..f4e52cf 100644 --- a/skills/neuron-framework-autoport/references/knowledge_base/model_architectures.md +++ b/skills/neuron-framework-autoport/references/knowledge_base/model_architectures.md @@ -40,7 +40,7 @@ Mistral is a decoder-only transformer model with the following architecture: - **Normalization**: RMSNorm for input and post-attention normalization - **Position Encoding**: Rotary Position Embeddings (RoPE) - **MLP**: Uses SwiGLU activation function (similar to LLaMA) -- **Special Features**: +- **Special Features**: - Sliding window attention mechanism to limit context for efficiency - Optimized for inference with efficient KV cache handling @@ -55,7 +55,7 @@ Mixtral is a Mixture-of-Experts (MoE) model based on the Mistral architecture: - Top-k routing (k=2) - each token is processed by the 2 most relevant experts - Router network determines which experts to use for each token - Expert outputs are weighted and combined -- **Parallelization**: +- **Parallelization**: - Expert parallelism for distributing experts across devices - Token shuffling for load balancing @@ -88,24 +88,28 @@ Qwen3-MoE is the Mixture-of-Experts variant of Qwen3: ### 5. DBRX DBRX is a transformer-based model with: + - Likely based on a decoder-only architecture - Specialized for distributed inference on Neuron hardware ### 6. DeepSeek DeepSeek model support includes: + - Transformer-based architecture - Optimized for Neuron hardware ### 7. T5 T5 is an encoder-decoder model, different from the decoder-only models above: + - **Architecture**: Encoder-decoder transformer - **Special Features**: Supports text-to-text tasks ### 8. CLIP CLIP is a multi-modal model: + - **Architecture**: Dual encoder (vision and text) - **Special Features**: Supports image-text tasks @@ -167,4 +171,4 @@ The MoE architecture in NeuronxDistributed consists of: The NeuronxDistributed framework provides a comprehensive implementation of modern transformer architectures optimized for AWS Neuron hardware. It supports both standard transformer models (Mistral, Qwen3) and Mixture-of-Experts models (Mixtral, Qwen3-MoE), with various parallelization strategies for efficient distributed training and inference. -The architecture emphasizes flexibility, allowing different parallelization strategies to be combined based on the specific requirements of the model and hardware configuration. The implementation includes specialized kernels and optimizations for the Neuron hardware, enabling efficient inference of large language models. \ No newline at end of file +The architecture emphasizes flexibility, allowing different parallelization strategies to be combined based on the specific requirements of the model and hardware configuration. The implementation includes specialized kernels and optimizations for the Neuron hardware, enabling efficient inference of large language models. diff --git a/skills/neuron-framework-autoport/references/knowledge_base/porting_exercise.md b/skills/neuron-framework-autoport/references/knowledge_base/porting_exercise.md index 01a8a3f..9896025 100644 --- a/skills/neuron-framework-autoport/references/knowledge_base/porting_exercise.md +++ b/skills/neuron-framework-autoport/references/knowledge_base/porting_exercise.md @@ -7,9 +7,11 @@ This document captures the comprehensive learnings from successfully porting Met ## Project Scope and Objectives ### Primary Goal + Port Meta's Llama3.2-1B model from the original CUDA implementation (`/home/ec2-user/source/llama3`) to run efficiently on AWS Neuron hardware using the NeuronxDistributed framework. ### Key Requirements + - Maintain architectural fidelity to the original Llama3 implementation - Support both original Llama3 checkpoint format and HuggingFace format - Follow established NeuronxDistributed framework patterns @@ -23,21 +25,23 @@ Port Meta's Llama3.2-1B model from the original CUDA implementation (`/home/ec2- From analyzing `/home/ec2-user/source/llama3/llama/model.py`, the key architectural components identified: #### Core Model Parameters (Llama3.2-1B) + ```json { - "dim": 2048, // Hidden size - "n_heads": 32, // Query attention heads - "n_kv_heads": 8, // Key-value heads (4:1 GQA ratio) - "n_layers": 16, // Transformer layers - "vocab_size": 128256, // Vocabulary size - "ffn_dim_multiplier": 1.5, // MLP dimension multiplier - "rope_theta": 500000.0, // RoPE base frequency - "use_scaled_rope": true, // Scaled RoPE for long context - "norm_eps": 1e-05 // RMSNorm epsilon + "dim": 2048, // Hidden size + "n_heads": 32, // Query attention heads + "n_kv_heads": 8, // Key-value heads (4:1 GQA ratio) + "n_layers": 16, // Transformer layers + "vocab_size": 128256, // Vocabulary size + "ffn_dim_multiplier": 1.5, // MLP dimension multiplier + "rope_theta": 500000.0, // RoPE base frequency + "use_scaled_rope": true, // Scaled RoPE for long context + "norm_eps": 1e-5 // RMSNorm epsilon } ``` #### Key Architectural Features + 1. **Grouped-Query Attention (GQA)**: 32 query heads share 8 key-value heads (4:1 ratio) 2. **Rotary Position Embeddings (RoPE)**: With θ=500,000 and scaling support 3. **SwiGLU Activation**: `w2(silu(w1(x)) * w3(x))` in MLP layers @@ -45,6 +49,7 @@ From analyzing `/home/ec2-user/source/llama3/llama/model.py`, the key architectu 5. **KV Caching**: Explicit key-value caching for autoregressive generation #### Original Implementation Structure + ```python # From source/llama3/llama/model.py class Transformer: @@ -63,17 +68,19 @@ class TransformerBlock: ### NeuronxDistributed Framework Patterns #### Framework Architecture Understanding + The NeuronxDistributed framework follows specific patterns that were crucial to understand: 1. **Base Class Hierarchy**: + ```python NeuronApplicationBase └── NeuronBaseForCausalLM └── NeuronLlama3ForCausalLM - + NeuronBaseModel └── NeuronLlama3Model - + NeuronAttentionBase └── NeuronLlama3Attention ``` @@ -103,7 +110,7 @@ def from_pretrained(cls, model_path: str, **kwargs): # Try original format first params_path = os.path.join(model_path, "params.json") config_path = os.path.join(model_path, "config.json") - + if os.path.exists(params_path): # Load original Llama3 format # Map: dim -> hidden_size, n_heads -> num_attention_heads, etc. @@ -119,6 +126,7 @@ def from_pretrained(cls, model_path: str, **kwargs): **Problem**: Different frameworks use different parameter naming conventions. **Mapping Table**: + ```python # Original Llama3 → NeuronxDistributed "tok_embeddings.weight" → "embed_tokens.weight" @@ -139,16 +147,17 @@ def from_pretrained(cls, model_path: str, **kwargs): **Problem**: NeuronxDistributed requires additional metadata for distributed execution. **Solution**: Added rank information for tensor parallelism: + ```python def convert_hf_to_neuron_state_dict(state_dict, config): tp_degree = config.neuron_config.tp_degree - + # Add rank information for attention layers for i in range(config.num_hidden_layers): state_dict[f"layers.{i}.self_attn.rank_util.rank"] = torch.arange( 0, tp_degree, dtype=torch.int32 ) - + # Add rank information for base model state_dict["rank_util.rank"] = torch.arange(0, tp_degree, dtype=torch.int32) ``` @@ -158,6 +167,7 @@ def convert_hf_to_neuron_state_dict(state_dict, config): **Problem**: Original Llama3 uses complex intermediate size calculation with multipliers. **Original Logic**: + ```python # From source/llama3/llama/model.py FeedForward.__init__ hidden_dim = int(2 * hidden_dim / 3) # Base calculation @@ -173,6 +183,7 @@ hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of) # Ro **Problem**: NeuronApplicationBase requires specific constructor parameters and method implementations. **Key Methods Required**: + - `from_pretrained(model_path)`: Load compiled models - `from_config(config)`: Create models from configuration - `compile(output_path)`: Compile for Neuron hardware @@ -193,7 +204,7 @@ class Llama3InferenceConfig(InferenceConfig): self.output_hidden_states = False self.use_return_dict = True self.use_cache = True - + @classmethod def from_pretrained(cls, model_path, **kwargs): # Support both original and HuggingFace formats @@ -210,7 +221,7 @@ class NeuronLlama3Attention(NeuronAttentionBase): max_position_embeddings=config.max_position_embeddings, base=config.rope_theta, # 500000.0 ) - + super().__init__( config=config, hidden_size=config.hidden_size, # 2048 @@ -256,6 +267,7 @@ neuron_config = NeuronConfig( ### Compilation Success Indicators The successful compilation showed: + - **GQA Handling**: Framework automatically converted GQA to MHA when TP=1 - **Model Loading**: All 164 parameters loaded correctly - **HLO Generation**: Both context encoding and token generation models compiled @@ -316,7 +328,7 @@ Created comprehensive tooling suite: ✅ **Weight Conversion**: Proper parameter name mapping and metadata addition ✅ **Model Compilation**: Successful compilation with GQA handling ✅ **Model Loading**: Compiled artifacts load correctly -✅ **Architecture Integrity**: All Llama3 components properly implemented +✅ **Architecture Integrity**: All Llama3 components properly implemented ### Key Validation Insights @@ -348,6 +360,7 @@ Created comprehensive tooling suite: **Problem**: The configuration loading failed with `FileNotFoundError` when using `~/.llama/checkpoints/Llama3.2-1B` path. **Error Message**: + ``` FileNotFoundError: No configuration file found in ~/.llama/checkpoints/Llama3.2-1B. Expected either params.json or config.json ``` @@ -355,12 +368,13 @@ FileNotFoundError: No configuration file found in ~/.llama/checkpoints/Llama3.2- **Root Cause**: The `~` home directory expansion wasn't being handled in the `from_pretrained` method. **Solution**: Added proper path expansion in the configuration loading: + ```python @classmethod def from_pretrained(cls, model_path: str, **kwargs): # Expand user home directory if needed model_path = os.path.expanduser(model_path) - + params_path = os.path.join(model_path, "params.json") config_path = os.path.join(model_path, "config.json") ``` @@ -372,6 +386,7 @@ def from_pretrained(cls, model_path: str, **kwargs): **Problem**: Weight conversion failed with `AttributeError: 'NoneType' object has no attribute 'vocab_parallel'`. **Error Message**: + ```python File "modeling_llama3.py", line 614, in convert_hf_to_neuron_state_dict if neuron_config.vocab_parallel: @@ -381,6 +396,7 @@ AttributeError: 'NoneType' object has no attribute 'vocab_parallel' **Root Cause**: The `convert_to_neuron_state_dict` method was being called without a proper `neuron_config` in the configuration object. **Solution**: Modified the checkpoint conversion script to create a minimal `NeuronConfig`: + ```python # Create a minimal neuron config for conversion from neuronx_distributed_inference.models.config import NeuronConfig @@ -403,6 +419,7 @@ config = Llama3InferenceConfig.from_pretrained(input_path, neuron_config=neuron_ **Problem**: Model compilation failed with `AttributeError: type object 'NeuronLlama3ForCausalLM' has no attribute 'from_config'`. **Error Message**: + ```python File "compile_model.py", line 139, in compile_model model = NeuronLlama3ForCausalLM.from_config(config) @@ -412,15 +429,16 @@ AttributeError: type object 'NeuronLlama3ForCausalLM' has no attribute 'from_con **Root Cause**: The framework expected a `from_config` class method that wasn't implemented. **Solution**: Added the required class method: + ```python @classmethod def from_config(cls, config): """ Create a model from a configuration. - + Args: config: Model configuration - + Returns: NeuronLlama3ForCausalLM: Model instance """ @@ -434,6 +452,7 @@ def from_config(cls, config): **Problem**: Model creation failed with `TypeError: NeuronApplicationBase.__init__() missing 1 required positional argument: 'model_path'`. **Error Message**: + ```python File "modeling_llama3.py", line 592, in from_config return cls(config=config) @@ -443,6 +462,7 @@ TypeError: NeuronApplicationBase.__init__() missing 1 required positional argume **Root Cause**: The `NeuronApplicationBase` constructor requires a `model_path` parameter, but `from_config` was only passing `config`. **Solution**: Updated the compilation script to use the constructor with `model_path`: + ```python # Create model instance model = NeuronLlama3ForCausalLM(model_path=args.checkpoint_path, config=config) @@ -455,6 +475,7 @@ model = NeuronLlama3ForCausalLM(model_path=args.checkpoint_path, config=config) **Problem**: Compilation failed with `AttributeError: 'NeuronLlama3ForCausalLM' object has no attribute 'compile_model'`. **Error Message**: + ```python File "compile_model.py", line 169, in compile_model model.compile_model() @@ -464,6 +485,7 @@ AttributeError: 'NeuronLlama3ForCausalLM' object has no attribute 'compile_model **Root Cause**: The method name was incorrect - the framework uses `compile()` not `compile_model()`. **Solution**: Fixed the method call and updated the save logic: + ```python # Compile model for Neuron hardware model.compile(args.output_path) @@ -479,6 +501,7 @@ logger.info(f"Compiled model saved to {args.output_path}") **Problem**: Inference failed with `AttributeError: type object 'NeuronLlama3ForCausalLM' has no attribute 'from_pretrained'`. **Error Message**: + ```python File "run_inference.py", line 60, in load_model_and_tokenizer model = NeuronLlama3ForCausalLM.from_pretrained(model_path) @@ -488,16 +511,17 @@ AttributeError: type object 'NeuronLlama3ForCausalLM' has no attribute 'from_pre **Root Cause**: The `from_pretrained` method for loading compiled models wasn't implemented. **Solution**: Added the required class method: + ```python @classmethod def from_pretrained(cls, model_path: str, **kwargs): """ Load a compiled model from a directory. - + Args: model_path: Path to compiled model directory **kwargs: Additional arguments - + Returns: NeuronLlama3ForCausalLM: Loaded model instance """ @@ -511,6 +535,7 @@ def from_pretrained(cls, model_path: str, **kwargs): **Problem**: Forward pass failed with `AttributeError: 'Llama3InferenceConfig' object has no attribute 'output_attentions'`. **Error Message**: + ```python File "model_base.py", line 3290, in _setup_func_config else self.text_config.output_attentions @@ -520,6 +545,7 @@ AttributeError: 'Llama3InferenceConfig' object has no attribute 'output_attentio **Root Cause**: The framework expected certain configuration attributes that weren't defined in the custom config class. **Solution**: Added missing attributes in the `add_derived_config` method: + ```python def add_derived_config(self): # Add missing configuration attributes expected by the framework @@ -540,6 +566,7 @@ def add_derived_config(self): **Problem**: Forward pass failed with `RuntimeError: Forward called before load. Run load() or load_state_dict() making calling forward`. **Error Message**: + ```python File "model_wrapper.py", line 1430, in forward raise RuntimeError("Forward called before load. Run load() or load_state_dict() making calling forward") @@ -549,6 +576,7 @@ RuntimeError: Forward called before load. Run load() or load_state_dict() making **Root Cause**: The compiled model needs to be explicitly loaded before inference can be performed. **Solution**: Added model loading step: + ```python # Load the compiled model model.load() @@ -564,6 +592,7 @@ outputs = model(inputs['input_ids'], position_ids=position_ids) **Problem**: Forward pass failed with `AssertionError: need to call forward with position_ids if attention_mask is not provided`. **Error Message**: + ```python File "model_base.py", line 3311, in _infer_attention_mask position_ids is not None @@ -573,6 +602,7 @@ AssertionError: need to call forward with position_ids if attention_mask is not **Root Cause**: The framework requires either `attention_mask` or `position_ids` to be provided for forward pass. **Solution**: Created position_ids for the input: + ```python # Create position_ids seq_len = inputs['input_ids'].shape[1] @@ -589,6 +619,7 @@ outputs = model(inputs['input_ids'], position_ids=position_ids) **Problem**: Tokenizer loading failed with `HFValidationError: Repo id must be in the form 'repo_name' or 'namespace/repo_name'`. **Error Message**: + ```python huggingface_hub.errors.HFValidationError: Repo id must be in the form 'repo_name' or 'namespace/repo_name': '~/.llama/checkpoints/Llama3.2-1B/hf_converted_new' ``` @@ -596,6 +627,7 @@ huggingface_hub.errors.HFValidationError: Repo id must be in the form 'repo_name **Root Cause**: The HuggingFace tokenizer loading function expected a repository ID, not a local path with `~`. **Solution**: Expanded the path before passing to tokenizer: + ```python tokenizer_path = os.path.expanduser('~/.llama/checkpoints/Llama3.2-1B/hf_converted_new') tokenizer = AutoTokenizer.from_pretrained(tokenizer_path) @@ -608,6 +640,7 @@ tokenizer = AutoTokenizer.from_pretrained(tokenizer_path) **Problem**: Pipeline script failed with `ERROR: file:///home/ec2-user/NeuronxSDK does not appear to be a Python project`. **Error Message**: + ```bash ERROR: file:///home/ec2-user/NeuronxSDK does not appear to be a Python project: neither 'setup.py' nor 'pyproject.toml' found. ``` @@ -677,4 +710,4 @@ This porting exercise successfully demonstrates how to adapt a complex transform This implementation demonstrates the feasibility and methodology for porting complex AI models to specialized hardware frameworks. The comprehensive documentation and tooling created will accelerate future model porting efforts and serve as a reference implementation for the NeuronxDistributed community. -The successful compilation and validation prove that the NeuronxDistributed framework can effectively support modern transformer architectures while providing the hardware optimizations necessary for efficient inference on AWS Neuron hardware. \ No newline at end of file +The successful compilation and validation prove that the NeuronxDistributed framework can effectively support modern transformer architectures while providing the hardware optimizations necessary for efficient inference on AWS Neuron hardware. diff --git a/skills/neuron-framework-autoport/references/knowledge_base/successful_run.md b/skills/neuron-framework-autoport/references/knowledge_base/successful_run.md index 02d70e0..25ac082 100644 --- a/skills/neuron-framework-autoport/references/knowledge_base/successful_run.md +++ b/skills/neuron-framework-autoport/references/knowledge_base/successful_run.md @@ -3,10 +3,6 @@ New Session analyze both of thes... - - - - analyze both of these projects, and particularly the src and test directories as well as the docs directories, the project contains the NeuronSDK source covering architecture and model definitions for a set of models and some implementation guides for models. Understand the architecture of each model, and the various mechanisms used to build each model. including differences in attention, embedding, mlp, encoder, decoder, generation and sampling, and parallelization, and sharding, please describe all of the architectural details. Please give me an architectural description of each of the existing supported models and also identify common traits amongst models including relationships between one model and another. Kiro @@ -73,21 +69,21 @@ The NeuronSDK consists of two main frameworks: NeuronxDistributed: For distributed training NeuronxDistributedInference: For optimized inference on AWS Neuron hardware Supported Model Architectures + 1. Mistral (Decoder-Only Transformer) -Core Architecture: + Core Architecture: Attention: Grouped-Query Attention (GQA) with multiple query heads sharing key-value heads Position Encoding: Rotary Position Embeddings (RoPE) with θ=10,000 Normalization: RMSNorm for both input and post-attention normalization -MLP: SwiGLU activation function (gate_proj * silu(up_proj) -> down_proj) +MLP: SwiGLU activation function (gate_proj \* silu(up_proj) -> down_proj) Special Feature: Sliding window attention mechanism to limit context for efficiency Key Parameters: Hidden size: Variable (typically 4096) Attention heads: Variable with GQA ratio Layers: Variable (typically 32) -Sliding window: 4096 tokens -2. Mixtral (Mixture-of-Experts based on Mistral) +Sliding window: 4096 tokens 2. Mixtral (Mixture-of-Experts based on Mistral) Core Architecture: Base: Inherits all Mistral architecture components @@ -98,16 +94,17 @@ Load Balancing: Token shuffling and capacity factors for even distribution MoE Implementation: # Router determines expert selection + router_logits = linear_router(hidden_states) routing_weights, selected_experts = torch.topk(router_logits, top_k=2, dim=-1) routing_weights = F.softmax(routing_weights, dim=-1) # Expert processing with weighted combination + expert_outputs = [] for expert_id in selected_experts: - expert_output = experts[expert_id](hidden_states) - expert_outputs.append(expert_output * routing_weights[expert_id]) -3. Qwen3 (Advanced Decoder-Only Transformer) +expert_output = experts[expert_id](hidden_states) +expert_outputs.append(expert_output \* routing_weights[expert_id]) 3. Qwen3 (Advanced Decoder-Only Transformer) Core Architecture: Attention: GQA with Q-K normalization (applies RMSNorm to query and key vectors after projection) @@ -118,24 +115,25 @@ Special Feature: Q-K normalization for improved training stability Q-K Normalization: # After QKV projection, normalize Q and K + q_normalized = q_layernorm(query_states) k_normalized = k_layernorm(key_states) + # Then apply RoPE and attention + 4. Qwen3-MoE (Mixture-of-Experts Qwen3) -Core Architecture: + Core Architecture: Base: Inherits Qwen3 architecture with Q-K normalization MoE: Multiple experts per layer with normalized routing probabilities Router: Advanced routing with load balancing mechanisms -Expert MLPs: GLU-based expert networks -5. DBRX (Databricks Transformer) +Expert MLPs: GLU-based expert networks 5. DBRX (Databricks Transformer) Core Architecture: Attention: Standard multi-head attention with custom parameter naming MoE: Mixture-of-Experts with specialized router implementation Normalization: Layer normalization -Parameter Structure: Uses different naming conventions (wte, norm_f, blocks) -6. T5 (Text-to-Text Transfer Transformer) +Parameter Structure: Uses different naming conventions (wte, norm_f, blocks) 6. T5 (Text-to-Text Transfer Transformer) Core Architecture: Type: Encoder-Decoder transformer (unlike others which are decoder-only) @@ -146,24 +144,21 @@ Position Encoding: Relative position embeddings Unique Features: class T5LayerNorm(nn.Module): - def forward(self, hidden_states): - # T5 uses RMSNorm without mean subtraction - variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True) - hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) - return self.weight * hidden_states -7. CLIP (Contrastive Language-Image Pre-training) +def forward(self, hidden*states): # T5 uses RMSNorm without mean subtraction +variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True) +hidden_states = hidden_states * torch.rsqrt(variance + self.variance*epsilon) +return self.weight * hidden_states 7. CLIP (Contrastive Language-Image Pre-training) Core Architecture: Type: Dual encoder (vision and text) Attention: Standard multi-head attention Modality: Multi-modal (text and image processing) -Output: Projection to shared embedding space -8. Llama3 (Meta's Latest Architecture) +Output: Projection to shared embedding space 8. Llama3 (Meta's Latest Architecture) Core Architecture (from source analysis): Attention: Grouped-Query Attention with 4:1 ratio (32 query heads, 8 KV heads) Position Encoding: RoPE with θ=500,000 and scaling support -MLP: SwiGLU activation (w2(silu(w1(x)) * w3(x))) +MLP: SwiGLU activation (w2(silu(w1(x)) \* w3(x))) Normalization: RMSNorm KV Caching: Explicit key-value caching for autoregressive generation Parameters (Llama3.2-1B): @@ -195,7 +190,7 @@ Q-K Normalization: Qwen3-specific enhancement CustomRMSNorm: Neuron-optimized implementation MLP Architectures SwiGLU Activation: Modern standard -Formula: gate_proj * silu(up_proj) -> down_proj +Formula: gate_proj \* silu(up_proj) -> down_proj Used by: Mistral, Mixtral, Qwen3, Llama3 GLU Variants: Different gating mechanisms Dense-Act-Dense: Traditional approach (T5) @@ -270,72 +265,72 @@ Error(s) while reading file(s) Command ls -la /home/ec2-user/.llama/checkpoints/Llama3.2-1B/ total 2415908 -drwxr-xr-x. 3 ec2-user ec2-user 120 Jul 21 21: -45 . drwxr-xr-x. 3 ec2-user ec2-user 25 Jul 14 02: -07 .. -rw-r--r--. 1 ec2-user ec2-user 156 Jul 14 02: -07 checklist.chk -rw-r--r--. 1 ec2-user ec2-user 2471677246 Jul 14 02: -07 consolidated.00.pth drwxrwxr-x. 2 ec2-user ec2-user 16384 Jul 21 21: -45 hf_converted_new -rw-r--r--. 1 ec2-user ec2-user 220 Jul 14 02: -07 params.json -rw-r--r--. 1 ec2-user ec2-user 2183982 Jul 14 02: -07 tokenizer.model +drwxr-xr-x. 3 ec2-user ec2-user 120 Jul 21 21: +45 . drwxr-xr-x. 3 ec2-user ec2-user 25 Jul 14 02: +07 .. -rw-r--r--. 1 ec2-user ec2-user 156 Jul 14 02: +07 checklist.chk -rw-r--r--. 1 ec2-user ec2-user 2471677246 Jul 14 02: +07 consolidated.00.pth drwxrwxr-x. 2 ec2-user ec2-user 16384 Jul 21 21: +45 hf_converted_new -rw-r--r--. 1 ec2-user ec2-user 220 Jul 14 02: +07 params.json -rw-r--r--. 1 ec2-user ec2-user 2183982 Jul 14 02: +07 tokenizer.model Command cat /home/ec2-user/.llama/checkpoints/Llama3.2-1B/params.json { - "dim": 2048, - "ffn_dim_multiplier": 1.5, - "multiple_of": 256, - "n_heads": 32, - "n_kv_heads": 8, - "n_layers": 16, - "norm_eps": 1e-05, - "rope_theta": 500000.0, - "use_scaled_rope": true, - "vocab_size": 128256 +"dim": 2048, +"ffn_dim_multiplier": 1.5, +"multiple_of": 256, +"n_heads": 32, +"n_kv_heads": 8, +"n_layers": 16, +"norm_eps": 1e-05, +"rope_theta": 500000.0, +"use_scaled_rope": true, +"vocab_size": 128256 Command ls -la /home/ec2-user/.llama/checkpoints/Llama3.2-1B/hf_converted_new/ total 5855648 -drwxrwxr-x. 2 ec2-user ec2-user 16384 Jul 21 21: -45 . drwxr-xr-x. 3 ec2-user ec2-user 120 Jul 21 21: -45 .. -rw-rw-r--. 1 ec2-user ec2-user 729 Jul 21 21: -45 config.json -rw-rw-r--. 1 ec2-user ec2-user 137 Jul 21 21: -45 generation_config.json -rw-rw-r--. 1 ec2-user ec2-user 4943274328 Jul 21 21: -45 model-00001-of-00002.safetensors -rw-rw-r--. 1 ec2-user ec2-user 1050673280 Jul 21 21: -45 model-00002-of-00002.safetensors -rw-rw-r--. 1 ec2-user ec2-user 12097 Jul 21 21: -45 model.safetensors.index.json -rw-r--r--. 1 ec2-user ec2-user 2183982 Jul 21 21: -45 tokenizer.model -rw-rw-r--. 1 ec2-user ec2-user 86 Jul 21 21: -45 tokenizer_config.json +drwxrwxr-x. 2 ec2-user ec2-user 16384 Jul 21 21: +45 . drwxr-xr-x. 3 ec2-user ec2-user 120 Jul 21 21: +45 .. -rw-rw-r--. 1 ec2-user ec2-user 729 Jul 21 21: +45 config.json -rw-rw-r--. 1 ec2-user ec2-user 137 Jul 21 21: +45 generation_config.json -rw-rw-r--. 1 ec2-user ec2-user 4943274328 Jul 21 21: +45 model-00001-of-00002.safetensors -rw-rw-r--. 1 ec2-user ec2-user 1050673280 Jul 21 21: +45 model-00002-of-00002.safetensors -rw-rw-r--. 1 ec2-user ec2-user 12097 Jul 21 21: +45 model.safetensors.index.json -rw-r--r--. 1 ec2-user ec2-user 2183982 Jul 21 21: +45 tokenizer.model -rw-rw-r--. 1 ec2-user ec2-user 86 Jul 21 21: +45 tokenizer_config.json Command cat /home/ec2-user/.llama/checkpoints/Llama3.2-1B/hf_converted_new/config.json { - "architectures": [ - "LlamaForCausalLM" - ], - "attention_bias": false, - "attention_dropout": 0.0, - "bos_token_id": 1, - "eos_token_id": 2, - "head_dim": 64, - "hidden_act": "silu", - "hidden_size": 2048, - "initializer_range": 0.02, - "intermediate_size": 8192, - "max_position_embeddings": 2048, - "mlp_bias": false, - "model_type": "llama", - "num_attention_heads": 32, - "num_hidden_layers": 16, - "num_key_value_heads": 8, - "pad_token_id": 0, - "pretraining_tp": 1, - "rms_norm_eps": 1e-05, - "rope_scaling": null, - "rope_theta": 500000.0, - "tie_word_embeddings": false, - "torch_dtype": "float32", - "transformers_version": "4.54.0.dev0", - "use_cache": true, - "use_scaled_rope": true, - "vocab_size": 128256 +"architectures": [ +"LlamaForCausalLM" +], +"attention_bias": false, +"attention_dropout": 0.0, +"bos_token_id": 1, +"eos_token_id": 2, +"head_dim": 64, +"hidden_act": "silu", +"hidden_size": 2048, +"initializer_range": 0.02, +"intermediate_size": 8192, +"max_position_embeddings": 2048, +"mlp_bias": false, +"model_type": "llama", +"num_attention_heads": 32, +"num_hidden_layers": 16, +"num_key_value_heads": 8, +"pad_token_id": 0, +"pretraining_tp": 1, +"rms_norm_eps": 1e-05, +"rope_scaling": null, +"rope_theta": 500000.0, +"tie_word_embeddings": false, +"torch_dtype": "float32", +"transformers_version": "4.54.0.dev0", +"use_cache": true, +"use_scaled_rope": true, +"vocab_size": 128256 } Error(s) while reading file(s) Searched workspace @@ -347,44 +342,36 @@ Now I'll create a comprehensive Llama3 implementation for the NeuronxDistributed Created -__init__.py - +**init**.py Created modeling_llama3.py - Created setup.py - Created convert_checkpoint.py - Created compile_model.py - Created run_inference.py - Created test_model.py - Created run_pipeline.sh - Command chmod +x neuronx_llama3/run_pipeline.sh (aws_neuronx_venv_pytorch_2_7_nxd_inference) (aws_neu @@ -392,12 +379,10 @@ Created README.md - Created IMPLEMENTATION_DETAILS.md - Now let me create a comprehensive summary document: Error(s) while editing @@ -407,7 +392,6 @@ Created IMPLEMENTATION_SUMMARY.md - Summary I have successfully created a comprehensive Llama3 implementation for the NeuronxDistributed framework. Here's what I've delivered: @@ -415,7 +399,7 @@ I have successfully created a comprehensive Llama3 implementation for the Neuron Core Architecture Full Model Implementation: Complete port of Llama3 from /home/ec2-user/source/llama3 maintaining 100% architectural fidelity Grouped-Query Attention: 32 query heads, 8 key-value heads (4:1 GQA ratio) -SwiGLU MLP: Exact w2(silu(w1(x)) * w3(x)) implementation +SwiGLU MLP: Exact w2(silu(w1(x)) _ w3(x)) implementation RoPE with θ=500,000: Maintains original high-frequency RoPE configuration Multi-format Support: Handles original Llama3, HuggingFace, and PyTorch formats Framework Integration @@ -433,23 +417,23 @@ run_pipeline.sh: One-command complete pipeline automation Package Structure neuronx_llama3/ ├── src/neuronx_llama3/ -│ ├── __init__.py # Package exports -│ └── modeling_llama3.py # Complete 1000+ line implementation -├── convert_checkpoint.py # 300+ lines with comprehensive conversion -├── compile_model.py # 250+ lines with distributed setup -├── run_inference.py # 400+ lines with generation logic -├── test_model.py # 500+ lines with multi-level testing -├── run_pipeline.sh # Complete automation pipeline -├── setup.py # Package configuration -├── README.md # Comprehensive user guide -├── IMPLEMENTATION_DETAILS.md # Technical deep-dive -└── IMPLEMENTATION_SUMMARY.md # Executive summary +│ ├── **init**.py # Package exports +│ └── modeling_llama3.py # Complete 1000+ line implementation +├── convert_checkpoint.py # 300+ lines with comprehensive conversion +├── compile_model.py # 250+ lines with distributed setup +├── run_inference.py # 400+ lines with generation logic +├── test_model.py # 500+ lines with multi-level testing +├── run_pipeline.sh # Complete automation pipeline +├── setup.py # Package configuration +├── README.md # Comprehensive user guide +├── IMPLEMENTATION_DETAILS.md # Technical deep-dive +└── IMPLEMENTATION_SUMMARY.md # Executive summary 🎯 Key Technical Achievements Architectural Fidelity Parameter Mapping: Complete 1:1 mapping from original to NeuronX format Forward Pass Logic: Exact replication of original computation graphs Configuration Handling: Supports both params.json and config.json formats -Weight Conversion: Handles consolidated.00.pth, model-*.safetensors, etc. +Weight Conversion: Handles consolidated.00.pth, model-_.safetensors, etc. Framework Compliance Base Class Integration: Follows all established framework patterns Error Handling: Comprehensive validation and graceful error recovery @@ -457,31 +441,39 @@ Documentation: Extensive inline documentation referencing original code Testing: Multi-level validation from config to inference 🚀 Usage Examples Quick Start + # Complete pipeline with one command + cd neuronx_llama3 ./run_pipeline.sh -i /home/ec2-user/.llama/checkpoints/Llama3.2-1B -v Step by Step + # 1. Install package + pip install -e . # 2. Convert checkpoint + python convert_checkpoint.py \ - --input_path /home/ec2-user/.llama/checkpoints/Llama3.2-1B \ - --output_path ./converted + --input_path /home/ec2-user/.llama/checkpoints/Llama3.2-1B \ + --output_path ./converted # 3. Test implementation + python test_model.py --checkpoint_path ./converted # 4. Compile model + python compile_model.py \ - --checkpoint_path ./converted \ - --output_path ./compiled + --checkpoint_path ./converted \ + --output_path ./compiled # 5. Run inference + python run_inference.py \ - --model_path ./compiled \ - --prompt "The meaning of life is" \ - --interactive + --model_path ./compiled \ + --prompt "The meaning of life is" \ + --interactive 📊 Validation Results The implementation includes comprehensive testing that validates: @@ -515,294 +507,297 @@ cd neuronx_llama3 && pip install -e . Command pip install -e . Looking in indexes: https://pypi.org/simple, https:// -pip.repos.neuron.amazonaws.com Obtaining file:///home/ec2-user/NeuronxSDK/neuronx_ll -ama3 Preparing metadata (setup.py) ... done +pip.repos.neuron.amazonaws.com Obtaining file:///home/ec2-user/NeuronxSDK/neuronx_ll +ama3 Preparing metadata (setup.py) ... done Requirement already satisfied: torch>=1.13.0 in /opt/ -aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from neuronx-llama3==1.0.0) (2.7.1) Requirement already satisfied: transformers>=4.20.0 i -n /opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from neuronx-llama3==1.0.0) (4.51.3) Requirement already satisfied: nvidia-curand-cu12==10 -.3.7.77 in /opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from torch>=1.13.0->neuronx-llama3==1.0.0) (10.3.7.77) Requirement already satisfied: nvidia-cusparselt-cu12 -==0.6.3 in /opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from torch>=1.13.0->neuronx-llama3==1.0.0) (0.6.3) Requirement already satisfied: fsspec in /opt/aws_neu -ronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from torch>=1.13.0->neuronx-llama3==1.0.0) (2025.3.0) Requirement already satisfied: nvidia-cuda-runtime-cu -12==12.6.77 in /opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from torch>=1.13.0->neuronx-llama3==1.0.0) (12.6.77) Requirement already satisfied: nvidia-cuda-cupti-cu12 -==12.6.80 in /opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from torch>=1.13.0->neuronx-llama3==1.0.0) (12.6.80) Requirement already satisfied: nvidia-cublas-cu12==12 -.6.4.1 in /opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from torch>=1.13.0->neuronx-llama3==1.0.0) (12.6.4.1) Requirement already satisfied: filelock in /opt/aws_n -euronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from torch>=1.13.0->neuronx-llama3==1.0.0) (3.18.0) Requirement already satisfied: jinja2 in /opt/aws_neu -ronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from torch>=1.13.0->neuronx-llama3==1.0.0) (3.1.6) Requirement already satisfied: triton==3.3.1 in /opt/ -aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from torch>=1.13.0->neuronx-llama3==1.0.0) (3.3.1) Requirement already satisfied: nvidia-nvjitlink-cu12= -=12.6.85 in /opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from torch>=1.13.0->neuronx-llama3==1.0.0) (12.6.85) Requirement already satisfied: nvidia-cusolver-cu12== -11.7.1.2 in /opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from torch>=1.13.0->neuronx-llama3==1.0.0) (11.7.1.2) Requirement already satisfied: nvidia-cufile-cu12==1. -11.1.6 in /opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from torch>=1.13.0->neuronx-llama3==1.0.0) (1.11.1.6) Requirement already satisfied: networkx in /opt/aws_n -euronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from torch>=1.13.0->neuronx-llama3==1.0.0) (2.8.8) Requirement already satisfied: nvidia-cudnn-cu12==9.5 -.1.17 in /opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from torch>=1.13.0->neuronx-llama3==1.0.0) (9.5.1.17) Requirement already satisfied: nvidia-cufft-cu12==11. -3.0.4 in /opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from torch>=1.13.0->neuronx-llama3==1.0.0) (11.3.0.4) Requirement already satisfied: nvidia-cusparse-cu12== -12.5.4.2 in /opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from torch>=1.13.0->neuronx-llama3==1.0.0) (12.5.4.2) Requirement already satisfied: nvidia-nvtx-cu12==12.6 -.77 in /opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from torch>=1.13.0->neuronx-llama3==1.0.0) (12.6.77) Requirement already satisfied: sympy>=1.13.3 in /opt/ -aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from torch>=1.13.0->neuronx-llama3==1.0.0) (1.14.0) Requirement already satisfied: nvidia-nccl-cu12==2.26 -.2 in /opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from torch>=1.13.0->neuronx-llama3==1.0.0) (2.26.2) Requirement already satisfied: nvidia-cuda-nvrtc-cu12 -==12.6.77 in /opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from torch>=1.13.0->neuronx-llama3==1.0.0) (12.6.77) Requirement already satisfied: typing-extensions>=4.1 -0.0 in /opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from torch>=1.13.0->neuronx-llama3==1.0.0) (4.14.0) Requirement already satisfied: setuptools>=40.8.0 in -/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from triton==3.3.1->torch>=1.13.0->neuronx-llama3==1.0.0) (80.9.0) Requirement already satisfied: pyyaml>=5.1 in /opt/aw -s_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from transformers>=4.20.0->neuronx-llama3==1.0.0) (6.0.2) Requirement already satisfied: regex!=2019.12.17 in / -opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from transformers>=4.20.0->neuronx-llama3==1.0.0) (2024.11.6) Requirement already satisfied: packaging>=20.0 in /op -t/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from transformers>=4.20.0->neuronx-llama3==1.0.0) (25.0) Requirement already satisfied: numpy>=1.17 in /opt/aw -s_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from transformers>=4.20.0->neuronx-llama3==1.0.0) (1.26.4) Requirement already satisfied: requests in /opt/aws_n -euronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from transformers>=4.20.0->neuronx-llama3==1.0.0) (2.32.4) Requirement already satisfied: safetensors>=0.4.3 in -/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from transformers>=4.20.0->neuronx-llama3==1.0.0) (0.5.3) Requirement already satisfied: huggingface-hub<1.0,>= -0.30.0 in /opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from transformers>=4.20.0->neuronx-llama3==1.0.0) (0.33.1) Requirement already satisfied: tqdm>=4.27 in /opt/aws -_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from transformers>=4.20.0->neuronx-llama3==1.0.0) (4.67.1) Requirement already satisfied: tokenizers<0.22,>=0.21 - in /opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from transformers>=4.20.0->neuronx-llama3==1.0.0) (0.21.2) Requirement already satisfied: hf-xet<2.0.0,>=1.1.2 i -n /opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from huggingface-hub<1.0,>=0.30.0->transformers>=4.20.0->neuronx-llama3==1.0.0) (1.1.5) Requirement already satisfied: mpmath<1.4,>=1.1.0 in -/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from sympy>=1.13.3->torch>=1.13.0->neuronx-llama3==1.0.0) (1.3.0) Requirement already satisfied: MarkupSafe>=2.0 in /op -t/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from jinja2->torch>=1.13.0->neuronx-llama3==1.0.0) (3.0.2) Requirement already satisfied: idna<4,>=2.5 in /opt/a -ws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from requests->transformers>=4.20.0->neuronx-llama3==1.0.0) (3.10) Requirement already satisfied: charset_normalizer<4,> -=2 in /opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from requests->transformers>=4.20.0->neuronx-llama3==1.0.0) (3.4.2) Requirement already satisfied: certifi>=2017.4.17 in -/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from requests->transformers>=4.20.0->neuronx-llama3==1.0.0) (2025.6.15) Requirement already satisfied: urllib3<3,>=1.21.1 in -/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from requests->transformers>=4.20.0->neuronx-llama3==1.0.0) (2.5.0) Installing collected packages: neuronx-llama3 - Attempting uninstall: neuronx-llama3 - Found existing installation: neuronx_llama3 0.1.0 - Uninstalling neuronx_llama3-0.1.0: - Successfully uninstalled neuronx_llama3-0.1.0 - Running setup.py develop for neuronx-llama3 +aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from neuronx-llama3==1.0.0) (2.7.1) Requirement already satisfied: transformers>=4.20.0 i +n /opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from neuronx-llama3==1.0.0) (4.51.3) Requirement already satisfied: nvidia-curand-cu12==10 +.3.7.77 in /opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from torch>=1.13.0->neuronx-llama3==1.0.0) (10.3.7.77) Requirement already satisfied: nvidia-cusparselt-cu12 +==0.6.3 in /opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from torch>=1.13.0->neuronx-llama3==1.0.0) (0.6.3) Requirement already satisfied: fsspec in /opt/aws_neu +ronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from torch>=1.13.0->neuronx-llama3==1.0.0) (2025.3.0) Requirement already satisfied: nvidia-cuda-runtime-cu +12==12.6.77 in /opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from torch>=1.13.0->neuronx-llama3==1.0.0) (12.6.77) Requirement already satisfied: nvidia-cuda-cupti-cu12 +==12.6.80 in /opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from torch>=1.13.0->neuronx-llama3==1.0.0) (12.6.80) Requirement already satisfied: nvidia-cublas-cu12==12 +.6.4.1 in /opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from torch>=1.13.0->neuronx-llama3==1.0.0) (12.6.4.1) Requirement already satisfied: filelock in /opt/aws_n +euronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from torch>=1.13.0->neuronx-llama3==1.0.0) (3.18.0) Requirement already satisfied: jinja2 in /opt/aws_neu +ronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from torch>=1.13.0->neuronx-llama3==1.0.0) (3.1.6) Requirement already satisfied: triton==3.3.1 in /opt/ +aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from torch>=1.13.0->neuronx-llama3==1.0.0) (3.3.1) Requirement already satisfied: nvidia-nvjitlink-cu12= +=12.6.85 in /opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from torch>=1.13.0->neuronx-llama3==1.0.0) (12.6.85) Requirement already satisfied: nvidia-cusolver-cu12== +11.7.1.2 in /opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from torch>=1.13.0->neuronx-llama3==1.0.0) (11.7.1.2) Requirement already satisfied: nvidia-cufile-cu12==1. +11.1.6 in /opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from torch>=1.13.0->neuronx-llama3==1.0.0) (1.11.1.6) Requirement already satisfied: networkx in /opt/aws_n +euronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from torch>=1.13.0->neuronx-llama3==1.0.0) (2.8.8) Requirement already satisfied: nvidia-cudnn-cu12==9.5 +.1.17 in /opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from torch>=1.13.0->neuronx-llama3==1.0.0) (9.5.1.17) Requirement already satisfied: nvidia-cufft-cu12==11. +3.0.4 in /opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from torch>=1.13.0->neuronx-llama3==1.0.0) (11.3.0.4) Requirement already satisfied: nvidia-cusparse-cu12== +12.5.4.2 in /opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from torch>=1.13.0->neuronx-llama3==1.0.0) (12.5.4.2) Requirement already satisfied: nvidia-nvtx-cu12==12.6 +.77 in /opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from torch>=1.13.0->neuronx-llama3==1.0.0) (12.6.77) Requirement already satisfied: sympy>=1.13.3 in /opt/ +aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from torch>=1.13.0->neuronx-llama3==1.0.0) (1.14.0) Requirement already satisfied: nvidia-nccl-cu12==2.26 +.2 in /opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from torch>=1.13.0->neuronx-llama3==1.0.0) (2.26.2) Requirement already satisfied: nvidia-cuda-nvrtc-cu12 +==12.6.77 in /opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from torch>=1.13.0->neuronx-llama3==1.0.0) (12.6.77) Requirement already satisfied: typing-extensions>=4.1 +0.0 in /opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from torch>=1.13.0->neuronx-llama3==1.0.0) (4.14.0) Requirement already satisfied: setuptools>=40.8.0 in +/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from triton==3.3.1->torch>=1.13.0->neuronx-llama3==1.0.0) (80.9.0) Requirement already satisfied: pyyaml>=5.1 in /opt/aw +s_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from transformers>=4.20.0->neuronx-llama3==1.0.0) (6.0.2) Requirement already satisfied: regex!=2019.12.17 in / +opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from transformers>=4.20.0->neuronx-llama3==1.0.0) (2024.11.6) Requirement already satisfied: packaging>=20.0 in /op +t/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from transformers>=4.20.0->neuronx-llama3==1.0.0) (25.0) Requirement already satisfied: numpy>=1.17 in /opt/aw +s_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from transformers>=4.20.0->neuronx-llama3==1.0.0) (1.26.4) Requirement already satisfied: requests in /opt/aws_n +euronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from transformers>=4.20.0->neuronx-llama3==1.0.0) (2.32.4) Requirement already satisfied: safetensors>=0.4.3 in +/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from transformers>=4.20.0->neuronx-llama3==1.0.0) (0.5.3) Requirement already satisfied: huggingface-hub<1.0,>= +0.30.0 in /opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from transformers>=4.20.0->neuronx-llama3==1.0.0) (0.33.1) Requirement already satisfied: tqdm>=4.27 in /opt/aws +\_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from transformers>=4.20.0->neuronx-llama3==1.0.0) (4.67.1) Requirement already satisfied: tokenizers<0.22,>=0.21 +in /opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from transformers>=4.20.0->neuronx-llama3==1.0.0) (0.21.2) Requirement already satisfied: hf-xet<2.0.0,>=1.1.2 i +n /opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from huggingface-hub<1.0,>=0.30.0->transformers>=4.20.0->neuronx-llama3==1.0.0) (1.1.5) Requirement already satisfied: mpmath<1.4,>=1.1.0 in +/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from sympy>=1.13.3->torch>=1.13.0->neuronx-llama3==1.0.0) (1.3.0) Requirement already satisfied: MarkupSafe>=2.0 in /op +t/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from jinja2->torch>=1.13.0->neuronx-llama3==1.0.0) (3.0.2) Requirement already satisfied: idna<4,>=2.5 in /opt/a +ws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from requests->transformers>=4.20.0->neuronx-llama3==1.0.0) (3.10) Requirement already satisfied: charset_normalizer<4,> +=2 in /opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from requests->transformers>=4.20.0->neuronx-llama3==1.0.0) (3.4.2) Requirement already satisfied: certifi>=2017.4.17 in +/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from requests->transformers>=4.20.0->neuronx-llama3==1.0.0) (2025.6.15) Requirement already satisfied: urllib3<3,>=1.21.1 in +/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages (from requests->transformers>=4.20.0->neuronx-llama3==1.0.0) (2.5.0) Installing collected packages: neuronx-llama3 +Attempting uninstall: neuronx-llama3 +Found existing installation: neuronx_llama3 0.1.0 +Uninstalling neuronx_llama3-0.1.0: +Successfully uninstalled neuronx_llama3-0.1.0 +Running setup.py develop for neuronx-llama3 Successfully installed neuronx-llama3-1.0.0 [notice] A new release of pip is available: 23.0.1 -> - 25.1.1 [notice] To update, run: pip install --upgrade pip +25.1.1 [notice] To update, run: pip install --upgrade pip Great! The package is installed successfully. Now let's proceed with the pipeline. Step 2: Convert Checkpoint Let's convert the original Llama3 checkpoint to NeuronX format: Command -python convert_checkpoint.py --input_path /home/ec2-user/.llama/checkpoints/Llama3.2-1B --output_path ./llama3_neuron_checkpoint --verbose +python convert*checkpoint.py --input_path /home/ec2-user/.llama/checkpoints/Llama3.2-1B --output_path ./llama3_neuron_checkpoint --verbose /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributedInference -/src/neuronx_distributed_inference/modules/attention/utils.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.custom_c -alls import neuron_cumsum /home/ec2-user/NeuronxSDK/NeuronxDistributedInference -/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference -/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference -/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_ -llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_ -llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_ -llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.attention_base import NeuronAttentionBase DEBUG:__main__:Verbose logging enabled -INFO:__main__:Starting Llama3 checkpoint conversion -INFO:__main__:Input path: /home/ec2-user/.llama/check -points/Llama3.2-1B INFO:__main__:Output path: ./llama3_neuron_checkpoint -INFO:__main__:Auto-detected format: original -INFO:__main__:Created NeuronConfig: tp_degree=1, batc -h_size=1, seq_len=128 INFO:__main__:Loading model configuration... -INFO:neuronx_llama3.modeling_llama3:Loading original -Llama3 configuration from /home/ec2-user/.llama/checkpoints/Llama3.2-1B/params.json INFO:neuronx_llama3.modeling_llama3:Calculated interm -ediate_size: 8192 INFO:neuronx_llama3.modeling_llama3:Loaded Llama3 con -figuration: {'hidden_size': 2048, 'num_attention_heads': 32, 'num_hidden_layers': 16, 'num_key_value_heads': 8, 'vocab_size': 128256, 'rope_theta': 500000.0, 'rms_norm_eps': 1e-05, 'use_scaled_rope': True, 'ffn_dim_multiplier': 1.5, 'multiple_of': 256, 'max_seq_len': 2048} INFO:__main__:Loaded configuration: hidden_size=2048, - num_layers=16, num_heads=32, num_kv_heads=8 INFO:__main__:Loading checkpoint weights... -INFO:neuronx_llama3.modeling_llama3:Loading original -Llama3 checkpoint from /home/ec2-user/.llama/checkpoints/Llama3.2-1B/consolidated.00.pth INFO:__main__:Loaded checkpoint with 147 parameters -INFO:__main__: tok_embeddings.weight: torch.Size([12 -8256, 2048]) INFO:__main__:Converting to NeuronX format... +/src/neuronx_distributed_inference/modules/attention/utils.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.custom_c +alls import neuron_cumsum /home/ec2-user/NeuronxSDK/NeuronxDistributedInference +/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference +/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference +/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx* +llama3/modeling*llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx* +llama3/modeling*llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx* +llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.attention_base import NeuronAttentionBase DEBUG:**main**:Verbose logging enabled +INFO:**main**:Starting Llama3 checkpoint conversion +INFO:**main**:Input path: /home/ec2-user/.llama/check +points/Llama3.2-1B INFO:**main**:Output path: ./llama3_neuron_checkpoint +INFO:**main**:Auto-detected format: original +INFO:**main**:Created NeuronConfig: tp_degree=1, batc +h_size=1, seq_len=128 INFO:**main**:Loading model configuration... +INFO:neuronx_llama3.modeling_llama3:Loading original +Llama3 configuration from /home/ec2-user/.llama/checkpoints/Llama3.2-1B/params.json INFO:neuronx_llama3.modeling_llama3:Calculated interm +ediate_size: 8192 INFO:neuronx_llama3.modeling_llama3:Loaded Llama3 con +figuration: {'hidden_size': 2048, 'num_attention_heads': 32, 'num_hidden_layers': 16, 'num_key_value_heads': 8, 'vocab_size': 128256, 'rope_theta': 500000.0, 'rms_norm_eps': 1e-05, 'use_scaled_rope': True, 'ffn_dim_multiplier': 1.5, 'multiple_of': 256, 'max_seq_len': 2048} INFO:**main**:Loaded configuration: hidden_size=2048, +num_layers=16, num_heads=32, num_kv_heads=8 INFO:**main**:Loading checkpoint weights... +INFO:neuronx_llama3.modeling_llama3:Loading original +Llama3 checkpoint from /home/ec2-user/.llama/checkpoints/Llama3.2-1B/consolidated.00.pth INFO:**main**:Loaded checkpoint with 147 parameters +INFO:**main**: tok_embeddings.weight: torch.Size([12 +8256, 2048]) INFO:**main**:Converting to NeuronX format... INFO:neuronx_llama3.modeling_llama3:Converting from o -riginal Llama3 format INFO:neuronx_llama3.modeling_llama3:Converted state d -ict with 164 parameters INFO:__main__:Converted to NeuronX format with 164 pa -rameters INFO:__main__: embed_tokens.weight: torch.Size([1282 -56, 2048]) INFO:__main__: norm.weight: torch.Size([2048]) -INFO:__main__: lm_head.weight: torch.Size([128256, 2 -048]) INFO:__main__:Created output directory: ./llama3_neur -on_checkpoint INFO:__main__:Saving converted checkpoint to ./llama3 -_neuron_checkpoint/pytorch_model.bin INFO:__main__:Successfully saved converted checkpoint -INFO:__main__:Saving configuration to ./llama3_neuron -_checkpoint/config.json INFO:__main__:Successfully saved configuration -INFO:__main__:Copying tokenizer.model to output direc -tory INFO:__main__:Checkpoint conversion completed success -fully! INFO:__main__:Converted checkpoint saved to: ./llama3 -_neuron_checkpoint +riginal Llama3 format INFO:neuronx_llama3.modeling_llama3:Converted state d +ict with 164 parameters INFO:**main**:Converted to NeuronX format with 164 pa +rameters INFO:**main**: embed_tokens.weight: torch.Size([1282 +56, 2048]) INFO:**main**: norm.weight: torch.Size([2048]) +INFO:**main**: lm_head.weight: torch.Size([128256, 2 +048]) INFO:**main**:Created output directory: ./llama3_neur +on_checkpoint INFO:**main**:Saving converted checkpoint to ./llama3 +\_neuron_checkpoint/pytorch_model.bin INFO:**main**:Successfully saved converted checkpoint +INFO:**main**:Saving configuration to ./llama3_neuron +\_checkpoint/config.json INFO:**main**:Successfully saved configuration +INFO:**main**:Copying tokenizer.model to output direc +tory INFO:**main**:Checkpoint conversion completed success +fully! INFO:**main**:Converted checkpoint saved to: ./llama3 +\_neuron_checkpoint ===================================================== -======= CONVERSION SUMMARY +======= CONVERSION SUMMARY ===================================================== -======= Input format: original +======= Input format: original Input path: /home/ec2-user/.llama/checkpoints/Llama3. -2-1B Output path: ./llama3_neuron_checkpoint +2-1B Output path: ./llama3_neuron_checkpoint Model parameters: - - Hidden size: 2048 - - Attention heads: 32 - - Key-value heads: 8 - - Layers: 16 - - Vocabulary size: 128256 - - Intermediate size: 8192 -Neuron configuration: - - Tensor parallelism: 1 - - Batch size: 1 - - Sequence length: 128 - - Torch dtype: torch.float32 -Converted parameters: 164 -===================================================== -======= -Excellent! The checkpoint conversion was successful. Let's verify what was created: + +- Hidden size: 2048 +- Attention heads: 32 +- Key-value heads: 8 +- Layers: 16 +- Vocabulary size: 128256 +- Intermediate size: 8192 + Neuron configuration: +- Tensor parallelism: 1 +- Batch size: 1 +- Sequence length: 128 +- Torch dtype: torch.float32 + Converted parameters: 164 + ===================================================== + ======= + Excellent! The checkpoint conversion was successful. Let's verify what was created: Command ls -la llama3_neuron_checkpoint/ total 2928932 -drwxrwxr-x. 2 ec2-user ec2-user 73 Jul 26 00: -29 . drwxrwxr-x. 4 ec2-user ec2-user 16384 Jul 26 00: -29 .. -rw-rw-r--. 1 ec2-user ec2-user 365 Jul 26 00: -29 config.json -rw-rw-r--. 1 ec2-user ec2-user 2997017843 Jul 26 00: -29 pytorch_model.bin -rw-r--r--. 1 ec2-user ec2-user 2183982 Jul 14 02: -07 tokenizer.model +drwxrwxr-x. 2 ec2-user ec2-user 73 Jul 26 00: +29 . drwxrwxr-x. 4 ec2-user ec2-user 16384 Jul 26 00: +29 .. -rw-rw-r--. 1 ec2-user ec2-user 365 Jul 26 00: +29 config.json -rw-rw-r--. 1 ec2-user ec2-user 2997017843 Jul 26 00: +29 pytorch_model.bin -rw-r--r--. 1 ec2-user ec2-user 2183982 Jul 14 02: +07 tokenizer.model Perfect! Now let's test the model implementation before compilation: Step 3: Test Model Implementation Command -python test_model.py --checkpoint_path ./llama3_neuron_checkpoint --verbose +python test*model.py --checkpoint_path ./llama3_neuron_checkpoint --verbose /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributedInference -/src/neuronx_distributed_inference/modules/attention/utils.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.custom_c -alls import neuron_cumsum /home/ec2-user/NeuronxSDK/NeuronxDistributedInference -/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference -/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference -/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_ -llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_ -llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_ -llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.attention_base import NeuronAttentionBase DEBUG:__main__:Verbose logging enabled -INFO:__main__:Starting Llama3 model testing -INFO:__main__:Checkpoint path: ./llama3_neuron_checkp -oint INFO:__main__:Initializing distributed environment fo -r testing... [2025-07-26 00:30:17.232: I neuronx_distributed/paral -lel_layers/parallel_state.py:628] > initializing tensor model parallel with size 1 [2025-07-26 00:30:17.232: I neuronx_distributed/paral -lel_layers/parallel_state.py:629] > initializing pipeline model parallel with size 1 [2025-07-26 00:30:17.232: I neuronx_distributed/paral -lel_layers/parallel_state.py:630] > initializing context model parallel with size 1 [2025-07-26 00:30:17.232: I neuronx_distributed/paral -lel_layers/parallel_state.py:631] > initializing data parallel with size 1 [2025-07-26 00:30:17.232: I neuronx_distributed/paral -lel_layers/parallel_state.py:632] > initializing world size to 1 2025-07-26 00:30:17.000239: 338094 INFO ||NEURON_CC -_WRAPPER||: Call compiler with cmd: neuronx-cc compile --framework=XLA /tmp/ec2-user/neuroncc_compile_workdir/202e8f5d-05e8-4a62-b200-60531cad6f0c/model.MODULE_13537832482499839909+e30acd3a.hlo_module.pb --output /tmp/ec2-user/neuroncc_compile_workdir/202e8f5d-05e8-4a62-b200-60531cad6f0c/model.MODULE_13537832482499839909+e30acd3a.neff --target=trn1 --verbose=35 .Completed run_backend_driver. +/src/neuronx_distributed_inference/modules/attention/utils.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.custom_c +alls import neuron_cumsum /home/ec2-user/NeuronxSDK/NeuronxDistributedInference +/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference +/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference +/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx* +llama3/modeling*llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx* +llama3/modeling*llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx* +llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.attention_base import NeuronAttentionBase DEBUG:**main**:Verbose logging enabled +INFO:**main**:Starting Llama3 model testing +INFO:**main**:Checkpoint path: ./llama3_neuron_checkp +oint INFO:**main**:Initializing distributed environment fo +r testing... [2025-07-26 00:30:17.232: I neuronx_distributed/paral +lel_layers/parallel_state.py:628] > initializing tensor model parallel with size 1 [2025-07-26 00:30:17.232: I neuronx_distributed/paral +lel_layers/parallel_state.py:629] > initializing pipeline model parallel with size 1 [2025-07-26 00:30:17.232: I neuronx_distributed/paral +lel_layers/parallel_state.py:630] > initializing context model parallel with size 1 [2025-07-26 00:30:17.232: I neuronx_distributed/paral +lel_layers/parallel_state.py:631] > initializing data parallel with size 1 [2025-07-26 00:30:17.232: I neuronx_distributed/paral +lel_layers/parallel_state.py:632] > initializing world size to 1 2025-07-26 00:30:17.000239: 338094 INFO ||NEURON_CC +\_WRAPPER||: Call compiler with cmd: neuronx-cc compile --framework=XLA /tmp/ec2-user/neuroncc_compile_workdir/202e8f5d-05e8-4a62-b200-60531cad6f0c/model.MODULE_13537832482499839909+e30acd3a.hlo_module.pb --output /tmp/ec2-user/neuroncc_compile_workdir/202e8f5d-05e8-4a62-b200-60531cad6f0c/model.MODULE_13537832482499839909+e30acd3a.neff --target=trn1 --verbose=35 .Completed run_backend_driver. Compiler status PASS [2025-07-26 00:30:18.407: I neuronx_distributed/paral lel_layers/parallel_state.py:379] [rank_0_pp-1_tp-1_dp-1_cp-1] Chosen Logic for replica groups ret_logic=, 'Ascending Ring PG Group')> [2025-07-26 00:30:18.408: I neuronx_distributed/paral -lel_layers/parallel_state.py:668] [rank_0_pp-1_tp-1_dp-1_cp-1] tp_groups: replica_groups.tp_groups=[[0]] [2025-07-26 00:30:18.408: I neuronx_distributed/paral -lel_layers/parallel_state.py:669] [rank_0_pp-1_tp-1_dp-1_cp-1] dp_groups: replica_groups.dp_groups=[[0]] [2025-07-26 00:30:18.408: I neuronx_distributed/paral -lel_layers/parallel_state.py:670] [rank_0_pp-1_tp-1_dp-1_cp-1] pp_groups: replica_groups.pp_groups=[[0]] [2025-07-26 00:30:18.408: I neuronx_distributed/paral -lel_layers/parallel_state.py:671] [rank_0_pp-1_tp-1_dp-1_cp-1] cp_groups: replica_groups.cp_groups=[[0]] [2025-07-26 00:30:18.408: I neuronx_distributed/paral -lel_layers/parallel_state.py:672] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_model_groups: replica_groups.ep_model_groups=[[0]] [2025-07-26 00:30:18.408: I neuronx_distributed/paral -lel_layers/parallel_state.py:673] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_data_groups: replica_groups.ep_data_groups=[[0]] INFO:__main__:Distributed environment initialized +lel_layers/parallel_state.py:668] [rank_0_pp-1_tp-1_dp-1_cp-1] tp_groups: replica_groups.tp_groups=[[0]] [2025-07-26 00:30:18.408: I neuronx_distributed/paral +lel_layers/parallel_state.py:669] [rank_0_pp-1_tp-1_dp-1_cp-1] dp_groups: replica_groups.dp_groups=[[0]] [2025-07-26 00:30:18.408: I neuronx_distributed/paral +lel_layers/parallel_state.py:670] [rank_0_pp-1_tp-1_dp-1_cp-1] pp_groups: replica_groups.pp_groups=[[0]] [2025-07-26 00:30:18.408: I neuronx_distributed/paral +lel_layers/parallel_state.py:671] [rank_0_pp-1_tp-1_dp-1_cp-1] cp_groups: replica_groups.cp_groups=[[0]] [2025-07-26 00:30:18.408: I neuronx_distributed/paral +lel_layers/parallel_state.py:672] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_model_groups: replica_groups.ep_model_groups=[[0]] [2025-07-26 00:30:18.408: I neuronx_distributed/paral +lel_layers/parallel_state.py:673] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_data_groups: replica_groups.ep_data_groups=[[0]] INFO:**main**:Distributed environment initialized ===================================================== -======= LLAMA3 MODEL TESTING +======= LLAMA3 MODEL TESTING ===================================================== -======= +======= [TEST] Configuration Check ----------------------------------------- -INFO:__main__:Checking model configuration... -INFO:__main__:Found HuggingFace configuration: ./llam -a3_neuron_checkpoint/config.json INFO:__main__:HuggingFace configuration looks good: { -'hidden_size': 2048, 'num_attention_heads': 32, 'num_hidden_layers': 16, 'num_key_value_heads': 8, 'vocab_size': 128256, 'max_position_embeddings': 2048, 'rope_theta': 500000.0, 'rms_norm_eps': 1e-05, 'hidden_act': 'silu', 'tie_word_embeddings': False, 'intermediate_size': 8192, 'model_type': 'llama3_neuron', 'torch_dtype': 'float32'} Result: PASS - -[TEST] Checkpoint Files Check ----------------------------------------- -INFO:__main__:Checking checkpoint files... -INFO:__main__:Found checkpoint files: ['pytorch_model -.bin'] Result: PASS - -[TEST] Weight Loading Check ----------------------------------------- -INFO:__main__:Checking weight loading... + +--- + +INFO:**main**:Checking model configuration... +INFO:**main**:Found HuggingFace configuration: ./llam +a3_neuron_checkpoint/config.json INFO:**main**:HuggingFace configuration looks good: { +'hidden_size': 2048, 'num_attention_heads': 32, 'num_hidden_layers': 16, 'num_key_value_heads': 8, 'vocab_size': 128256, 'max_position_embeddings': 2048, 'rope_theta': 500000.0, 'rms_norm_eps': 1e-05, 'hidden_act': 'silu', 'tie_word_embeddings': False, 'intermediate_size': 8192, 'model_type': 'llama3_neuron', 'torch_dtype': 'float32'} Result: PASS + +## [TEST] Checkpoint Files Check + +INFO:**main**:Checking checkpoint files... +INFO:**main**:Found checkpoint files: ['pytorch_model +.bin'] Result: PASS + +## [TEST] Weight Loading Check + +INFO:**main**:Checking weight loading... INFO:neuronx_llama3.modeling_llama3:Loading PyTorch c -heckpoint from ./llama3_neuron_checkpoint/pytorch_model.bin INFO:__main__:Successfully loaded checkpoint with 164 - parameters INFO:__main__:Sample keys: ['embed_tokens.weight', 'n -orm.weight', 'lm_head.weight', 'layers.0.self_attn.q_proj.weight', 'layers.0.self_attn.k_proj.weight'] INFO:__main__: embed_tokens.weight: torch.Size([1282 -56, 2048]) (torch.bfloat16) INFO:__main__: norm.weight: torch.Size([2048]) (torc -h.bfloat16) INFO:__main__: lm_head.weight: torch.Size([128256, 2 -048]) (torch.bfloat16) Result: PASS - -[TEST] Model Creation Test ----------------------------------------- -INFO:__main__:Testing model creation... +heckpoint from ./llama3_neuron_checkpoint/pytorch_model.bin INFO:**main**:Successfully loaded checkpoint with 164 +parameters INFO:**main**:Sample keys: ['embed_tokens.weight', 'n +orm.weight', 'lm_head.weight', 'layers.0.self_attn.q_proj.weight', 'layers.0.self_attn.k_proj.weight'] INFO:**main**: embed_tokens.weight: torch.Size([1282 +56, 2048]) (torch.bfloat16) INFO:**main**: norm.weight: torch.Size([2048]) (torc +h.bfloat16) INFO:**main**: lm_head.weight: torch.Size([128256, 2 +048]) (torch.bfloat16) Result: PASS + +## [TEST] Model Creation Test + +INFO:**main**:Testing model creation... INFO:neuronx_llama3.modeling_llama3:Loading HuggingFa -ce configuration from ./llama3_neuron_checkpoint/config.json INFO:neuronx_llama3.modeling_llama3:Loaded Llama3 con -figuration: {'hidden_size': 2048, 'num_attention_heads': 32, 'num_hidden_layers': 16, 'num_key_value_heads': 8, 'vocab_size': 128256, 'max_position_embeddings': 2048, 'rope_theta': 500000.0, 'rms_norm_eps': 1e-05, 'hidden_act': 'silu', 'tie_word_embeddings': False, 'use_scaled_rope': True, 'intermediate_size': 8192} INFO:__main__:Configuration loaded successfully: -INFO:__main__: - Hidden size: 2048 -INFO:__main__: - Attention heads: 32 -INFO:__main__: - Key-value heads: 8 -INFO:__main__: - Layers: 16 -INFO:__main__: - Vocabulary size: 128256 -INFO:__main__: - Intermediate size: 8192 -ERROR:__main__:Error creating model: NeuronApplicatio -nBase.__init__() missing 1 required positional argument: 'model_path' Traceback (most recent call last): - File "/home/ec2-user/NeuronxSDK/neuronx_llama3/test -_model.py", line 287, in test_model_creation model = NeuronLlama3ForCausalLM.from_config(confi -g) File "/home/ec2-user/NeuronxSDK/neuronx_llama3/src/ -neuronx_llama3/modeling_llama3.py", line 742, in from_config return cls(config=config) - File "/home/ec2-user/NeuronxSDK/neuronx_llama3/src/ -neuronx_llama3/modeling_llama3.py", line 726, in __init__ super().__init__(config=config, **kwargs) - File "/home/ec2-user/NeuronxSDK/NeuronxDistributedI -nference/src/neuronx_distributed_inference/models/model_base.py", line 2844, in __init__ super().__init__(*args, **kwargs) -TypeError: NeuronApplicationBase.__init__() missing 1 - required positional argument: 'model_path' Result: FAIL - -[TEST] Forward Pass Test ----------------------------------------- -INFO:__main__:Testing forward pass... +ce configuration from ./llama3_neuron_checkpoint/config.json INFO:neuronx_llama3.modeling_llama3:Loaded Llama3 con +figuration: {'hidden_size': 2048, 'num_attention_heads': 32, 'num_hidden_layers': 16, 'num_key_value_heads': 8, 'vocab_size': 128256, 'max_position_embeddings': 2048, 'rope_theta': 500000.0, 'rms_norm_eps': 1e-05, 'hidden_act': 'silu', 'tie_word_embeddings': False, 'use_scaled_rope': True, 'intermediate_size': 8192} INFO:**main**:Configuration loaded successfully: +INFO:**main**: - Hidden size: 2048 +INFO:**main**: - Attention heads: 32 +INFO:**main**: - Key-value heads: 8 +INFO:**main**: - Layers: 16 +INFO:**main**: - Vocabulary size: 128256 +INFO:**main**: - Intermediate size: 8192 +ERROR:**main**:Error creating model: NeuronApplicatio +nBase.**init**() missing 1 required positional argument: 'model_path' Traceback (most recent call last): +File "/home/ec2-user/NeuronxSDK/neuronx_llama3/test +\_model.py", line 287, in test_model_creation model = NeuronLlama3ForCausalLM.from_config(confi +g) File "/home/ec2-user/NeuronxSDK/neuronx_llama3/src/ +neuronx_llama3/modeling_llama3.py", line 742, in from_config return cls(config=config) +File "/home/ec2-user/NeuronxSDK/neuronx_llama3/src/ +neuronx_llama3/modeling_llama3.py", line 726, in **init** super().**init**(config=config, **kwargs) +File "/home/ec2-user/NeuronxSDK/NeuronxDistributedI +nference/src/neuronx_distributed_inference/models/model_base.py", line 2844, in **init** super().**init**(\*args, **kwargs) +TypeError: NeuronApplicationBase.**init**() missing 1 +required positional argument: 'model_path' Result: FAIL + +## [TEST] Forward Pass Test + +INFO:**main**:Testing forward pass... INFO:neuronx_llama3.modeling_llama3:Loading HuggingFa -ce configuration from ./llama3_neuron_checkpoint/config.json INFO:neuronx_llama3.modeling_llama3:Loaded Llama3 con -figuration: {'hidden_size': 2048, 'num_attention_heads': 32, 'num_hidden_layers': 16, 'num_key_value_heads': 8, 'vocab_size': 128256, 'max_position_embeddings': 2048, 'rope_theta': 500000.0, 'rms_norm_eps': 1e-05, 'hidden_act': 'silu', 'tie_word_embeddings': False, 'use_scaled_rope': True, 'intermediate_size': 8192} ERROR:__main__:Error in forward pass: NeuronApplicati -onBase.__init__() missing 1 required positional argument: 'model_path' Traceback (most recent call last): - File "/home/ec2-user/NeuronxSDK/neuronx_llama3/test -_model.py", line 327, in test_forward_pass model = NeuronLlama3ForCausalLM.from_config(confi -g) File "/home/ec2-user/NeuronxSDK/neuronx_llama3/src/ -neuronx_llama3/modeling_llama3.py", line 742, in from_config return cls(config=config) - File "/home/ec2-user/NeuronxSDK/neuronx_llama3/src/ -neuronx_llama3/modeling_llama3.py", line 726, in __init__ super().__init__(config=config, **kwargs) - File "/home/ec2-user/NeuronxSDK/NeuronxDistributedI -nference/src/neuronx_distributed_inference/models/model_base.py", line 2844, in __init__ super().__init__(*args, **kwargs) -TypeError: NeuronApplicationBase.__init__() missing 1 - required positional argument: 'model_path' Result: FAIL +ce configuration from ./llama3_neuron_checkpoint/config.json INFO:neuronx_llama3.modeling_llama3:Loaded Llama3 con +figuration: {'hidden_size': 2048, 'num_attention_heads': 32, 'num_hidden_layers': 16, 'num_key_value_heads': 8, 'vocab_size': 128256, 'max_position_embeddings': 2048, 'rope_theta': 500000.0, 'rms_norm_eps': 1e-05, 'hidden_act': 'silu', 'tie_word_embeddings': False, 'use_scaled_rope': True, 'intermediate_size': 8192} ERROR:**main**:Error in forward pass: NeuronApplicati +onBase.**init**() missing 1 required positional argument: 'model_path' Traceback (most recent call last): +File "/home/ec2-user/NeuronxSDK/neuronx_llama3/test +\_model.py", line 327, in test_forward_pass model = NeuronLlama3ForCausalLM.from_config(confi +g) File "/home/ec2-user/NeuronxSDK/neuronx_llama3/src/ +neuronx_llama3/modeling_llama3.py", line 742, in from_config return cls(config=config) +File "/home/ec2-user/NeuronxSDK/neuronx_llama3/src/ +neuronx_llama3/modeling_llama3.py", line 726, in **init** super().**init**(config=config, **kwargs) +File "/home/ec2-user/NeuronxSDK/NeuronxDistributedI +nference/src/neuronx_distributed_inference/models/model_base.py", line 2844, in **init** super().**init**(\*args, **kwargs) +TypeError: NeuronApplicationBase.**init**() missing 1 +required positional argument: 'model_path' Result: FAIL ===================================================== -======= TEST SUMMARY +======= TEST SUMMARY ===================================================== -======= Configuration Check: PASS +======= Configuration Check: PASS Checkpoint Files Check: PASS Weight Loading Check: PASS Model Creation Test: FAIL @@ -810,114 +805,115 @@ Forward Pass Test: FAIL Overall: 3/5 tests passed ❌ Some tests failed. Please check the implementation -. +. I see the issue. The from_config method needs to provide a model_path parameter. Let me fix this: Accepted edits to modeling_llama3.py - Now let's test again: Command -python test_model.py --checkpoint_path ./llama3_neuron_checkpoint --batch_size 1 --seq_len 32 +python test*model.py --checkpoint_path ./llama3_neuron_checkpoint --batch_size 1 --seq_len 32 /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributedInference -/src/neuronx_distributed_inference/modules/attention/utils.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.custom_c -alls import neuron_cumsum /home/ec2-user/NeuronxSDK/NeuronxDistributedInference -/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference -/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference -/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_ -llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_ -llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_ -llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.attention_base import NeuronAttentionBase [2025-07-26 00:30:52.924: I neuronx_distributed/paral -lel_layers/parallel_state.py:628] > initializing tensor model parallel with size 1 [2025-07-26 00:30:52.924: I neuronx_distributed/paral -lel_layers/parallel_state.py:629] > initializing pipeline model parallel with size 1 [2025-07-26 00:30:52.925: I neuronx_distributed/paral -lel_layers/parallel_state.py:630] > initializing context model parallel with size 1 [2025-07-26 00:30:52.925: I neuronx_distributed/paral -lel_layers/parallel_state.py:631] > initializing data parallel with size 1 [2025-07-26 00:30:52.925: I neuronx_distributed/paral -lel_layers/parallel_state.py:632] > initializing world size to 1 2025-07-26 00:30:52.000930: 339608 INFO ||NEURON_CC -_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/neuronxcc-2.19.8089.0+8ab9f450/MODULE_13537832482499839909+e30acd3a/model.neff [2025-07-26 00:30:52.971: I neuronx_distributed/paral +/src/neuronx_distributed_inference/modules/attention/utils.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.custom_c +alls import neuron_cumsum /home/ec2-user/NeuronxSDK/NeuronxDistributedInference +/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference +/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference +/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx* +llama3/modeling*llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx* +llama3/modeling*llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx* +llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.attention_base import NeuronAttentionBase [2025-07-26 00:30:52.924: I neuronx_distributed/paral +lel_layers/parallel_state.py:628] > initializing tensor model parallel with size 1 [2025-07-26 00:30:52.924: I neuronx_distributed/paral +lel_layers/parallel_state.py:629] > initializing pipeline model parallel with size 1 [2025-07-26 00:30:52.925: I neuronx_distributed/paral +lel_layers/parallel_state.py:630] > initializing context model parallel with size 1 [2025-07-26 00:30:52.925: I neuronx_distributed/paral +lel_layers/parallel_state.py:631] > initializing data parallel with size 1 [2025-07-26 00:30:52.925: I neuronx_distributed/paral +lel_layers/parallel_state.py:632] > initializing world size to 1 2025-07-26 00:30:52.000930: 339608 INFO ||NEURON_CC +\_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/neuronxcc-2.19.8089.0+8ab9f450/MODULE_13537832482499839909+e30acd3a/model.neff [2025-07-26 00:30:52.971: I neuronx_distributed/paral lel_layers/parallel_state.py:379] [rank_0_pp-1_tp-1_dp-1_cp-1] Chosen Logic for replica groups ret_logic=, 'Ascending Ring PG Group')> [2025-07-26 00:30:52.972: I neuronx_distributed/paral -lel_layers/parallel_state.py:668] [rank_0_pp-1_tp-1_dp-1_cp-1] tp_groups: replica_groups.tp_groups=[[0]] [2025-07-26 00:30:52.972: I neuronx_distributed/paral -lel_layers/parallel_state.py:669] [rank_0_pp-1_tp-1_dp-1_cp-1] dp_groups: replica_groups.dp_groups=[[0]] [2025-07-26 00:30:52.972: I neuronx_distributed/paral -lel_layers/parallel_state.py:670] [rank_0_pp-1_tp-1_dp-1_cp-1] pp_groups: replica_groups.pp_groups=[[0]] [2025-07-26 00:30:52.972: I neuronx_distributed/paral -lel_layers/parallel_state.py:671] [rank_0_pp-1_tp-1_dp-1_cp-1] cp_groups: replica_groups.cp_groups=[[0]] [2025-07-26 00:30:52.972: I neuronx_distributed/paral -lel_layers/parallel_state.py:672] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_model_groups: replica_groups.ep_model_groups=[[0]] [2025-07-26 00:30:52.972: I neuronx_distributed/paral -lel_layers/parallel_state.py:673] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_data_groups: replica_groups.ep_data_groups=[[0]] +lel_layers/parallel_state.py:668] [rank_0_pp-1_tp-1_dp-1_cp-1] tp_groups: replica_groups.tp_groups=[[0]] [2025-07-26 00:30:52.972: I neuronx_distributed/paral +lel_layers/parallel_state.py:669] [rank_0_pp-1_tp-1_dp-1_cp-1] dp_groups: replica_groups.dp_groups=[[0]] [2025-07-26 00:30:52.972: I neuronx_distributed/paral +lel_layers/parallel_state.py:670] [rank_0_pp-1_tp-1_dp-1_cp-1] pp_groups: replica_groups.pp_groups=[[0]] [2025-07-26 00:30:52.972: I neuronx_distributed/paral +lel_layers/parallel_state.py:671] [rank_0_pp-1_tp-1_dp-1_cp-1] cp_groups: replica_groups.cp_groups=[[0]] [2025-07-26 00:30:52.972: I neuronx_distributed/paral +lel_layers/parallel_state.py:672] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_model_groups: replica_groups.ep_model_groups=[[0]] [2025-07-26 00:30:52.972: I neuronx_distributed/paral +lel_layers/parallel_state.py:673] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_data_groups: replica_groups.ep_data_groups=[[0]] ===================================================== -======= LLAMA3 MODEL TESTING +======= LLAMA3 MODEL TESTING ===================================================== -======= +======= [TEST] Configuration Check ----------------------------------------- + +--- + Result: PASS -[TEST] Checkpoint Files Check ----------------------------------------- +## [TEST] Checkpoint Files Check + Result: PASS -[TEST] Weight Loading Check ----------------------------------------- +## [TEST] Weight Loading Check + Result: PASS -[TEST] Model Creation Test ----------------------------------------- -ERROR:__main__:Error creating model: NeuronApplicatio -nBase.__init__() missing 1 required positional argument: 'model_path' Traceback (most recent call last): - File "/home/ec2-user/NeuronxSDK/neuronx_llama3/test -_model.py", line 287, in test_model_creation model = NeuronLlama3ForCausalLM.from_config(confi -g) File "/home/ec2-user/NeuronxSDK/neuronx_llama3/src/ -neuronx_llama3/modeling_llama3.py", line 744, in from_config return cls(model_path="", config=config) - File "/home/ec2-user/NeuronxSDK/neuronx_llama3/src/ -neuronx_llama3/modeling_llama3.py", line 726, in __init__ super().__init__(config=config, **kwargs) - File "/home/ec2-user/NeuronxSDK/NeuronxDistributedI -nference/src/neuronx_distributed_inference/models/model_base.py", line 2844, in __init__ super().__init__(*args, **kwargs) -TypeError: NeuronApplicationBase.__init__() missing 1 - required positional argument: 'model_path' Result: FAIL - -[TEST] Forward Pass Test ----------------------------------------- -ERROR:__main__:Error in forward pass: NeuronApplicati -onBase.__init__() missing 1 required positional argument: 'model_path' Traceback (most recent call last): - File "/home/ec2-user/NeuronxSDK/neuronx_llama3/test -_model.py", line 327, in test_forward_pass model = NeuronLlama3ForCausalLM.from_config(confi -g) File "/home/ec2-user/NeuronxSDK/neuronx_llama3/src/ -neuronx_llama3/modeling_llama3.py", line 744, in from_config return cls(model_path="", config=config) - File "/home/ec2-user/NeuronxSDK/neuronx_llama3/src/ -neuronx_llama3/modeling_llama3.py", line 726, in __init__ super().__init__(config=config, **kwargs) - File "/home/ec2-user/NeuronxSDK/NeuronxDistributedI -nference/src/neuronx_distributed_inference/models/model_base.py", line 2844, in __init__ super().__init__(*args, **kwargs) -TypeError: NeuronApplicationBase.__init__() missing 1 - required positional argument: 'model_path' Result: FAIL +## [TEST] Model Creation Test + +ERROR:**main**:Error creating model: NeuronApplicatio +nBase.**init**() missing 1 required positional argument: 'model_path' Traceback (most recent call last): +File "/home/ec2-user/NeuronxSDK/neuronx_llama3/test +\_model.py", line 287, in test_model_creation model = NeuronLlama3ForCausalLM.from_config(confi +g) File "/home/ec2-user/NeuronxSDK/neuronx_llama3/src/ +neuronx_llama3/modeling_llama3.py", line 744, in from_config return cls(model_path="", config=config) +File "/home/ec2-user/NeuronxSDK/neuronx_llama3/src/ +neuronx_llama3/modeling_llama3.py", line 726, in **init** super().**init**(config=config, **kwargs) +File "/home/ec2-user/NeuronxSDK/NeuronxDistributedI +nference/src/neuronx_distributed_inference/models/model_base.py", line 2844, in **init** super().**init**(\*args, **kwargs) +TypeError: NeuronApplicationBase.**init**() missing 1 +required positional argument: 'model_path' Result: FAIL + +## [TEST] Forward Pass Test + +ERROR:**main**:Error in forward pass: NeuronApplicati +onBase.**init**() missing 1 required positional argument: 'model_path' Traceback (most recent call last): +File "/home/ec2-user/NeuronxSDK/neuronx_llama3/test +\_model.py", line 327, in test_forward_pass model = NeuronLlama3ForCausalLM.from_config(confi +g) File "/home/ec2-user/NeuronxSDK/neuronx_llama3/src/ +neuronx_llama3/modeling_llama3.py", line 744, in from_config return cls(model_path="", config=config) +File "/home/ec2-user/NeuronxSDK/neuronx_llama3/src/ +neuronx_llama3/modeling_llama3.py", line 726, in **init** super().**init**(config=config, **kwargs) +File "/home/ec2-user/NeuronxSDK/NeuronxDistributedI +nference/src/neuronx_distributed_inference/models/model_base.py", line 2844, in **init** super().**init**(\*args, **kwargs) +TypeError: NeuronApplicationBase.**init**() missing 1 +required positional argument: 'model_path' Result: FAIL ===================================================== -======= TEST SUMMARY +======= TEST SUMMARY ===================================================== -======= Configuration Check: PASS +======= Configuration Check: PASS Checkpoint Files Check: PASS Weight Loading Check: PASS Model Creation Test: FAIL @@ -925,113 +921,113 @@ Forward Pass Test: FAIL Overall: 3/5 tests passed ❌ Some tests failed. Please check the implementation -. +. I need to fix the constructor to properly handle the model_path parameter. Let me check how other models handle this: - 2 edits to file modeling_llama3.py - Now let's test again: Command -python test_model.py --checkpoint_path ./llama3_neuron_checkpoint --batch_size 1 --seq_len 32 +python test*model.py --checkpoint_path ./llama3_neuron_checkpoint --batch_size 1 --seq_len 32 /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributedInference -/src/neuronx_distributed_inference/modules/attention/utils.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.custom_c -alls import neuron_cumsum /home/ec2-user/NeuronxSDK/NeuronxDistributedInference -/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference -/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference -/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_ -llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_ -llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_ -llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.attention_base import NeuronAttentionBase [2025-07-26 00:31:39.393: I neuronx_distributed/paral -lel_layers/parallel_state.py:628] > initializing tensor model parallel with size 1 [2025-07-26 00:31:39.393: I neuronx_distributed/paral -lel_layers/parallel_state.py:629] > initializing pipeline model parallel with size 1 [2025-07-26 00:31:39.393: I neuronx_distributed/paral -lel_layers/parallel_state.py:630] > initializing context model parallel with size 1 [2025-07-26 00:31:39.393: I neuronx_distributed/paral -lel_layers/parallel_state.py:631] > initializing data parallel with size 1 [2025-07-26 00:31:39.393: I neuronx_distributed/paral -lel_layers/parallel_state.py:632] > initializing world size to 1 2025-07-26 00:31:39.000399: 341076 INFO ||NEURON_CC -_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/neuronxcc-2.19.8089.0+8ab9f450/MODULE_13537832482499839909+e30acd3a/model.neff [2025-07-26 00:31:39.439: I neuronx_distributed/paral +/src/neuronx_distributed_inference/modules/attention/utils.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.custom_c +alls import neuron_cumsum /home/ec2-user/NeuronxSDK/NeuronxDistributedInference +/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference +/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference +/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx* +llama3/modeling*llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx* +llama3/modeling*llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx* +llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.attention_base import NeuronAttentionBase [2025-07-26 00:31:39.393: I neuronx_distributed/paral +lel_layers/parallel_state.py:628] > initializing tensor model parallel with size 1 [2025-07-26 00:31:39.393: I neuronx_distributed/paral +lel_layers/parallel_state.py:629] > initializing pipeline model parallel with size 1 [2025-07-26 00:31:39.393: I neuronx_distributed/paral +lel_layers/parallel_state.py:630] > initializing context model parallel with size 1 [2025-07-26 00:31:39.393: I neuronx_distributed/paral +lel_layers/parallel_state.py:631] > initializing data parallel with size 1 [2025-07-26 00:31:39.393: I neuronx_distributed/paral +lel_layers/parallel_state.py:632] > initializing world size to 1 2025-07-26 00:31:39.000399: 341076 INFO ||NEURON_CC +\_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/neuronxcc-2.19.8089.0+8ab9f450/MODULE_13537832482499839909+e30acd3a/model.neff [2025-07-26 00:31:39.439: I neuronx_distributed/paral lel_layers/parallel_state.py:379] [rank_0_pp-1_tp-1_dp-1_cp-1] Chosen Logic for replica groups ret_logic=, 'Ascending Ring PG Group')> [2025-07-26 00:31:39.439: I neuronx_distributed/paral -lel_layers/parallel_state.py:668] [rank_0_pp-1_tp-1_dp-1_cp-1] tp_groups: replica_groups.tp_groups=[[0]] [2025-07-26 00:31:39.439: I neuronx_distributed/paral -lel_layers/parallel_state.py:669] [rank_0_pp-1_tp-1_dp-1_cp-1] dp_groups: replica_groups.dp_groups=[[0]] [2025-07-26 00:31:39.439: I neuronx_distributed/paral -lel_layers/parallel_state.py:670] [rank_0_pp-1_tp-1_dp-1_cp-1] pp_groups: replica_groups.pp_groups=[[0]] [2025-07-26 00:31:39.439: I neuronx_distributed/paral -lel_layers/parallel_state.py:671] [rank_0_pp-1_tp-1_dp-1_cp-1] cp_groups: replica_groups.cp_groups=[[0]] [2025-07-26 00:31:39.439: I neuronx_distributed/paral -lel_layers/parallel_state.py:672] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_model_groups: replica_groups.ep_model_groups=[[0]] [2025-07-26 00:31:39.439: I neuronx_distributed/paral -lel_layers/parallel_state.py:673] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_data_groups: replica_groups.ep_data_groups=[[0]] +lel_layers/parallel_state.py:668] [rank_0_pp-1_tp-1_dp-1_cp-1] tp_groups: replica_groups.tp_groups=[[0]] [2025-07-26 00:31:39.439: I neuronx_distributed/paral +lel_layers/parallel_state.py:669] [rank_0_pp-1_tp-1_dp-1_cp-1] dp_groups: replica_groups.dp_groups=[[0]] [2025-07-26 00:31:39.439: I neuronx_distributed/paral +lel_layers/parallel_state.py:670] [rank_0_pp-1_tp-1_dp-1_cp-1] pp_groups: replica_groups.pp_groups=[[0]] [2025-07-26 00:31:39.439: I neuronx_distributed/paral +lel_layers/parallel_state.py:671] [rank_0_pp-1_tp-1_dp-1_cp-1] cp_groups: replica_groups.cp_groups=[[0]] [2025-07-26 00:31:39.439: I neuronx_distributed/paral +lel_layers/parallel_state.py:672] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_model_groups: replica_groups.ep_model_groups=[[0]] [2025-07-26 00:31:39.439: I neuronx_distributed/paral +lel_layers/parallel_state.py:673] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_data_groups: replica_groups.ep_data_groups=[[0]] ===================================================== -======= LLAMA3 MODEL TESTING +======= LLAMA3 MODEL TESTING ===================================================== -======= +======= [TEST] Configuration Check ----------------------------------------- + +--- + Result: PASS -[TEST] Checkpoint Files Check ----------------------------------------- +## [TEST] Checkpoint Files Check + Result: PASS -[TEST] Weight Loading Check ----------------------------------------- +## [TEST] Weight Loading Check + Result: PASS -[TEST] Model Creation Test ----------------------------------------- +## [TEST] Model Creation Test + Result: PASS -[TEST] Forward Pass Test ----------------------------------------- -WARNING:__main__:Unexpected keys: 164 -ERROR:__main__:Error in forward pass: Forward called -before load. Run load() or load_state_dict() making calling forward Traceback (most recent call last): - File "/home/ec2-user/NeuronxSDK/neuronx_llama3/test -_model.py", line 356, in test_forward_pass outputs = model(input_ids, position_ids=position_ -ids) File "/opt/aws_neuronx_venv_pytorch_2_7_nxd_inferen -ce/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl return self._call_impl(*args, **kwargs) - File "/opt/aws_neuronx_venv_pytorch_2_7_nxd_inferen -ce/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl return forward_call(*args, **kwargs) - File "/home/ec2-user/NeuronxSDK/NeuronxDistributedI -nference/src/neuronx_distributed_inference/models/model_base.py", line 3248, in forward outputs, is_run_on_neuron = self._get_model_outpu -ts( File "/home/ec2-user/NeuronxSDK/NeuronxDistributedI -nference/src/neuronx_distributed_inference/models/model_base.py", line 3520, in _get_model_outputs outputs = self.context_encoding_model( - File "/opt/aws_neuronx_venv_pytorch_2_7_nxd_inferen -ce/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl return self._call_impl(*args, **kwargs) - File "/opt/aws_neuronx_venv_pytorch_2_7_nxd_inferen -ce/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1762, in _call_impl return forward_call(*args, **kwargs) - File "/home/ec2-user/NeuronxSDK/NeuronxDistributedI -nference/src/neuronx_distributed_inference/models/model_wrapper.py", line 1430, in forward raise RuntimeError( -RuntimeError: Forward called before load. Run load() -or load_state_dict() making calling forward Result: FAIL +## [TEST] Forward Pass Test + +WARNING:**main**:Unexpected keys: 164 +ERROR:**main**:Error in forward pass: Forward called +before load. Run load() or load*state_dict() making calling forward Traceback (most recent call last): +File "/home/ec2-user/NeuronxSDK/neuronx_llama3/test +\_model.py", line 356, in test_forward_pass outputs = model(input_ids, position_ids=position* +ids) File "/opt/aws_neuronx_venv_pytorch_2_7_nxd_inferen +ce/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1751, in \_wrapped_call_impl return self.\_call_impl(*args, \*\*kwargs) +File "/opt/aws_neuronx_venv_pytorch_2_7_nxd_inferen +ce/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1762, in \_call_impl return forward_call(*args, **kwargs) +File "/home/ec2-user/NeuronxSDK/NeuronxDistributedI +nference/src/neuronx_distributed_inference/models/model_base.py", line 3248, in forward outputs, is_run_on_neuron = self.\_get_model_outpu +ts( File "/home/ec2-user/NeuronxSDK/NeuronxDistributedI +nference/src/neuronx_distributed_inference/models/model_base.py", line 3520, in \_get_model_outputs outputs = self.context_encoding_model( +File "/opt/aws_neuronx_venv_pytorch_2_7_nxd_inferen +ce/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1751, in \_wrapped_call_impl return self.\_call_impl(\*args, **kwargs) +File "/opt/aws_neuronx_venv_pytorch_2_7_nxd_inferen +ce/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1762, in \_call_impl return forward_call(\*args, \*\*kwargs) +File "/home/ec2-user/NeuronxSDK/NeuronxDistributedI +nference/src/neuronx_distributed_inference/models/model_wrapper.py", line 1430, in forward raise RuntimeError( +RuntimeError: Forward called before load. Run load() +or load_state_dict() making calling forward Result: FAIL ===================================================== -======= TEST SUMMARY +======= TEST SUMMARY ===================================================== -======= Configuration Check: PASS +======= Configuration Check: PASS Checkpoint Files Check: PASS Weight Loading Check: PASS Model Creation Test: PASS @@ -1039,661 +1035,652 @@ Forward Pass Test: FAIL Overall: 4/5 tests passed ❌ Some tests failed. Please check the implementation -. +. Good progress! The model creation is now working. The forward pass test is failing because it's trying to use the compiled model interface. For testing purposes, let me modify the test to use the underlying model directly. But first, let's proceed with compilation since the basic model creation is working. Step 4: Compile Model Command -python compile_model.py --checkpoint_path ./llama3_neuron_checkpoint --output_path ./llama3_compiled --verbose +python compile*model.py --checkpoint_path ./llama3_neuron_checkpoint --output_path ./llama3_compiled --verbose /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributedInference -/src/neuronx_distributed_inference/modules/attention/utils.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.custom_c -alls import neuron_cumsum /home/ec2-user/NeuronxSDK/NeuronxDistributedInference -/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference -/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference -/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_ -llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_ -llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_ -llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.attention_base import NeuronAttentionBase DEBUG:__main__:Verbose logging enabled -INFO:__main__:Starting Llama3 model compilation -INFO:__main__:Checkpoint path: ./llama3_neuron_checkp -oint INFO:__main__:Output path: ./llama3_compiled -INFO:__main__:Initializing distributed environment... -INFO:__main__:Initializing distributed process group -INFO:__main__:Initializing model parallel with tp_deg -ree=1 [2025-07-26 00:32:07.000: I neuronx_distributed/paral -lel_layers/parallel_state.py:628] > initializing tensor model parallel with size 1 [2025-07-26 00:32:07.000: I neuronx_distributed/paral -lel_layers/parallel_state.py:629] > initializing pipeline model parallel with size 1 [2025-07-26 00:32:07.000: I neuronx_distributed/paral -lel_layers/parallel_state.py:630] > initializing context model parallel with size 1 [2025-07-26 00:32:07.000: I neuronx_distributed/paral -lel_layers/parallel_state.py:631] > initializing data parallel with size 1 [2025-07-26 00:32:07.001: I neuronx_distributed/paral -lel_layers/parallel_state.py:632] > initializing world size to 1 2025-07-26 00:32:07.000006: 342140 INFO ||NEURON_CC -_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/neuronxcc-2.19.8089.0+8ab9f450/MODULE_13537832482499839909+e30acd3a/model.neff [2025-07-26 00:32:07.046: I neuronx_distributed/paral +/src/neuronx_distributed_inference/modules/attention/utils.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.custom_c +alls import neuron_cumsum /home/ec2-user/NeuronxSDK/NeuronxDistributedInference +/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference +/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference +/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx* +llama3/modeling*llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx* +llama3/modeling*llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx* +llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.attention_base import NeuronAttentionBase DEBUG:**main**:Verbose logging enabled +INFO:**main**:Starting Llama3 model compilation +INFO:**main**:Checkpoint path: ./llama3_neuron_checkp +oint INFO:**main**:Output path: ./llama3_compiled +INFO:**main**:Initializing distributed environment... +INFO:**main**:Initializing distributed process group +INFO:**main**:Initializing model parallel with tp_deg +ree=1 [2025-07-26 00:32:07.000: I neuronx_distributed/paral +lel_layers/parallel_state.py:628] > initializing tensor model parallel with size 1 [2025-07-26 00:32:07.000: I neuronx_distributed/paral +lel_layers/parallel_state.py:629] > initializing pipeline model parallel with size 1 [2025-07-26 00:32:07.000: I neuronx_distributed/paral +lel_layers/parallel_state.py:630] > initializing context model parallel with size 1 [2025-07-26 00:32:07.000: I neuronx_distributed/paral +lel_layers/parallel_state.py:631] > initializing data parallel with size 1 [2025-07-26 00:32:07.001: I neuronx_distributed/paral +lel_layers/parallel_state.py:632] > initializing world size to 1 2025-07-26 00:32:07.000006: 342140 INFO ||NEURON_CC +\_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/neuronxcc-2.19.8089.0+8ab9f450/MODULE_13537832482499839909+e30acd3a/model.neff [2025-07-26 00:32:07.046: I neuronx_distributed/paral lel_layers/parallel_state.py:379] [rank_0_pp-1_tp-1_dp-1_cp-1] Chosen Logic for replica groups ret_logic=, 'Ascending Ring PG Group')> [2025-07-26 00:32:07.047: I neuronx_distributed/paral -lel_layers/parallel_state.py:668] [rank_0_pp-1_tp-1_dp-1_cp-1] tp_groups: replica_groups.tp_groups=[[0]] [2025-07-26 00:32:07.047: I neuronx_distributed/paral -lel_layers/parallel_state.py:669] [rank_0_pp-1_tp-1_dp-1_cp-1] dp_groups: replica_groups.dp_groups=[[0]] [2025-07-26 00:32:07.047: I neuronx_distributed/paral -lel_layers/parallel_state.py:670] [rank_0_pp-1_tp-1_dp-1_cp-1] pp_groups: replica_groups.pp_groups=[[0]] [2025-07-26 00:32:07.047: I neuronx_distributed/paral -lel_layers/parallel_state.py:671] [rank_0_pp-1_tp-1_dp-1_cp-1] cp_groups: replica_groups.cp_groups=[[0]] [2025-07-26 00:32:07.047: I neuronx_distributed/paral -lel_layers/parallel_state.py:672] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_model_groups: replica_groups.ep_model_groups=[[0]] [2025-07-26 00:32:07.047: I neuronx_distributed/paral -lel_layers/parallel_state.py:673] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_data_groups: replica_groups.ep_data_groups=[[0]] INFO:__main__:Distributed environment initialized suc -cessfully INFO:__main__:Creating NeuronConfig with minimal stab -le settings INFO:__main__:Created NeuronConfig: INFO:__main__:Loading model configuration... +lel_layers/parallel_state.py:668] [rank_0_pp-1_tp-1_dp-1_cp-1] tp_groups: replica_groups.tp_groups=[[0]] [2025-07-26 00:32:07.047: I neuronx_distributed/paral +lel_layers/parallel_state.py:669] [rank_0_pp-1_tp-1_dp-1_cp-1] dp_groups: replica_groups.dp_groups=[[0]] [2025-07-26 00:32:07.047: I neuronx_distributed/paral +lel_layers/parallel_state.py:670] [rank_0_pp-1_tp-1_dp-1_cp-1] pp_groups: replica_groups.pp_groups=[[0]] [2025-07-26 00:32:07.047: I neuronx_distributed/paral +lel_layers/parallel_state.py:671] [rank_0_pp-1_tp-1_dp-1_cp-1] cp_groups: replica_groups.cp_groups=[[0]] [2025-07-26 00:32:07.047: I neuronx_distributed/paral +lel_layers/parallel_state.py:672] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_model_groups: replica_groups.ep_model_groups=[[0]] [2025-07-26 00:32:07.047: I neuronx_distributed/paral +lel_layers/parallel_state.py:673] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_data_groups: replica_groups.ep_data_groups=[[0]] INFO:**main**:Distributed environment initialized suc +cessfully INFO:**main**:Creating NeuronConfig with minimal stab +le settings INFO:**main**:Created NeuronConfig: INFO:**main**:Loading model configuration... INFO:neuronx_llama3.modeling_llama3:Loading HuggingFa -ce configuration from ./llama3_neuron_checkpoint/config.json INFO:neuronx_llama3.modeling_llama3:Loaded Llama3 con -figuration: {'hidden_size': 2048, 'num_attention_heads': 32, 'num_hidden_layers': 16, 'num_key_value_heads': 8, 'vocab_size': 128256, 'max_position_embeddings': 2048, 'rope_theta': 500000.0, 'rms_norm_eps': 1e-05, 'hidden_act': 'silu', 'tie_word_embeddings': False, 'use_scaled_rope': True, 'intermediate_size': 8192} INFO:__main__:Loaded configuration: hidden_size=2048, - num_layers=16, vocab_size=128256 INFO:__main__:Creating model instance... +ce configuration from ./llama3_neuron_checkpoint/config.json INFO:neuronx_llama3.modeling_llama3:Loaded Llama3 con +figuration: {'hidden_size': 2048, 'num_attention_heads': 32, 'num_hidden_layers': 16, 'num_key_value_heads': 8, 'vocab_size': 128256, 'max_position_embeddings': 2048, 'rope_theta': 500000.0, 'rms_norm_eps': 1e-05, 'hidden_act': 'silu', 'tie_word_embeddings': False, 'use_scaled_rope': True, 'intermediate_size': 8192} INFO:**main**:Loaded configuration: hidden_size=2048, +num_layers=16, vocab_size=128256 INFO:**main**:Creating model instance... INFO:root:neuronx-cc compiler_args are: --auto-cast=n -one --model-type=transformer --tensorizer-options='--enable-ccop-compute-overlap --cc-pipeline-tiling-factor=2 --vectorize-strided-dma ' --lnc=1 -O1 --internal-hlo2tensorizer-options=' --modular-flow-mac-threshold=10 --verify-hlo=true' INFO:root:neuronx-cc compiler_args are: --auto-cast=n -one --model-type=transformer --tensorizer-options='--enable-ccop-compute-overlap --cc-pipeline-tiling-factor=1 --vectorize-strided-dma ' --lnc=1 -O2 --internal-hlo2tensorizer-options='--verify-hlo=true' INFO:__main__:Loading model weights... -INFO:__main__:Loaded 164 parameters from ./llama3_neu -ron_checkpoint/pytorch_model.bin WARNING:__main__:Unexpected keys in state dict: ['emb -ed_tokens.weight', 'norm.weight', 'lm_head.weight', 'layers.0.self_attn.q_proj.weight', 'layers.0.self_attn.k_proj.weight', 'layers.0.self_attn.v_proj.weight', 'layers.0.self_attn.o_proj.weight', 'layers.0.mlp.gate_proj.weight', 'layers.0.mlp.down_proj.weight', 'layers.0.mlp.up_proj.weight', 'layers.0.input_layernorm.weight', 'layers.0.post_attention_layernorm.weight', 'layers.1.self_attn.q_proj.weight', 'layers.1.self_attn.k_proj.weight', 'layers.1.self_attn.v_proj.weight', 'layers.1.self_attn.o_proj.weight', 'layers.1.mlp.gate_proj.weight', 'layers.1.mlp.down_proj.weight', 'layers.1.mlp.up_proj.weight', 'layers.1.input_layernorm.weight', 'layers.1.post_attention_layernorm.weight', 'layers.2.self_attn.q_proj.weight', 'layers.2.self_attn.k_proj.weight', 'layers.2.self_attn.v_proj.weight', 'layers.2.self_attn.o_proj.weight', 'layers.2.mlp.gate_proj.weight', 'layers.2.mlp.down_proj.weight', 'layers.2.mlp.up_proj.weight', 'layers.2.input_layernorm.weight', 'layers.2.post_attention_layernorm.weight', 'layers.3.self_attn.q_proj.weight', 'layers.3.self_attn.k_proj.weight', 'layers.3.self_attn.v_proj.weight', 'layers.3.self_attn.o_proj.weight', 'layers.3.mlp.gate_proj.weight', 'layers.3.mlp.down_proj.weight', 'layers.3.mlp.up_proj.weight', 'layers.3.input_layernorm.weight', 'layers.3.post_attention_layernorm.weight', 'layers.4.self_attn.q_proj.weight', 'layers.4.self_attn.k_proj.weight', 'layers.4.self_attn.v_proj.weight', 'layers.4.self_attn.o_proj.weight', 'layers.4.mlp.gate_proj.weight', 'layers.4.mlp.down_proj.weight', 'layers.4.mlp.up_proj.weight', 'layers.4.input_layernorm.weight', 'layers.4.post_attention_layernorm.weight', 'layers.5.self_attn.q_proj.weight', 'layers.5.self_attn.k_proj.weight', 'layers.5.self_attn.v_proj.weight', 'layers.5.self_attn.o_proj.weight', 'layers.5.mlp.gate_proj.weight', 'layers.5.mlp.down_proj.weight', 'layers.5.mlp.up_proj.weight', 'layers.5.input_layernorm.weight', 'layers.5.post_attention_layernorm.weight', 'layers.6.self_attn.q_proj.weight', 'layers.6.self_attn.k_proj.weight', 'layers.6.self_attn.v_proj.weight', 'layers.6.self_attn.o_proj.weight', 'layers.6.mlp.gate_proj.weight', 'layers.6.mlp.down_proj.weight', 'layers.6.mlp.up_proj.weight', 'layers.6.input_layernorm.weight', 'layers.6.post_attention_layernorm.weight', 'layers.7.self_attn.q_proj.weight', 'layers.7.self_attn.k_proj.weight', 'layers.7.self_attn.v_proj.weight', 'layers.7.self_attn.o_proj.weight', 'layers.7.mlp.gate_proj.weight', 'layers.7.mlp.down_proj.weight', 'layers.7.mlp.up_proj.weight', 'layers.7.input_layernorm.weight', 'layers.7.post_attention_layernorm.weight', 'layers.8.self_attn.q_proj.weight', 'layers.8.self_attn.k_proj.weight', 'layers.8.self_attn.v_proj.weight', 'layers.8.self_attn.o_proj.weight', 'layers.8.mlp.gate_proj.weight', 'layers.8.mlp.down_proj.weight', 'layers.8.mlp.up_proj.weight', 'layers.8.input_layernorm.weight', 'layers.8.post_attention_layernorm.weight', 'layers.9.self_attn.q_proj.weight', 'layers.9.self_attn.k_proj.weight', 'layers.9.self_attn.v_proj.weight', 'layers.9.self_attn.o_proj.weight', 'layers.9.mlp.gate_proj.weight', 'layers.9.mlp.down_proj.weight', 'layers.9.mlp.up_proj.weight', 'layers.9.input_layernorm.weight', 'layers.9.post_attention_layernorm.weight', 'layers.10.self_attn.q_proj.weight', 'layers.10.self_attn.k_proj.weight', 'layers.10.self_attn.v_proj.weight', 'layers.10.self_attn.o_proj.weight', 'layers.10.mlp.gate_proj.weight', 'layers.10.mlp.down_proj.weight', 'layers.10.mlp.up_proj.weight', 'layers.10.input_layernorm.weight', 'layers.10.post_attention_layernorm.weight', 'layers.11.self_attn.q_proj.weight', 'layers.11.self_attn.k_proj.weight', 'layers.11.self_attn.v_proj.weight', 'layers.11.self_attn.o_proj.weight', 'layers.11.mlp.gate_proj.weight', 'layers.11.mlp.down_proj.weight', 'layers.11.mlp.up_proj.weight', 'layers.11.input_layernorm.weight', 'layers.11.post_attention_layernorm.weight', 'layers.12.self_attn.q_proj.weight', 'layers.12.self_attn.k_proj.weight', 'layers.12.self_attn.v_proj.weight', 'layers.12.self_attn.o_proj.weight', 'layers.12.mlp.gate_proj.weight', 'layers.12.mlp.down_proj.weight', 'layers.12.mlp.up_proj.weight', 'layers.12.input_layernorm.weight', 'layers.12.post_attention_layernorm.weight', 'layers.13.self_attn.q_proj.weight', 'layers.13.self_attn.k_proj.weight', 'layers.13.self_attn.v_proj.weight', 'layers.13.self_attn.o_proj.weight', 'layers.13.mlp.gate_proj.weight', 'layers.13.mlp.down_proj.weight', 'layers.13.mlp.up_proj.weight', 'layers.13.input_layernorm.weight', 'layers.13.post_attention_layernorm.weight', 'layers.14.self_attn.q_proj.weight', 'layers.14.self_attn.k_proj.weight', 'layers.14.self_attn.v_proj.weight', 'layers.14.self_attn.o_proj.weight', 'layers.14.mlp.gate_proj.weight', 'layers.14.mlp.down_proj.weight', 'layers.14.mlp.up_proj.weight', 'layers.14.input_layernorm.weight', 'layers.14.post_attention_layernorm.weight', 'layers.15.self_attn.q_proj.weight', 'layers.15.self_attn.k_proj.weight', 'layers.15.self_attn.v_proj.weight', 'layers.15.self_attn.o_proj.weight', 'layers.15.mlp.gate_proj.weight', 'layers.15.mlp.down_proj.weight', 'layers.15.mlp.up_proj.weight', 'layers.15.input_layernorm.weight', 'layers.15.post_attention_layernorm.weight', 'layers.0.self_attn.rank_util.rank', 'layers.1.self_attn.rank_util.rank', 'layers.2.self_attn.rank_util.rank', 'layers.3.self_attn.rank_util.rank', 'layers.4.self_attn.rank_util.rank', 'layers.5.self_attn.rank_util.rank', 'layers.6.self_attn.rank_util.rank', 'layers.7.self_attn.rank_util.rank', 'layers.8.self_attn.rank_util.rank', 'layers.9.self_attn.rank_util.rank', 'layers.10.self_attn.rank_util.rank', 'layers.11.self_attn.rank_util.rank', 'layers.12.self_attn.rank_util.rank', 'layers.13.self_attn.rank_util.rank', 'layers.14.self_attn.rank_util.rank', 'layers.15.self_attn.rank_util.rank', 'rank_util.rank'] INFO:__main__:Model and weights loaded successfully -INFO:__main__:Created output directory: ./llama3_comp -iled INFO:__main__:Compiling model for Neuron hardware... -INFO:__main__:This may take several minutes... +one --model-type=transformer --tensorizer-options='--enable-ccop-compute-overlap --cc-pipeline-tiling-factor=2 --vectorize-strided-dma ' --lnc=1 -O1 --internal-hlo2tensorizer-options=' --modular-flow-mac-threshold=10 --verify-hlo=true' INFO:root:neuronx-cc compiler_args are: --auto-cast=n +one --model-type=transformer --tensorizer-options='--enable-ccop-compute-overlap --cc-pipeline-tiling-factor=1 --vectorize-strided-dma ' --lnc=1 -O2 --internal-hlo2tensorizer-options='--verify-hlo=true' INFO:**main**:Loading model weights... +INFO:**main**:Loaded 164 parameters from ./llama3_neu +ron_checkpoint/pytorch_model.bin WARNING:**main**:Unexpected keys in state dict: ['emb +ed_tokens.weight', 'norm.weight', 'lm_head.weight', 'layers.0.self_attn.q_proj.weight', 'layers.0.self_attn.k_proj.weight', 'layers.0.self_attn.v_proj.weight', 'layers.0.self_attn.o_proj.weight', 'layers.0.mlp.gate_proj.weight', 'layers.0.mlp.down_proj.weight', 'layers.0.mlp.up_proj.weight', 'layers.0.input_layernorm.weight', 'layers.0.post_attention_layernorm.weight', 'layers.1.self_attn.q_proj.weight', 'layers.1.self_attn.k_proj.weight', 'layers.1.self_attn.v_proj.weight', 'layers.1.self_attn.o_proj.weight', 'layers.1.mlp.gate_proj.weight', 'layers.1.mlp.down_proj.weight', 'layers.1.mlp.up_proj.weight', 'layers.1.input_layernorm.weight', 'layers.1.post_attention_layernorm.weight', 'layers.2.self_attn.q_proj.weight', 'layers.2.self_attn.k_proj.weight', 'layers.2.self_attn.v_proj.weight', 'layers.2.self_attn.o_proj.weight', 'layers.2.mlp.gate_proj.weight', 'layers.2.mlp.down_proj.weight', 'layers.2.mlp.up_proj.weight', 'layers.2.input_layernorm.weight', 'layers.2.post_attention_layernorm.weight', 'layers.3.self_attn.q_proj.weight', 'layers.3.self_attn.k_proj.weight', 'layers.3.self_attn.v_proj.weight', 'layers.3.self_attn.o_proj.weight', 'layers.3.mlp.gate_proj.weight', 'layers.3.mlp.down_proj.weight', 'layers.3.mlp.up_proj.weight', 'layers.3.input_layernorm.weight', 'layers.3.post_attention_layernorm.weight', 'layers.4.self_attn.q_proj.weight', 'layers.4.self_attn.k_proj.weight', 'layers.4.self_attn.v_proj.weight', 'layers.4.self_attn.o_proj.weight', 'layers.4.mlp.gate_proj.weight', 'layers.4.mlp.down_proj.weight', 'layers.4.mlp.up_proj.weight', 'layers.4.input_layernorm.weight', 'layers.4.post_attention_layernorm.weight', 'layers.5.self_attn.q_proj.weight', 'layers.5.self_attn.k_proj.weight', 'layers.5.self_attn.v_proj.weight', 'layers.5.self_attn.o_proj.weight', 'layers.5.mlp.gate_proj.weight', 'layers.5.mlp.down_proj.weight', 'layers.5.mlp.up_proj.weight', 'layers.5.input_layernorm.weight', 'layers.5.post_attention_layernorm.weight', 'layers.6.self_attn.q_proj.weight', 'layers.6.self_attn.k_proj.weight', 'layers.6.self_attn.v_proj.weight', 'layers.6.self_attn.o_proj.weight', 'layers.6.mlp.gate_proj.weight', 'layers.6.mlp.down_proj.weight', 'layers.6.mlp.up_proj.weight', 'layers.6.input_layernorm.weight', 'layers.6.post_attention_layernorm.weight', 'layers.7.self_attn.q_proj.weight', 'layers.7.self_attn.k_proj.weight', 'layers.7.self_attn.v_proj.weight', 'layers.7.self_attn.o_proj.weight', 'layers.7.mlp.gate_proj.weight', 'layers.7.mlp.down_proj.weight', 'layers.7.mlp.up_proj.weight', 'layers.7.input_layernorm.weight', 'layers.7.post_attention_layernorm.weight', 'layers.8.self_attn.q_proj.weight', 'layers.8.self_attn.k_proj.weight', 'layers.8.self_attn.v_proj.weight', 'layers.8.self_attn.o_proj.weight', 'layers.8.mlp.gate_proj.weight', 'layers.8.mlp.down_proj.weight', 'layers.8.mlp.up_proj.weight', 'layers.8.input_layernorm.weight', 'layers.8.post_attention_layernorm.weight', 'layers.9.self_attn.q_proj.weight', 'layers.9.self_attn.k_proj.weight', 'layers.9.self_attn.v_proj.weight', 'layers.9.self_attn.o_proj.weight', 'layers.9.mlp.gate_proj.weight', 'layers.9.mlp.down_proj.weight', 'layers.9.mlp.up_proj.weight', 'layers.9.input_layernorm.weight', 'layers.9.post_attention_layernorm.weight', 'layers.10.self_attn.q_proj.weight', 'layers.10.self_attn.k_proj.weight', 'layers.10.self_attn.v_proj.weight', 'layers.10.self_attn.o_proj.weight', 'layers.10.mlp.gate_proj.weight', 'layers.10.mlp.down_proj.weight', 'layers.10.mlp.up_proj.weight', 'layers.10.input_layernorm.weight', 'layers.10.post_attention_layernorm.weight', 'layers.11.self_attn.q_proj.weight', 'layers.11.self_attn.k_proj.weight', 'layers.11.self_attn.v_proj.weight', 'layers.11.self_attn.o_proj.weight', 'layers.11.mlp.gate_proj.weight', 'layers.11.mlp.down_proj.weight', 'layers.11.mlp.up_proj.weight', 'layers.11.input_layernorm.weight', 'layers.11.post_attention_layernorm.weight', 'layers.12.self_attn.q_proj.weight', 'layers.12.self_attn.k_proj.weight', 'layers.12.self_attn.v_proj.weight', 'layers.12.self_attn.o_proj.weight', 'layers.12.mlp.gate_proj.weight', 'layers.12.mlp.down_proj.weight', 'layers.12.mlp.up_proj.weight', 'layers.12.input_layernorm.weight', 'layers.12.post_attention_layernorm.weight', 'layers.13.self_attn.q_proj.weight', 'layers.13.self_attn.k_proj.weight', 'layers.13.self_attn.v_proj.weight', 'layers.13.self_attn.o_proj.weight', 'layers.13.mlp.gate_proj.weight', 'layers.13.mlp.down_proj.weight', 'layers.13.mlp.up_proj.weight', 'layers.13.input_layernorm.weight', 'layers.13.post_attention_layernorm.weight', 'layers.14.self_attn.q_proj.weight', 'layers.14.self_attn.k_proj.weight', 'layers.14.self_attn.v_proj.weight', 'layers.14.self_attn.o_proj.weight', 'layers.14.mlp.gate_proj.weight', 'layers.14.mlp.down_proj.weight', 'layers.14.mlp.up_proj.weight', 'layers.14.input_layernorm.weight', 'layers.14.post_attention_layernorm.weight', 'layers.15.self_attn.q_proj.weight', 'layers.15.self_attn.k_proj.weight', 'layers.15.self_attn.v_proj.weight', 'layers.15.self_attn.o_proj.weight', 'layers.15.mlp.gate_proj.weight', 'layers.15.mlp.down_proj.weight', 'layers.15.mlp.up_proj.weight', 'layers.15.input_layernorm.weight', 'layers.15.post_attention_layernorm.weight', 'layers.0.self_attn.rank_util.rank', 'layers.1.self_attn.rank_util.rank', 'layers.2.self_attn.rank_util.rank', 'layers.3.self_attn.rank_util.rank', 'layers.4.self_attn.rank_util.rank', 'layers.5.self_attn.rank_util.rank', 'layers.6.self_attn.rank_util.rank', 'layers.7.self_attn.rank_util.rank', 'layers.8.self_attn.rank_util.rank', 'layers.9.self_attn.rank_util.rank', 'layers.10.self_attn.rank_util.rank', 'layers.11.self_attn.rank_util.rank', 'layers.12.self_attn.rank_util.rank', 'layers.13.self_attn.rank_util.rank', 'layers.14.self_attn.rank_util.rank', 'layers.15.self_attn.rank_util.rank', 'rank_util.rank'] INFO:**main**:Model and weights loaded successfully +INFO:**main**:Created output directory: ./llama3_comp +iled INFO:**main**:Compiling model for Neuron hardware... +INFO:**main**:This may take several minutes... DEBUG:root:Saving config: { - "fused_spec_config": null, - "hidden_act": "silu", - "hidden_size": 2048, - "intermediate_size": 8192, - "max_position_embeddings": 2048, - "metadata": null, - "neuron_config": { - "activation_quantization_type": null, - "allow_input_truncation": false, - "apply_seq_ids_mask": false, - "async_mode": false, - "attention_dp_degree": 1, - "attention_dtype": null, - "attn_block_cte_nki_kernel_enabled": false, - "attn_block_tkg_nki_kernel_cache_update": false, - "attn_block_tkg_nki_kernel_enabled": false, - "attn_cls": "NeuronLlamaAttention", - "attn_kernel_enabled": null, - "attn_tkg_builtin_kernel_enabled": false, - "attn_tkg_nki_kernel_enabled": false, - "batch_size": 1, - "bucket_n_active_tokens": false, - "buckets": [ - 128 - ], - "cast_type": "config", - "cc_pipeline_tiling_factor": 2, - "chunked_prefill_config": null, - "context_encoding_buckets": null, - "cp_degree": 1, - "ctx_batch_size": 1, - "disable_kv_cache_tiling": false, - "draft_model_modules_to_not_convert": null, - "enable_bucketing": false, - "enable_cte_modular_flow": false, - "enable_eagle_draft_input_norm": false, - "enable_eagle_speculation": false, - "enable_fused_speculation": false, - "enable_long_context_mode": false, - "enable_output_completion_notifications": false, - "enable_spill_reload_dge": false, - "enable_token_tree": false, - "ep_degree": 1, - "expert_mlp_nki_kernel_enabled": null, - "flash_decoding_enabled": false, - "fused_qkv": false, - "fused_rmsnorm_skip_gamma": false, - "is_block_kv_layout": false, - "is_chunked_prefill": false, - "is_continuous_batching": false, - "is_eagle_draft": false, - "is_medusa": false, - "is_prefill_stage": null, - "is_prefix_caching": false, - "k_cache_transposed": false, - "kv_cache_batch_size": 1, - "kv_cache_padding_size": 0, - "kv_cache_quant": false, - "kv_cache_tiling": false, - "layer_boundary_markers": false, - "lm_head_pad": false, - "lm_head_pad_alignment_size": 1, - "local_ranks_size": 1, - "logical_nc_config": 1, - "lora_config": null, - "max_batch_size": 1, - "max_context_length": 128, - "max_length": 128, - "max_new_tokens": null, - "medusa_speculation_length": 0, - "medusa_tree": null, - "mlp_kernel_enabled": false, - "mlp_kernel_fuse_residual_add": false, - "modules_to_not_convert": null, - "moe_fused_nki_kernel_enabled": null, - "n_active_tokens": 128, - "n_positions": 128, - "num_medusa_heads": 0, - "on_cpu": false, - "on_device_sampling_config": null, - "output_logits": false, - "overrides_torch_dtype": true, - "pa_block_size": 128, - "pa_num_blocks": 1, - "padding_side": "right", - "pp_degree": 1, - "prefix_buckets": null, - "qk_layernorm": false, - "qkv_kernel_enabled": false, - "qkv_kernel_fuse_residual_add": false, - "qkv_kernel_nbsd_layout": false, - "quantization_dtype": "int8", - "quantization_type": "per_tensor_symmetric", - "quantize_clamp_bound": Infinity, - "quantized": false, - "quantized_checkpoints_path": null, - "quantized_mlp_kernel_enabled": false, - "rmsnorm_quantize_kernel_enabled": false, - "router_topk_nki_kernel_enabled": null, - "rpl_reduce_dtype": null, - "save_sharded_checkpoint": false, - "scratchpad_page_size": null, - "seq_len": 128, - "seq_len_threshold_for_cc_tiling": 16384, - "sequence_parallel_enabled": false, - "shared_mlp_nki_kernel_enabled": null, - "skip_sharding": false, - "skip_warmup": false, - "spec_batch_size": 1, - "speculation_length": 0, - "start_rank_id": 0, - "strided_context_parallel_kernel_enabled": false, - "target": null, - "tensor_capture_config": null, - "tile_cc": false, - "tkg_batch_size": 1, - "token_generation_buckets": null, - "token_tree_config": null, - "torch_dtype": "float32", - "tp_degree": 1, - "vocab_parallel": false, - "weight_gather_seq_len_threshold": 32768, - "weights_to_skip_layout_optimization": [], - "world_size": 1 - }, - "num_attention_heads": 32, - "num_cores_per_group": 1, - "num_hidden_layers": 16, - "num_key_value_heads": 8, - "output_attentions": false, - "output_hidden_states": false, - "rms_norm_eps": 1e-05, - "rope_theta": 500000.0, - "tie_word_embeddings": false, - "use_cache": true, - "use_return_dict": true, - "use_scaled_rope": true, - "vocab_size": 128256 +"fused_spec_config": null, +"hidden_act": "silu", +"hidden_size": 2048, +"intermediate_size": 8192, +"max_position_embeddings": 2048, +"metadata": null, +"neuron_config": { +"activation_quantization_type": null, +"allow_input_truncation": false, +"apply_seq_ids_mask": false, +"async_mode": false, +"attention_dp_degree": 1, +"attention_dtype": null, +"attn_block_cte_nki_kernel_enabled": false, +"attn_block_tkg_nki_kernel_cache_update": false, +"attn_block_tkg_nki_kernel_enabled": false, +"attn_cls": "NeuronLlamaAttention", +"attn_kernel_enabled": null, +"attn_tkg_builtin_kernel_enabled": false, +"attn_tkg_nki_kernel_enabled": false, +"batch_size": 1, +"bucket_n_active_tokens": false, +"buckets": [ +128 +], +"cast_type": "config", +"cc_pipeline_tiling_factor": 2, +"chunked_prefill_config": null, +"context_encoding_buckets": null, +"cp_degree": 1, +"ctx_batch_size": 1, +"disable_kv_cache_tiling": false, +"draft_model_modules_to_not_convert": null, +"enable_bucketing": false, +"enable_cte_modular_flow": false, +"enable_eagle_draft_input_norm": false, +"enable_eagle_speculation": false, +"enable_fused_speculation": false, +"enable_long_context_mode": false, +"enable_output_completion_notifications": false, +"enable_spill_reload_dge": false, +"enable_token_tree": false, +"ep_degree": 1, +"expert_mlp_nki_kernel_enabled": null, +"flash_decoding_enabled": false, +"fused_qkv": false, +"fused_rmsnorm_skip_gamma": false, +"is_block_kv_layout": false, +"is_chunked_prefill": false, +"is_continuous_batching": false, +"is_eagle_draft": false, +"is_medusa": false, +"is_prefill_stage": null, +"is_prefix_caching": false, +"k_cache_transposed": false, +"kv_cache_batch_size": 1, +"kv_cache_padding_size": 0, +"kv_cache_quant": false, +"kv_cache_tiling": false, +"layer_boundary_markers": false, +"lm_head_pad": false, +"lm_head_pad_alignment_size": 1, +"local_ranks_size": 1, +"logical_nc_config": 1, +"lora_config": null, +"max_batch_size": 1, +"max_context_length": 128, +"max_length": 128, +"max_new_tokens": null, +"medusa_speculation_length": 0, +"medusa_tree": null, +"mlp_kernel_enabled": false, +"mlp_kernel_fuse_residual_add": false, +"modules_to_not_convert": null, +"moe_fused_nki_kernel_enabled": null, +"n_active_tokens": 128, +"n_positions": 128, +"num_medusa_heads": 0, +"on_cpu": false, +"on_device_sampling_config": null, +"output_logits": false, +"overrides_torch_dtype": true, +"pa_block_size": 128, +"pa_num_blocks": 1, +"padding_side": "right", +"pp_degree": 1, +"prefix_buckets": null, +"qk_layernorm": false, +"qkv_kernel_enabled": false, +"qkv_kernel_fuse_residual_add": false, +"qkv_kernel_nbsd_layout": false, +"quantization_dtype": "int8", +"quantization_type": "per_tensor_symmetric", +"quantize_clamp_bound": Infinity, +"quantized": false, +"quantized_checkpoints_path": null, +"quantized_mlp_kernel_enabled": false, +"rmsnorm_quantize_kernel_enabled": false, +"router_topk_nki_kernel_enabled": null, +"rpl_reduce_dtype": null, +"save_sharded_checkpoint": false, +"scratchpad_page_size": null, +"seq_len": 128, +"seq_len_threshold_for_cc_tiling": 16384, +"sequence_parallel_enabled": false, +"shared_mlp_nki_kernel_enabled": null, +"skip_sharding": false, +"skip_warmup": false, +"spec_batch_size": 1, +"speculation_length": 0, +"start_rank_id": 0, +"strided_context_parallel_kernel_enabled": false, +"target": null, +"tensor_capture_config": null, +"tile_cc": false, +"tkg_batch_size": 1, +"token_generation_buckets": null, +"token_tree_config": null, +"torch_dtype": "float32", +"tp_degree": 1, +"vocab_parallel": false, +"weight_gather_seq_len_threshold": 32768, +"weights_to_skip_layout_optimization": [], +"world_size": 1 +}, +"num_attention_heads": 32, +"num_cores_per_group": 1, +"num_hidden_layers": 16, +"num_key_value_heads": 8, +"output_attentions": false, +"output_hidden_states": false, +"rms_norm_eps": 1e-05, +"rope_theta": 500000.0, +"tie_word_embeddings": false, +"use_cache": true, +"use_return_dict": true, +"use_scaled_rope": true, +"vocab_size": 128256 } INFO:Neuron:Saving the neuron_config to llama3_compil -ed/ INFO:Neuron:Generating HLOs for the following models: - ['context_encoding_model', 'token_generation_model'][2025-07-26 00:32:08.310: I neuronx_distributed/paral -lel_layers/parallel_state.py:628] > initializing tensor model parallel with size 1 [2025-07-26 00:32:08.311: I neuronx_distributed/paral -lel_layers/parallel_state.py:629] > initializing pipeline model parallel with size 1 [2025-07-26 00:32:08.311: I neuronx_distributed/paral -lel_layers/parallel_state.py:630] > initializing context model parallel with size 1 [2025-07-26 00:32:08.311: I neuronx_distributed/paral -lel_layers/parallel_state.py:631] > initializing data parallel with size 1 [2025-07-26 00:32:08.311: I neuronx_distributed/paral -lel_layers/parallel_state.py:632] > initializing world size to 1 [2025-07-26 00:32:08.312: I neuronx_distributed/paral -lel_layers/parallel_state.py:379] [rank_0_pp0_tp0_dp0_cp0] Chosen Logic for replica groups ret_logic=, 'Ascending Ring PG Group')> [2025-07-26 00:32:08.312: I neuronx_distributed/paral -lel_layers/parallel_state.py:668] [rank_0_pp0_tp0_dp0_cp0] tp_groups: replica_groups.tp_groups=[[0]] [2025-07-26 00:32:08.312: I neuronx_distributed/paral -lel_layers/parallel_state.py:669] [rank_0_pp0_tp0_dp0_cp0] dp_groups: replica_groups.dp_groups=[[0]] [2025-07-26 00:32:08.312: I neuronx_distributed/paral -lel_layers/parallel_state.py:670] [rank_0_pp0_tp0_dp0_cp0] pp_groups: replica_groups.pp_groups=[[0]] [2025-07-26 00:32:08.312: I neuronx_distributed/paral -lel_layers/parallel_state.py:671] [rank_0_pp0_tp0_dp0_cp0] cp_groups: replica_groups.cp_groups=[[0]] [2025-07-26 00:32:08.312: I neuronx_distributed/paral -lel_layers/parallel_state.py:672] [rank_0_pp0_tp0_dp0_cp0] ep_model_groups: replica_groups.ep_model_groups=[[0]] [2025-07-26 00:32:08.312: I neuronx_distributed/paral -lel_layers/parallel_state.py:673] [rank_0_pp0_tp0_dp0_cp0] ep_data_groups: replica_groups.ep_data_groups=[[0]] INFO:Neuron:Generating 1 hlos for key: context_encodi -ng_model INFO:Neuron:Started loading module context_encoding_m -odel ERROR:__main__:Model compilation failed: setup_attr_f -or_model() is not implemented ERROR:__main__:Compilation failed: setup_attr_for_mod -el() is not implemented Traceback (most recent call last): - File "/home/ec2-user/NeuronxSDK/neuronx_llama3/comp -ile_model.py", line 353, in main compile_model(args) - File "/home/ec2-user/NeuronxSDK/neuronx_llama3/comp -ile_model.py", line 291, in compile_model model.compile(output_path) - File "/home/ec2-user/NeuronxSDK/NeuronxDistributedI -nference/src/neuronx_distributed_inference/models/application_base.py", line 272, in compile traced_model = self.get_builder(debug).trace( - File "/home/ec2-user/NeuronxSDK/NeuronxDistributed/ -src/neuronx_distributed/trace/model_builder.py", line 1416, in trace hlo_artifact_collection = self._generate_hlo(key) - File "/home/ec2-user/NeuronxSDK/NeuronxDistributed/ -src/neuronx_distributed/trace/model_builder.py", line 1702, in _generate_hlo model_input_container.model_instance.load_module( -) File "/home/ec2-user/NeuronxSDK/NeuronxDistributedI -nference/src/neuronx_distributed_inference/models/model_wrapper.py", line 1543, in load_module float_model = self.model_cls(self.config, **self. -kwargs) File "/home/ec2-user/NeuronxSDK/neuronx_llama3/src/ -neuronx_llama3/modeling_llama3.py", line 555, in __init__ super().__init__(config) - File "/home/ec2-user/NeuronxSDK/NeuronxDistributedI -nference/src/neuronx_distributed_inference/models/model_base.py", line 120, in __init__ self.setup_attr_for_model(config) - File "/home/ec2-user/NeuronxSDK/NeuronxDistributedI -nference/src/neuronx_distributed_inference/models/model_base.py", line 147, in setup_attr_for_model raise NotImplementedError("setup_attr_for_model() - is not implemented") NotImplementedError: setup_attr_for_model() is not im -plemented +ed/ INFO:Neuron:Generating HLOs for the following models: +['context_encoding_model', 'token_generation_model'][2025-07-26 00:32:08.310: I neuronx_distributed/paral lel_layers/parallel_state.py:628] > initializing tensor model parallel with size 1 [2025-07-26 00:32:08.311: I neuronx_distributed/paral +lel_layers/parallel_state.py:629] > initializing pipeline model parallel with size 1 [2025-07-26 00:32:08.311: I neuronx_distributed/paral +lel_layers/parallel_state.py:630] > initializing context model parallel with size 1 [2025-07-26 00:32:08.311: I neuronx_distributed/paral +lel_layers/parallel_state.py:631] > initializing data parallel with size 1 [2025-07-26 00:32:08.311: I neuronx_distributed/paral +lel_layers/parallel_state.py:632] > initializing world size to 1 [2025-07-26 00:32:08.312: I neuronx_distributed/paral +lel_layers/parallel_state.py:379] [rank_0_pp0_tp0_dp0_cp0] Chosen Logic for replica groups ret_logic=, 'Ascending Ring PG Group')> [2025-07-26 00:32:08.312: I neuronx_distributed/paral +lel_layers/parallel_state.py:668] [rank_0_pp0_tp0_dp0_cp0] tp_groups: replica_groups.tp_groups=[[0]] [2025-07-26 00:32:08.312: I neuronx_distributed/paral +lel_layers/parallel_state.py:669] [rank_0_pp0_tp0_dp0_cp0] dp_groups: replica_groups.dp_groups=[[0]] [2025-07-26 00:32:08.312: I neuronx_distributed/paral +lel_layers/parallel_state.py:670] [rank_0_pp0_tp0_dp0_cp0] pp_groups: replica_groups.pp_groups=[[0]] [2025-07-26 00:32:08.312: I neuronx_distributed/paral +lel_layers/parallel_state.py:671] [rank_0_pp0_tp0_dp0_cp0] cp_groups: replica_groups.cp_groups=[[0]] [2025-07-26 00:32:08.312: I neuronx_distributed/paral +lel_layers/parallel_state.py:672] [rank_0_pp0_tp0_dp0_cp0] ep_model_groups: replica_groups.ep_model_groups=[[0]] [2025-07-26 00:32:08.312: I neuronx_distributed/paral +lel_layers/parallel_state.py:673] [rank_0_pp0_tp0_dp0_cp0] ep_data_groups: replica_groups.ep_data_groups=[[0]] INFO:Neuron:Generating 1 hlos for key: context_encodi +ng_model INFO:Neuron:Started loading module context_encoding_m +odel ERROR:**main**:Model compilation failed: setup_attr_f +or_model() is not implemented ERROR:**main**:Compilation failed: setup_attr_for_mod +el() is not implemented Traceback (most recent call last): +File "/home/ec2-user/NeuronxSDK/neuronx_llama3/comp +ile_model.py", line 353, in main compile_model(args) +File "/home/ec2-user/NeuronxSDK/neuronx_llama3/comp +ile_model.py", line 291, in compile_model model.compile(output_path) +File "/home/ec2-user/NeuronxSDK/NeuronxDistributedI +nference/src/neuronx_distributed_inference/models/application_base.py", line 272, in compile traced_model = self.get_builder(debug).trace( +File "/home/ec2-user/NeuronxSDK/NeuronxDistributed/ +src/neuronx_distributed/trace/model_builder.py", line 1416, in trace hlo_artifact_collection = self.\_generate_hlo(key) +File "/home/ec2-user/NeuronxSDK/NeuronxDistributed/ +src/neuronx_distributed/trace/model_builder.py", line 1702, in \_generate_hlo model_input_container.model_instance.load_module( +) File "/home/ec2-user/NeuronxSDK/NeuronxDistributedI +nference/src/neuronx_distributed_inference/models/model_wrapper.py", line 1543, in load_module float_model = self.model_cls(self.config, \*\*self. +kwargs) File "/home/ec2-user/NeuronxSDK/neuronx_llama3/src/ +neuronx_llama3/modeling_llama3.py", line 555, in **init** super().**init**(config) +File "/home/ec2-user/NeuronxSDK/NeuronxDistributedI +nference/src/neuronx_distributed_inference/models/model_base.py", line 120, in **init** self.setup_attr_for_model(config) +File "/home/ec2-user/NeuronxSDK/NeuronxDistributedI +nference/src/neuronx_distributed_inference/models/model_base.py", line 147, in setup_attr_for_model raise NotImplementedError("setup_attr_for_model() +is not implemented") NotImplementedError: setup_attr_for_model() is not im +plemented I need to implement the setup_attr_for_model method in the NeuronLlama3Model class. Let me add this: Accepted edits to modeling_llama3.py - Now let's try compiling again: Command -python compile_model.py --checkpoint_path ./llama3_neuron_checkpoint --output_path ./llama3_compiled --batch_size 1 --seq_len 128 +python compile*model.py --checkpoint_path ./llama3_neuron_checkpoint --output_path ./llama3_compiled --batch_size 1 --seq_len 128 /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributedInference -/src/neuronx_distributed_inference/modules/attention/utils.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.custom_c -alls import neuron_cumsum /home/ec2-user/NeuronxSDK/NeuronxDistributedInference -/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference -/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference -/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_ -llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_ -llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_ -llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.attention_base import NeuronAttentionBase [2025-07-26 00:32:58.909: I neuronx_distributed/paral -lel_layers/parallel_state.py:628] > initializing tensor model parallel with size 1 [2025-07-26 00:32:58.910: I neuronx_distributed/paral -lel_layers/parallel_state.py:629] > initializing pipeline model parallel with size 1 [2025-07-26 00:32:58.910: I neuronx_distributed/paral -lel_layers/parallel_state.py:630] > initializing context model parallel with size 1 [2025-07-26 00:32:58.910: I neuronx_distributed/paral -lel_layers/parallel_state.py:631] > initializing data parallel with size 1 [2025-07-26 00:32:58.910: I neuronx_distributed/paral -lel_layers/parallel_state.py:632] > initializing world size to 1 2025-07-26 00:32:58.000916: 343763 INFO ||NEURON_CC -_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/neuronxcc-2.19.8089.0+8ab9f450/MODULE_13537832482499839909+e30acd3a/model.neff [2025-07-26 00:32:58.955: I neuronx_distributed/paral +/src/neuronx_distributed_inference/modules/attention/utils.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.custom_c +alls import neuron_cumsum /home/ec2-user/NeuronxSDK/NeuronxDistributedInference +/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference +/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference +/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx* +llama3/modeling*llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx* +llama3/modeling*llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx* +llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.attention_base import NeuronAttentionBase [2025-07-26 00:32:58.909: I neuronx_distributed/paral +lel_layers/parallel_state.py:628] > initializing tensor model parallel with size 1 [2025-07-26 00:32:58.910: I neuronx_distributed/paral +lel_layers/parallel_state.py:629] > initializing pipeline model parallel with size 1 [2025-07-26 00:32:58.910: I neuronx_distributed/paral +lel_layers/parallel_state.py:630] > initializing context model parallel with size 1 [2025-07-26 00:32:58.910: I neuronx_distributed/paral +lel_layers/parallel_state.py:631] > initializing data parallel with size 1 [2025-07-26 00:32:58.910: I neuronx_distributed/paral +lel_layers/parallel_state.py:632] > initializing world size to 1 2025-07-26 00:32:58.000916: 343763 INFO ||NEURON_CC +\_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/neuronxcc-2.19.8089.0+8ab9f450/MODULE_13537832482499839909+e30acd3a/model.neff [2025-07-26 00:32:58.955: I neuronx_distributed/paral lel_layers/parallel_state.py:379] [rank_0_pp-1_tp-1_dp-1_cp-1] Chosen Logic for replica groups ret_logic=, 'Ascending Ring PG Group')> [2025-07-26 00:32:58.955: I neuronx_distributed/paral -lel_layers/parallel_state.py:668] [rank_0_pp-1_tp-1_dp-1_cp-1] tp_groups: replica_groups.tp_groups=[[0]] [2025-07-26 00:32:58.955: I neuronx_distributed/paral -lel_layers/parallel_state.py:669] [rank_0_pp-1_tp-1_dp-1_cp-1] dp_groups: replica_groups.dp_groups=[[0]] [2025-07-26 00:32:58.955: I neuronx_distributed/paral -lel_layers/parallel_state.py:670] [rank_0_pp-1_tp-1_dp-1_cp-1] pp_groups: replica_groups.pp_groups=[[0]] [2025-07-26 00:32:58.955: I neuronx_distributed/paral -lel_layers/parallel_state.py:671] [rank_0_pp-1_tp-1_dp-1_cp-1] cp_groups: replica_groups.cp_groups=[[0]] [2025-07-26 00:32:58.955: I neuronx_distributed/paral -lel_layers/parallel_state.py:672] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_model_groups: replica_groups.ep_model_groups=[[0]] [2025-07-26 00:32:58.955: I neuronx_distributed/paral -lel_layers/parallel_state.py:673] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_data_groups: replica_groups.ep_data_groups=[[0]] WARNING:__main__:Unexpected keys in state dict: ['emb -ed_tokens.weight', 'norm.weight', 'lm_head.weight', 'layers.0.self_attn.q_proj.weight', 'layers.0.self_attn.k_proj.weight', 'layers.0.self_attn.v_proj.weight', 'layers.0.self_attn.o_proj.weight', 'layers.0.mlp.gate_proj.weight', 'layers.0.mlp.down_proj.weight', 'layers.0.mlp.up_proj.weight', 'layers.0.input_layernorm.weight', 'layers.0.post_attention_layernorm.weight', 'layers.1.self_attn.q_proj.weight', 'layers.1.self_attn.k_proj.weight', 'layers.1.self_attn.v_proj.weight', 'layers.1.self_attn.o_proj.weight', 'layers.1.mlp.gate_proj.weight', 'layers.1.mlp.down_proj.weight', 'layers.1.mlp.up_proj.weight', 'layers.1.input_layernorm.weight', 'layers.1.post_attention_layernorm.weight', 'layers.2.self_attn.q_proj.weight', 'layers.2.self_attn.k_proj.weight', 'layers.2.self_attn.v_proj.weight', 'layers.2.self_attn.o_proj.weight', 'layers.2.mlp.gate_proj.weight', 'layers.2.mlp.down_proj.weight', 'layers.2.mlp.up_proj.weight', 'layers.2.input_layernorm.weight', 'layers.2.post_attention_layernorm.weight', 'layers.3.self_attn.q_proj.weight', 'layers.3.self_attn.k_proj.weight', 'layers.3.self_attn.v_proj.weight', 'layers.3.self_attn.o_proj.weight', 'layers.3.mlp.gate_proj.weight', 'layers.3.mlp.down_proj.weight', 'layers.3.mlp.up_proj.weight', 'layers.3.input_layernorm.weight', 'layers.3.post_attention_layernorm.weight', 'layers.4.self_attn.q_proj.weight', 'layers.4.self_attn.k_proj.weight', 'layers.4.self_attn.v_proj.weight', 'layers.4.self_attn.o_proj.weight', 'layers.4.mlp.gate_proj.weight', 'layers.4.mlp.down_proj.weight', 'layers.4.mlp.up_proj.weight', 'layers.4.input_layernorm.weight', 'layers.4.post_attention_layernorm.weight', 'layers.5.self_attn.q_proj.weight', 'layers.5.self_attn.k_proj.weight', 'layers.5.self_attn.v_proj.weight', 'layers.5.self_attn.o_proj.weight', 'layers.5.mlp.gate_proj.weight', 'layers.5.mlp.down_proj.weight', 'layers.5.mlp.up_proj.weight', 'layers.5.input_layernorm.weight', 'layers.5.post_attention_layernorm.weight', 'layers.6.self_attn.q_proj.weight', 'layers.6.self_attn.k_proj.weight', 'layers.6.self_attn.v_proj.weight', 'layers.6.self_attn.o_proj.weight', 'layers.6.mlp.gate_proj.weight', 'layers.6.mlp.down_proj.weight', 'layers.6.mlp.up_proj.weight', 'layers.6.input_layernorm.weight', 'layers.6.post_attention_layernorm.weight', 'layers.7.self_attn.q_proj.weight', 'layers.7.self_attn.k_proj.weight', 'layers.7.self_attn.v_proj.weight', 'layers.7.self_attn.o_proj.weight', 'layers.7.mlp.gate_proj.weight', 'layers.7.mlp.down_proj.weight', 'layers.7.mlp.up_proj.weight', 'layers.7.input_layernorm.weight', 'layers.7.post_attention_layernorm.weight', 'layers.8.self_attn.q_proj.weight', 'layers.8.self_attn.k_proj.weight', 'layers.8.self_attn.v_proj.weight', 'layers.8.self_attn.o_proj.weight', 'layers.8.mlp.gate_proj.weight', 'layers.8.mlp.down_proj.weight', 'layers.8.mlp.up_proj.weight', 'layers.8.input_layernorm.weight', 'layers.8.post_attention_layernorm.weight', 'layers.9.self_attn.q_proj.weight', 'layers.9.self_attn.k_proj.weight', 'layers.9.self_attn.v_proj.weight', 'layers.9.self_attn.o_proj.weight', 'layers.9.mlp.gate_proj.weight', 'layers.9.mlp.down_proj.weight', 'layers.9.mlp.up_proj.weight', 'layers.9.input_layernorm.weight', 'layers.9.post_attention_layernorm.weight', 'layers.10.self_attn.q_proj.weight', 'layers.10.self_attn.k_proj.weight', 'layers.10.self_attn.v_proj.weight', 'layers.10.self_attn.o_proj.weight', 'layers.10.mlp.gate_proj.weight', 'layers.10.mlp.down_proj.weight', 'layers.10.mlp.up_proj.weight', 'layers.10.input_layernorm.weight', 'layers.10.post_attention_layernorm.weight', 'layers.11.self_attn.q_proj.weight', 'layers.11.self_attn.k_proj.weight', 'layers.11.self_attn.v_proj.weight', 'layers.11.self_attn.o_proj.weight', 'layers.11.mlp.gate_proj.weight', 'layers.11.mlp.down_proj.weight', 'layers.11.mlp.up_proj.weight', 'layers.11.input_layernorm.weight', 'layers.11.post_attention_layernorm.weight', 'layers.12.self_attn.q_proj.weight', 'layers.12.self_attn.k_proj.weight', 'layers.12.self_attn.v_proj.weight', 'layers.12.self_attn.o_proj.weight', 'layers.12.mlp.gate_proj.weight', 'layers.12.mlp.down_proj.weight', 'layers.12.mlp.up_proj.weight', 'layers.12.input_layernorm.weight', 'layers.12.post_attention_layernorm.weight', 'layers.13.self_attn.q_proj.weight', 'layers.13.self_attn.k_proj.weight', 'layers.13.self_attn.v_proj.weight', 'layers.13.self_attn.o_proj.weight', 'layers.13.mlp.gate_proj.weight', 'layers.13.mlp.down_proj.weight', 'layers.13.mlp.up_proj.weight', 'layers.13.input_layernorm.weight', 'layers.13.post_attention_layernorm.weight', 'layers.14.self_attn.q_proj.weight', 'layers.14.self_attn.k_proj.weight', 'layers.14.self_attn.v_proj.weight', 'layers.14.self_attn.o_proj.weight', 'layers.14.mlp.gate_proj.weight', 'layers.14.mlp.down_proj.weight', 'layers.14.mlp.up_proj.weight', 'layers.14.input_layernorm.weight', 'layers.14.post_attention_layernorm.weight', 'layers.15.self_attn.q_proj.weight', 'layers.15.self_attn.k_proj.weight', 'layers.15.self_attn.v_proj.weight', 'layers.15.self_attn.o_proj.weight', 'layers.15.mlp.gate_proj.weight', 'layers.15.mlp.down_proj.weight', 'layers.15.mlp.up_proj.weight', 'layers.15.input_layernorm.weight', 'layers.15.post_attention_layernorm.weight', 'layers.0.self_attn.rank_util.rank', 'layers.1.self_attn.rank_util.rank', 'layers.2.self_attn.rank_util.rank', 'layers.3.self_attn.rank_util.rank', 'layers.4.self_attn.rank_util.rank', 'layers.5.self_attn.rank_util.rank', 'layers.6.self_attn.rank_util.rank', 'layers.7.self_attn.rank_util.rank', 'layers.8.self_attn.rank_util.rank', 'layers.9.self_attn.rank_util.rank', 'layers.10.self_attn.rank_util.rank', 'layers.11.self_attn.rank_util.rank', 'layers.12.self_attn.rank_util.rank', 'layers.13.self_attn.rank_util.rank', 'layers.14.self_attn.rank_util.rank', 'layers.15.self_attn.rank_util.rank', 'rank_util.rank'] INFO:Neuron:Saving the neuron_config to llama3_compil -ed/ INFO:Neuron:Generating HLOs for the following models: - ['context_encoding_model', 'token_generation_model'][2025-07-26 00:33:00.384: I neuronx_distributed/paral -lel_layers/parallel_state.py:628] > initializing tensor model parallel with size 1 [2025-07-26 00:33:00.385: I neuronx_distributed/paral -lel_layers/parallel_state.py:629] > initializing pipeline model parallel with size 1 [2025-07-26 00:33:00.385: I neuronx_distributed/paral -lel_layers/parallel_state.py:630] > initializing context model parallel with size 1 [2025-07-26 00:33:00.385: I neuronx_distributed/paral -lel_layers/parallel_state.py:631] > initializing data parallel with size 1 [2025-07-26 00:33:00.385: I neuronx_distributed/paral -lel_layers/parallel_state.py:632] > initializing world size to 1 [2025-07-26 00:33:00.386: I neuronx_distributed/paral -lel_layers/parallel_state.py:379] [rank_0_pp0_tp0_dp0_cp0] Chosen Logic for replica groups ret_logic=, 'Ascending Ring PG Group')> [2025-07-26 00:33:00.386: I neuronx_distributed/paral -lel_layers/parallel_state.py:668] [rank_0_pp0_tp0_dp0_cp0] tp_groups: replica_groups.tp_groups=[[0]] [2025-07-26 00:33:00.386: I neuronx_distributed/paral -lel_layers/parallel_state.py:669] [rank_0_pp0_tp0_dp0_cp0] dp_groups: replica_groups.dp_groups=[[0]] [2025-07-26 00:33:00.386: I neuronx_distributed/paral -lel_layers/parallel_state.py:670] [rank_0_pp0_tp0_dp0_cp0] pp_groups: replica_groups.pp_groups=[[0]] [2025-07-26 00:33:00.386: I neuronx_distributed/paral -lel_layers/parallel_state.py:671] [rank_0_pp0_tp0_dp0_cp0] cp_groups: replica_groups.cp_groups=[[0]] [2025-07-26 00:33:00.386: I neuronx_distributed/paral -lel_layers/parallel_state.py:672] [rank_0_pp0_tp0_dp0_cp0] ep_model_groups: replica_groups.ep_model_groups=[[0]] [2025-07-26 00:33:00.386: I neuronx_distributed/paral -lel_layers/parallel_state.py:673] [rank_0_pp0_tp0_dp0_cp0] ep_data_groups: replica_groups.ep_data_groups=[[0]] INFO:Neuron:Generating 1 hlos for key: context_encodi -ng_model INFO:Neuron:Started loading module context_encoding_m -odel ERROR:__main__:Model compilation failed: init_model() - is not implemented ERROR:__main__:Compilation failed: init_model() is no -t implemented +lel_layers/parallel_state.py:668] [rank_0_pp-1_tp-1_dp-1_cp-1] tp_groups: replica_groups.tp_groups=[[0]] [2025-07-26 00:32:58.955: I neuronx_distributed/paral +lel_layers/parallel_state.py:669] [rank_0_pp-1_tp-1_dp-1_cp-1] dp_groups: replica_groups.dp_groups=[[0]] [2025-07-26 00:32:58.955: I neuronx_distributed/paral +lel_layers/parallel_state.py:670] [rank_0_pp-1_tp-1_dp-1_cp-1] pp_groups: replica_groups.pp_groups=[[0]] [2025-07-26 00:32:58.955: I neuronx_distributed/paral +lel_layers/parallel_state.py:671] [rank_0_pp-1_tp-1_dp-1_cp-1] cp_groups: replica_groups.cp_groups=[[0]] [2025-07-26 00:32:58.955: I neuronx_distributed/paral +lel_layers/parallel_state.py:672] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_model_groups: replica_groups.ep_model_groups=[[0]] [2025-07-26 00:32:58.955: I neuronx_distributed/paral +lel_layers/parallel_state.py:673] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_data_groups: replica_groups.ep_data_groups=[[0]] WARNING:**main**:Unexpected keys in state dict: ['emb +ed_tokens.weight', 'norm.weight', 'lm_head.weight', 'layers.0.self_attn.q_proj.weight', 'layers.0.self_attn.k_proj.weight', 'layers.0.self_attn.v_proj.weight', 'layers.0.self_attn.o_proj.weight', 'layers.0.mlp.gate_proj.weight', 'layers.0.mlp.down_proj.weight', 'layers.0.mlp.up_proj.weight', 'layers.0.input_layernorm.weight', 'layers.0.post_attention_layernorm.weight', 'layers.1.self_attn.q_proj.weight', 'layers.1.self_attn.k_proj.weight', 'layers.1.self_attn.v_proj.weight', 'layers.1.self_attn.o_proj.weight', 'layers.1.mlp.gate_proj.weight', 'layers.1.mlp.down_proj.weight', 'layers.1.mlp.up_proj.weight', 'layers.1.input_layernorm.weight', 'layers.1.post_attention_layernorm.weight', 'layers.2.self_attn.q_proj.weight', 'layers.2.self_attn.k_proj.weight', 'layers.2.self_attn.v_proj.weight', 'layers.2.self_attn.o_proj.weight', 'layers.2.mlp.gate_proj.weight', 'layers.2.mlp.down_proj.weight', 'layers.2.mlp.up_proj.weight', 'layers.2.input_layernorm.weight', 'layers.2.post_attention_layernorm.weight', 'layers.3.self_attn.q_proj.weight', 'layers.3.self_attn.k_proj.weight', 'layers.3.self_attn.v_proj.weight', 'layers.3.self_attn.o_proj.weight', 'layers.3.mlp.gate_proj.weight', 'layers.3.mlp.down_proj.weight', 'layers.3.mlp.up_proj.weight', 'layers.3.input_layernorm.weight', 'layers.3.post_attention_layernorm.weight', 'layers.4.self_attn.q_proj.weight', 'layers.4.self_attn.k_proj.weight', 'layers.4.self_attn.v_proj.weight', 'layers.4.self_attn.o_proj.weight', 'layers.4.mlp.gate_proj.weight', 'layers.4.mlp.down_proj.weight', 'layers.4.mlp.up_proj.weight', 'layers.4.input_layernorm.weight', 'layers.4.post_attention_layernorm.weight', 'layers.5.self_attn.q_proj.weight', 'layers.5.self_attn.k_proj.weight', 'layers.5.self_attn.v_proj.weight', 'layers.5.self_attn.o_proj.weight', 'layers.5.mlp.gate_proj.weight', 'layers.5.mlp.down_proj.weight', 'layers.5.mlp.up_proj.weight', 'layers.5.input_layernorm.weight', 'layers.5.post_attention_layernorm.weight', 'layers.6.self_attn.q_proj.weight', 'layers.6.self_attn.k_proj.weight', 'layers.6.self_attn.v_proj.weight', 'layers.6.self_attn.o_proj.weight', 'layers.6.mlp.gate_proj.weight', 'layers.6.mlp.down_proj.weight', 'layers.6.mlp.up_proj.weight', 'layers.6.input_layernorm.weight', 'layers.6.post_attention_layernorm.weight', 'layers.7.self_attn.q_proj.weight', 'layers.7.self_attn.k_proj.weight', 'layers.7.self_attn.v_proj.weight', 'layers.7.self_attn.o_proj.weight', 'layers.7.mlp.gate_proj.weight', 'layers.7.mlp.down_proj.weight', 'layers.7.mlp.up_proj.weight', 'layers.7.input_layernorm.weight', 'layers.7.post_attention_layernorm.weight', 'layers.8.self_attn.q_proj.weight', 'layers.8.self_attn.k_proj.weight', 'layers.8.self_attn.v_proj.weight', 'layers.8.self_attn.o_proj.weight', 'layers.8.mlp.gate_proj.weight', 'layers.8.mlp.down_proj.weight', 'layers.8.mlp.up_proj.weight', 'layers.8.input_layernorm.weight', 'layers.8.post_attention_layernorm.weight', 'layers.9.self_attn.q_proj.weight', 'layers.9.self_attn.k_proj.weight', 'layers.9.self_attn.v_proj.weight', 'layers.9.self_attn.o_proj.weight', 'layers.9.mlp.gate_proj.weight', 'layers.9.mlp.down_proj.weight', 'layers.9.mlp.up_proj.weight', 'layers.9.input_layernorm.weight', 'layers.9.post_attention_layernorm.weight', 'layers.10.self_attn.q_proj.weight', 'layers.10.self_attn.k_proj.weight', 'layers.10.self_attn.v_proj.weight', 'layers.10.self_attn.o_proj.weight', 'layers.10.mlp.gate_proj.weight', 'layers.10.mlp.down_proj.weight', 'layers.10.mlp.up_proj.weight', 'layers.10.input_layernorm.weight', 'layers.10.post_attention_layernorm.weight', 'layers.11.self_attn.q_proj.weight', 'layers.11.self_attn.k_proj.weight', 'layers.11.self_attn.v_proj.weight', 'layers.11.self_attn.o_proj.weight', 'layers.11.mlp.gate_proj.weight', 'layers.11.mlp.down_proj.weight', 'layers.11.mlp.up_proj.weight', 'layers.11.input_layernorm.weight', 'layers.11.post_attention_layernorm.weight', 'layers.12.self_attn.q_proj.weight', 'layers.12.self_attn.k_proj.weight', 'layers.12.self_attn.v_proj.weight', 'layers.12.self_attn.o_proj.weight', 'layers.12.mlp.gate_proj.weight', 'layers.12.mlp.down_proj.weight', 'layers.12.mlp.up_proj.weight', 'layers.12.input_layernorm.weight', 'layers.12.post_attention_layernorm.weight', 'layers.13.self_attn.q_proj.weight', 'layers.13.self_attn.k_proj.weight', 'layers.13.self_attn.v_proj.weight', 'layers.13.self_attn.o_proj.weight', 'layers.13.mlp.gate_proj.weight', 'layers.13.mlp.down_proj.weight', 'layers.13.mlp.up_proj.weight', 'layers.13.input_layernorm.weight', 'layers.13.post_attention_layernorm.weight', 'layers.14.self_attn.q_proj.weight', 'layers.14.self_attn.k_proj.weight', 'layers.14.self_attn.v_proj.weight', 'layers.14.self_attn.o_proj.weight', 'layers.14.mlp.gate_proj.weight', 'layers.14.mlp.down_proj.weight', 'layers.14.mlp.up_proj.weight', 'layers.14.input_layernorm.weight', 'layers.14.post_attention_layernorm.weight', 'layers.15.self_attn.q_proj.weight', 'layers.15.self_attn.k_proj.weight', 'layers.15.self_attn.v_proj.weight', 'layers.15.self_attn.o_proj.weight', 'layers.15.mlp.gate_proj.weight', 'layers.15.mlp.down_proj.weight', 'layers.15.mlp.up_proj.weight', 'layers.15.input_layernorm.weight', 'layers.15.post_attention_layernorm.weight', 'layers.0.self_attn.rank_util.rank', 'layers.1.self_attn.rank_util.rank', 'layers.2.self_attn.rank_util.rank', 'layers.3.self_attn.rank_util.rank', 'layers.4.self_attn.rank_util.rank', 'layers.5.self_attn.rank_util.rank', 'layers.6.self_attn.rank_util.rank', 'layers.7.self_attn.rank_util.rank', 'layers.8.self_attn.rank_util.rank', 'layers.9.self_attn.rank_util.rank', 'layers.10.self_attn.rank_util.rank', 'layers.11.self_attn.rank_util.rank', 'layers.12.self_attn.rank_util.rank', 'layers.13.self_attn.rank_util.rank', 'layers.14.self_attn.rank_util.rank', 'layers.15.self_attn.rank_util.rank', 'rank_util.rank'] INFO:Neuron:Saving the neuron_config to llama3_compil +ed/ INFO:Neuron:Generating HLOs for the following models: +['context_encoding_model', 'token_generation_model'][2025-07-26 00:33:00.384: I neuronx_distributed/paral lel_layers/parallel_state.py:628] > initializing tensor model parallel with size 1 [2025-07-26 00:33:00.385: I neuronx_distributed/paral +lel_layers/parallel_state.py:629] > initializing pipeline model parallel with size 1 [2025-07-26 00:33:00.385: I neuronx_distributed/paral +lel_layers/parallel_state.py:630] > initializing context model parallel with size 1 [2025-07-26 00:33:00.385: I neuronx_distributed/paral +lel_layers/parallel_state.py:631] > initializing data parallel with size 1 [2025-07-26 00:33:00.385: I neuronx_distributed/paral +lel_layers/parallel_state.py:632] > initializing world size to 1 [2025-07-26 00:33:00.386: I neuronx_distributed/paral +lel_layers/parallel_state.py:379] [rank_0_pp0_tp0_dp0_cp0] Chosen Logic for replica groups ret_logic=, 'Ascending Ring PG Group')> [2025-07-26 00:33:00.386: I neuronx_distributed/paral +lel_layers/parallel_state.py:668] [rank_0_pp0_tp0_dp0_cp0] tp_groups: replica_groups.tp_groups=[[0]] [2025-07-26 00:33:00.386: I neuronx_distributed/paral +lel_layers/parallel_state.py:669] [rank_0_pp0_tp0_dp0_cp0] dp_groups: replica_groups.dp_groups=[[0]] [2025-07-26 00:33:00.386: I neuronx_distributed/paral +lel_layers/parallel_state.py:670] [rank_0_pp0_tp0_dp0_cp0] pp_groups: replica_groups.pp_groups=[[0]] [2025-07-26 00:33:00.386: I neuronx_distributed/paral +lel_layers/parallel_state.py:671] [rank_0_pp0_tp0_dp0_cp0] cp_groups: replica_groups.cp_groups=[[0]] [2025-07-26 00:33:00.386: I neuronx_distributed/paral +lel_layers/parallel_state.py:672] [rank_0_pp0_tp0_dp0_cp0] ep_model_groups: replica_groups.ep_model_groups=[[0]] [2025-07-26 00:33:00.386: I neuronx_distributed/paral +lel_layers/parallel_state.py:673] [rank_0_pp0_tp0_dp0_cp0] ep_data_groups: replica_groups.ep_data_groups=[[0]] INFO:Neuron:Generating 1 hlos for key: context_encodi +ng_model INFO:Neuron:Started loading module context_encoding_m +odel ERROR:**main**:Model compilation failed: init_model() +is not implemented ERROR:**main**:Compilation failed: init_model() is no +t implemented I need to implement the init_model method as well. Let me add this: Accepted edits to modeling_llama3.py - Now let's try compiling again: Command -python compile_model.py --checkpoint_path ./llama3_neuron_checkpoint --output_path ./llama3_compiled --batch_size 1 --seq_len 128 +python compile*model.py --checkpoint_path ./llama3_neuron_checkpoint --output_path ./llama3_compiled --batch_size 1 --seq_len 128 /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributedInference -/src/neuronx_distributed_inference/modules/attention/utils.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.custom_c -alls import neuron_cumsum /home/ec2-user/NeuronxSDK/NeuronxDistributedInference -/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference -/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference -/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_ -llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_ -llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_ -llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.attention_base import NeuronAttentionBase [2025-07-26 00:33:51.360: I neuronx_distributed/paral -lel_layers/parallel_state.py:628] > initializing tensor model parallel with size 1 [2025-07-26 00:33:51.361: I neuronx_distributed/paral -lel_layers/parallel_state.py:629] > initializing pipeline model parallel with size 1 [2025-07-26 00:33:51.361: I neuronx_distributed/paral -lel_layers/parallel_state.py:630] > initializing context model parallel with size 1 [2025-07-26 00:33:51.361: I neuronx_distributed/paral -lel_layers/parallel_state.py:631] > initializing data parallel with size 1 [2025-07-26 00:33:51.361: I neuronx_distributed/paral -lel_layers/parallel_state.py:632] > initializing world size to 1 2025-07-26 00:33:51.000366: 345240 INFO ||NEURON_CC -_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/neuronxcc-2.19.8089.0+8ab9f450/MODULE_13537832482499839909+e30acd3a/model.neff [2025-07-26 00:33:51.406: I neuronx_distributed/paral +/src/neuronx_distributed_inference/modules/attention/utils.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.custom_c +alls import neuron_cumsum /home/ec2-user/NeuronxSDK/NeuronxDistributedInference +/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference +/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference +/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx* +llama3/modeling*llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx* +llama3/modeling*llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx* +llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.attention_base import NeuronAttentionBase [2025-07-26 00:33:51.360: I neuronx_distributed/paral +lel_layers/parallel_state.py:628] > initializing tensor model parallel with size 1 [2025-07-26 00:33:51.361: I neuronx_distributed/paral +lel_layers/parallel_state.py:629] > initializing pipeline model parallel with size 1 [2025-07-26 00:33:51.361: I neuronx_distributed/paral +lel_layers/parallel_state.py:630] > initializing context model parallel with size 1 [2025-07-26 00:33:51.361: I neuronx_distributed/paral +lel_layers/parallel_state.py:631] > initializing data parallel with size 1 [2025-07-26 00:33:51.361: I neuronx_distributed/paral +lel_layers/parallel_state.py:632] > initializing world size to 1 2025-07-26 00:33:51.000366: 345240 INFO ||NEURON_CC +\_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/neuronxcc-2.19.8089.0+8ab9f450/MODULE_13537832482499839909+e30acd3a/model.neff [2025-07-26 00:33:51.406: I neuronx_distributed/paral lel_layers/parallel_state.py:379] [rank_0_pp-1_tp-1_dp-1_cp-1] Chosen Logic for replica groups ret_logic=, 'Ascending Ring PG Group')> [2025-07-26 00:33:51.407: I neuronx_distributed/paral -lel_layers/parallel_state.py:668] [rank_0_pp-1_tp-1_dp-1_cp-1] tp_groups: replica_groups.tp_groups=[[0]] [2025-07-26 00:33:51.407: I neuronx_distributed/paral -lel_layers/parallel_state.py:669] [rank_0_pp-1_tp-1_dp-1_cp-1] dp_groups: replica_groups.dp_groups=[[0]] [2025-07-26 00:33:51.407: I neuronx_distributed/paral -lel_layers/parallel_state.py:670] [rank_0_pp-1_tp-1_dp-1_cp-1] pp_groups: replica_groups.pp_groups=[[0]] [2025-07-26 00:33:51.407: I neuronx_distributed/paral -lel_layers/parallel_state.py:671] [rank_0_pp-1_tp-1_dp-1_cp-1] cp_groups: replica_groups.cp_groups=[[0]] [2025-07-26 00:33:51.407: I neuronx_distributed/paral -lel_layers/parallel_state.py:672] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_model_groups: replica_groups.ep_model_groups=[[0]] [2025-07-26 00:33:51.407: I neuronx_distributed/paral -lel_layers/parallel_state.py:673] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_data_groups: replica_groups.ep_data_groups=[[0]] WARNING:__main__:Unexpected keys in state dict: ['emb -ed_tokens.weight', 'norm.weight', 'lm_head.weight', 'layers.0.self_attn.q_proj.weight', 'layers.0.self_attn.k_proj.weight', 'layers.0.self_attn.v_proj.weight', 'layers.0.self_attn.o_proj.weight', 'layers.0.mlp.gate_proj.weight', 'layers.0.mlp.down_proj.weight', 'layers.0.mlp.up_proj.weight', 'layers.0.input_layernorm.weight', 'layers.0.post_attention_layernorm.weight', 'layers.1.self_attn.q_proj.weight', 'layers.1.self_attn.k_proj.weight', 'layers.1.self_attn.v_proj.weight', 'layers.1.self_attn.o_proj.weight', 'layers.1.mlp.gate_proj.weight', 'layers.1.mlp.down_proj.weight', 'layers.1.mlp.up_proj.weight', 'layers.1.input_layernorm.weight', 'layers.1.post_attention_layernorm.weight', 'layers.2.self_attn.q_proj.weight', 'layers.2.self_attn.k_proj.weight', 'layers.2.self_attn.v_proj.weight', 'layers.2.self_attn.o_proj.weight', 'layers.2.mlp.gate_proj.weight', 'layers.2.mlp.down_proj.weight', 'layers.2.mlp.up_proj.weight', 'layers.2.input_layernorm.weight', 'layers.2.post_attention_layernorm.weight', 'layers.3.self_attn.q_proj.weight', 'layers.3.self_attn.k_proj.weight', 'layers.3.self_attn.v_proj.weight', 'layers.3.self_attn.o_proj.weight', 'layers.3.mlp.gate_proj.weight', 'layers.3.mlp.down_proj.weight', 'layers.3.mlp.up_proj.weight', 'layers.3.input_layernorm.weight', 'layers.3.post_attention_layernorm.weight', 'layers.4.self_attn.q_proj.weight', 'layers.4.self_attn.k_proj.weight', 'layers.4.self_attn.v_proj.weight', 'layers.4.self_attn.o_proj.weight', 'layers.4.mlp.gate_proj.weight', 'layers.4.mlp.down_proj.weight', 'layers.4.mlp.up_proj.weight', 'layers.4.input_layernorm.weight', 'layers.4.post_attention_layernorm.weight', 'layers.5.self_attn.q_proj.weight', 'layers.5.self_attn.k_proj.weight', 'layers.5.self_attn.v_proj.weight', 'layers.5.self_attn.o_proj.weight', 'layers.5.mlp.gate_proj.weight', 'layers.5.mlp.down_proj.weight', 'layers.5.mlp.up_proj.weight', 'layers.5.input_layernorm.weight', 'layers.5.post_attention_layernorm.weight', 'layers.6.self_attn.q_proj.weight', 'layers.6.self_attn.k_proj.weight', 'layers.6.self_attn.v_proj.weight', 'layers.6.self_attn.o_proj.weight', 'layers.6.mlp.gate_proj.weight', 'layers.6.mlp.down_proj.weight', 'layers.6.mlp.up_proj.weight', 'layers.6.input_layernorm.weight', 'layers.6.post_attention_layernorm.weight', 'layers.7.self_attn.q_proj.weight', 'layers.7.self_attn.k_proj.weight', 'layers.7.self_attn.v_proj.weight', 'layers.7.self_attn.o_proj.weight', 'layers.7.mlp.gate_proj.weight', 'layers.7.mlp.down_proj.weight', 'layers.7.mlp.up_proj.weight', 'layers.7.input_layernorm.weight', 'layers.7.post_attention_layernorm.weight', 'layers.8.self_attn.q_proj.weight', 'layers.8.self_attn.k_proj.weight', 'layers.8.self_attn.v_proj.weight', 'layers.8.self_attn.o_proj.weight', 'layers.8.mlp.gate_proj.weight', 'layers.8.mlp.down_proj.weight', 'layers.8.mlp.up_proj.weight', 'layers.8.input_layernorm.weight', 'layers.8.post_attention_layernorm.weight', 'layers.9.self_attn.q_proj.weight', 'layers.9.self_attn.k_proj.weight', 'layers.9.self_attn.v_proj.weight', 'layers.9.self_attn.o_proj.weight', 'layers.9.mlp.gate_proj.weight', 'layers.9.mlp.down_proj.weight', 'layers.9.mlp.up_proj.weight', 'layers.9.input_layernorm.weight', 'layers.9.post_attention_layernorm.weight', 'layers.10.self_attn.q_proj.weight', 'layers.10.self_attn.k_proj.weight', 'layers.10.self_attn.v_proj.weight', 'layers.10.self_attn.o_proj.weight', 'layers.10.mlp.gate_proj.weight', 'layers.10.mlp.down_proj.weight', 'layers.10.mlp.up_proj.weight', 'layers.10.input_layernorm.weight', 'layers.10.post_attention_layernorm.weight', 'layers.11.self_attn.q_proj.weight', 'layers.11.self_attn.k_proj.weight', 'layers.11.self_attn.v_proj.weight', 'layers.11.self_attn.o_proj.weight', 'layers.11.mlp.gate_proj.weight', 'layers.11.mlp.down_proj.weight', 'layers.11.mlp.up_proj.weight', 'layers.11.input_layernorm.weight', 'layers.11.post_attention_layernorm.weight', 'layers.12.self_attn.q_proj.weight', 'layers.12.self_attn.k_proj.weight', 'layers.12.self_attn.v_proj.weight', 'layers.12.self_attn.o_proj.weight', 'layers.12.mlp.gate_proj.weight', 'layers.12.mlp.down_proj.weight', 'layers.12.mlp.up_proj.weight', 'layers.12.input_layernorm.weight', 'layers.12.post_attention_layernorm.weight', 'layers.13.self_attn.q_proj.weight', 'layers.13.self_attn.k_proj.weight', 'layers.13.self_attn.v_proj.weight', 'layers.13.self_attn.o_proj.weight', 'layers.13.mlp.gate_proj.weight', 'layers.13.mlp.down_proj.weight', 'layers.13.mlp.up_proj.weight', 'layers.13.input_layernorm.weight', 'layers.13.post_attention_layernorm.weight', 'layers.14.self_attn.q_proj.weight', 'layers.14.self_attn.k_proj.weight', 'layers.14.self_attn.v_proj.weight', 'layers.14.self_attn.o_proj.weight', 'layers.14.mlp.gate_proj.weight', 'layers.14.mlp.down_proj.weight', 'layers.14.mlp.up_proj.weight', 'layers.14.input_layernorm.weight', 'layers.14.post_attention_layernorm.weight', 'layers.15.self_attn.q_proj.weight', 'layers.15.self_attn.k_proj.weight', 'layers.15.self_attn.v_proj.weight', 'layers.15.self_attn.o_proj.weight', 'layers.15.mlp.gate_proj.weight', 'layers.15.mlp.down_proj.weight', 'layers.15.mlp.up_proj.weight', 'layers.15.input_layernorm.weight', 'layers.15.post_attention_layernorm.weight', 'layers.0.self_attn.rank_util.rank', 'layers.1.self_attn.rank_util.rank', 'layers.2.self_attn.rank_util.rank', 'layers.3.self_attn.rank_util.rank', 'layers.4.self_attn.rank_util.rank', 'layers.5.self_attn.rank_util.rank', 'layers.6.self_attn.rank_util.rank', 'layers.7.self_attn.rank_util.rank', 'layers.8.self_attn.rank_util.rank', 'layers.9.self_attn.rank_util.rank', 'layers.10.self_attn.rank_util.rank', 'layers.11.self_attn.rank_util.rank', 'layers.12.self_attn.rank_util.rank', 'layers.13.self_attn.rank_util.rank', 'layers.14.self_attn.rank_util.rank', 'layers.15.self_attn.rank_util.rank', 'rank_util.rank'] INFO:Neuron:Saving the neuron_config to llama3_compil -ed/ INFO:Neuron:Generating HLOs for the following models: - ['context_encoding_model', 'token_generation_model'][2025-07-26 00:33:52.835: I neuronx_distributed/paral -lel_layers/parallel_state.py:628] > initializing tensor model parallel with size 1 [2025-07-26 00:33:52.836: I neuronx_distributed/paral -lel_layers/parallel_state.py:629] > initializing pipeline model parallel with size 1 [2025-07-26 00:33:52.836: I neuronx_distributed/paral -lel_layers/parallel_state.py:630] > initializing context model parallel with size 1 [2025-07-26 00:33:52.836: I neuronx_distributed/paral -lel_layers/parallel_state.py:631] > initializing data parallel with size 1 [2025-07-26 00:33:52.836: I neuronx_distributed/paral -lel_layers/parallel_state.py:632] > initializing world size to 1 [2025-07-26 00:33:52.837: I neuronx_distributed/paral -lel_layers/parallel_state.py:379] [rank_0_pp0_tp0_dp0_cp0] Chosen Logic for replica groups ret_logic=, 'Ascending Ring PG Group')> [2025-07-26 00:33:52.837: I neuronx_distributed/paral -lel_layers/parallel_state.py:668] [rank_0_pp0_tp0_dp0_cp0] tp_groups: replica_groups.tp_groups=[[0]] [2025-07-26 00:33:52.837: I neuronx_distributed/paral -lel_layers/parallel_state.py:669] [rank_0_pp0_tp0_dp0_cp0] dp_groups: replica_groups.dp_groups=[[0]] [2025-07-26 00:33:52.837: I neuronx_distributed/paral -lel_layers/parallel_state.py:670] [rank_0_pp0_tp0_dp0_cp0] pp_groups: replica_groups.pp_groups=[[0]] [2025-07-26 00:33:52.837: I neuronx_distributed/paral -lel_layers/parallel_state.py:671] [rank_0_pp0_tp0_dp0_cp0] cp_groups: replica_groups.cp_groups=[[0]] [2025-07-26 00:33:52.837: I neuronx_distributed/paral -lel_layers/parallel_state.py:672] [rank_0_pp0_tp0_dp0_cp0] ep_model_groups: replica_groups.ep_model_groups=[[0]] [2025-07-26 00:33:52.837: I neuronx_distributed/paral -lel_layers/parallel_state.py:673] [rank_0_pp0_tp0_dp0_cp0] ep_data_groups: replica_groups.ep_data_groups=[[0]] INFO:Neuron:Generating 1 hlos for key: context_encodi -ng_model INFO:Neuron:Started loading module context_encoding_m -odel ERROR:__main__:Model compilation failed: NeuronLlama3 -Model.init_model() takes 1 positional argument but 2 were given ERROR:__main__:Compilation failed: NeuronLlama3Model. -init_model() takes 1 positional argument but 2 were given +lel_layers/parallel_state.py:668] [rank_0_pp-1_tp-1_dp-1_cp-1] tp_groups: replica_groups.tp_groups=[[0]] [2025-07-26 00:33:51.407: I neuronx_distributed/paral +lel_layers/parallel_state.py:669] [rank_0_pp-1_tp-1_dp-1_cp-1] dp_groups: replica_groups.dp_groups=[[0]] [2025-07-26 00:33:51.407: I neuronx_distributed/paral +lel_layers/parallel_state.py:670] [rank_0_pp-1_tp-1_dp-1_cp-1] pp_groups: replica_groups.pp_groups=[[0]] [2025-07-26 00:33:51.407: I neuronx_distributed/paral +lel_layers/parallel_state.py:671] [rank_0_pp-1_tp-1_dp-1_cp-1] cp_groups: replica_groups.cp_groups=[[0]] [2025-07-26 00:33:51.407: I neuronx_distributed/paral +lel_layers/parallel_state.py:672] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_model_groups: replica_groups.ep_model_groups=[[0]] [2025-07-26 00:33:51.407: I neuronx_distributed/paral +lel_layers/parallel_state.py:673] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_data_groups: replica_groups.ep_data_groups=[[0]] WARNING:**main**:Unexpected keys in state dict: ['emb +ed_tokens.weight', 'norm.weight', 'lm_head.weight', 'layers.0.self_attn.q_proj.weight', 'layers.0.self_attn.k_proj.weight', 'layers.0.self_attn.v_proj.weight', 'layers.0.self_attn.o_proj.weight', 'layers.0.mlp.gate_proj.weight', 'layers.0.mlp.down_proj.weight', 'layers.0.mlp.up_proj.weight', 'layers.0.input_layernorm.weight', 'layers.0.post_attention_layernorm.weight', 'layers.1.self_attn.q_proj.weight', 'layers.1.self_attn.k_proj.weight', 'layers.1.self_attn.v_proj.weight', 'layers.1.self_attn.o_proj.weight', 'layers.1.mlp.gate_proj.weight', 'layers.1.mlp.down_proj.weight', 'layers.1.mlp.up_proj.weight', 'layers.1.input_layernorm.weight', 'layers.1.post_attention_layernorm.weight', 'layers.2.self_attn.q_proj.weight', 'layers.2.self_attn.k_proj.weight', 'layers.2.self_attn.v_proj.weight', 'layers.2.self_attn.o_proj.weight', 'layers.2.mlp.gate_proj.weight', 'layers.2.mlp.down_proj.weight', 'layers.2.mlp.up_proj.weight', 'layers.2.input_layernorm.weight', 'layers.2.post_attention_layernorm.weight', 'layers.3.self_attn.q_proj.weight', 'layers.3.self_attn.k_proj.weight', 'layers.3.self_attn.v_proj.weight', 'layers.3.self_attn.o_proj.weight', 'layers.3.mlp.gate_proj.weight', 'layers.3.mlp.down_proj.weight', 'layers.3.mlp.up_proj.weight', 'layers.3.input_layernorm.weight', 'layers.3.post_attention_layernorm.weight', 'layers.4.self_attn.q_proj.weight', 'layers.4.self_attn.k_proj.weight', 'layers.4.self_attn.v_proj.weight', 'layers.4.self_attn.o_proj.weight', 'layers.4.mlp.gate_proj.weight', 'layers.4.mlp.down_proj.weight', 'layers.4.mlp.up_proj.weight', 'layers.4.input_layernorm.weight', 'layers.4.post_attention_layernorm.weight', 'layers.5.self_attn.q_proj.weight', 'layers.5.self_attn.k_proj.weight', 'layers.5.self_attn.v_proj.weight', 'layers.5.self_attn.o_proj.weight', 'layers.5.mlp.gate_proj.weight', 'layers.5.mlp.down_proj.weight', 'layers.5.mlp.up_proj.weight', 'layers.5.input_layernorm.weight', 'layers.5.post_attention_layernorm.weight', 'layers.6.self_attn.q_proj.weight', 'layers.6.self_attn.k_proj.weight', 'layers.6.self_attn.v_proj.weight', 'layers.6.self_attn.o_proj.weight', 'layers.6.mlp.gate_proj.weight', 'layers.6.mlp.down_proj.weight', 'layers.6.mlp.up_proj.weight', 'layers.6.input_layernorm.weight', 'layers.6.post_attention_layernorm.weight', 'layers.7.self_attn.q_proj.weight', 'layers.7.self_attn.k_proj.weight', 'layers.7.self_attn.v_proj.weight', 'layers.7.self_attn.o_proj.weight', 'layers.7.mlp.gate_proj.weight', 'layers.7.mlp.down_proj.weight', 'layers.7.mlp.up_proj.weight', 'layers.7.input_layernorm.weight', 'layers.7.post_attention_layernorm.weight', 'layers.8.self_attn.q_proj.weight', 'layers.8.self_attn.k_proj.weight', 'layers.8.self_attn.v_proj.weight', 'layers.8.self_attn.o_proj.weight', 'layers.8.mlp.gate_proj.weight', 'layers.8.mlp.down_proj.weight', 'layers.8.mlp.up_proj.weight', 'layers.8.input_layernorm.weight', 'layers.8.post_attention_layernorm.weight', 'layers.9.self_attn.q_proj.weight', 'layers.9.self_attn.k_proj.weight', 'layers.9.self_attn.v_proj.weight', 'layers.9.self_attn.o_proj.weight', 'layers.9.mlp.gate_proj.weight', 'layers.9.mlp.down_proj.weight', 'layers.9.mlp.up_proj.weight', 'layers.9.input_layernorm.weight', 'layers.9.post_attention_layernorm.weight', 'layers.10.self_attn.q_proj.weight', 'layers.10.self_attn.k_proj.weight', 'layers.10.self_attn.v_proj.weight', 'layers.10.self_attn.o_proj.weight', 'layers.10.mlp.gate_proj.weight', 'layers.10.mlp.down_proj.weight', 'layers.10.mlp.up_proj.weight', 'layers.10.input_layernorm.weight', 'layers.10.post_attention_layernorm.weight', 'layers.11.self_attn.q_proj.weight', 'layers.11.self_attn.k_proj.weight', 'layers.11.self_attn.v_proj.weight', 'layers.11.self_attn.o_proj.weight', 'layers.11.mlp.gate_proj.weight', 'layers.11.mlp.down_proj.weight', 'layers.11.mlp.up_proj.weight', 'layers.11.input_layernorm.weight', 'layers.11.post_attention_layernorm.weight', 'layers.12.self_attn.q_proj.weight', 'layers.12.self_attn.k_proj.weight', 'layers.12.self_attn.v_proj.weight', 'layers.12.self_attn.o_proj.weight', 'layers.12.mlp.gate_proj.weight', 'layers.12.mlp.down_proj.weight', 'layers.12.mlp.up_proj.weight', 'layers.12.input_layernorm.weight', 'layers.12.post_attention_layernorm.weight', 'layers.13.self_attn.q_proj.weight', 'layers.13.self_attn.k_proj.weight', 'layers.13.self_attn.v_proj.weight', 'layers.13.self_attn.o_proj.weight', 'layers.13.mlp.gate_proj.weight', 'layers.13.mlp.down_proj.weight', 'layers.13.mlp.up_proj.weight', 'layers.13.input_layernorm.weight', 'layers.13.post_attention_layernorm.weight', 'layers.14.self_attn.q_proj.weight', 'layers.14.self_attn.k_proj.weight', 'layers.14.self_attn.v_proj.weight', 'layers.14.self_attn.o_proj.weight', 'layers.14.mlp.gate_proj.weight', 'layers.14.mlp.down_proj.weight', 'layers.14.mlp.up_proj.weight', 'layers.14.input_layernorm.weight', 'layers.14.post_attention_layernorm.weight', 'layers.15.self_attn.q_proj.weight', 'layers.15.self_attn.k_proj.weight', 'layers.15.self_attn.v_proj.weight', 'layers.15.self_attn.o_proj.weight', 'layers.15.mlp.gate_proj.weight', 'layers.15.mlp.down_proj.weight', 'layers.15.mlp.up_proj.weight', 'layers.15.input_layernorm.weight', 'layers.15.post_attention_layernorm.weight', 'layers.0.self_attn.rank_util.rank', 'layers.1.self_attn.rank_util.rank', 'layers.2.self_attn.rank_util.rank', 'layers.3.self_attn.rank_util.rank', 'layers.4.self_attn.rank_util.rank', 'layers.5.self_attn.rank_util.rank', 'layers.6.self_attn.rank_util.rank', 'layers.7.self_attn.rank_util.rank', 'layers.8.self_attn.rank_util.rank', 'layers.9.self_attn.rank_util.rank', 'layers.10.self_attn.rank_util.rank', 'layers.11.self_attn.rank_util.rank', 'layers.12.self_attn.rank_util.rank', 'layers.13.self_attn.rank_util.rank', 'layers.14.self_attn.rank_util.rank', 'layers.15.self_attn.rank_util.rank', 'rank_util.rank'] INFO:Neuron:Saving the neuron_config to llama3_compil +ed/ INFO:Neuron:Generating HLOs for the following models: +['context_encoding_model', 'token_generation_model'][2025-07-26 00:33:52.835: I neuronx_distributed/paral lel_layers/parallel_state.py:628] > initializing tensor model parallel with size 1 [2025-07-26 00:33:52.836: I neuronx_distributed/paral +lel_layers/parallel_state.py:629] > initializing pipeline model parallel with size 1 [2025-07-26 00:33:52.836: I neuronx_distributed/paral +lel_layers/parallel_state.py:630] > initializing context model parallel with size 1 [2025-07-26 00:33:52.836: I neuronx_distributed/paral +lel_layers/parallel_state.py:631] > initializing data parallel with size 1 [2025-07-26 00:33:52.836: I neuronx_distributed/paral +lel_layers/parallel_state.py:632] > initializing world size to 1 [2025-07-26 00:33:52.837: I neuronx_distributed/paral +lel_layers/parallel_state.py:379] [rank_0_pp0_tp0_dp0_cp0] Chosen Logic for replica groups ret_logic=, 'Ascending Ring PG Group')> [2025-07-26 00:33:52.837: I neuronx_distributed/paral +lel_layers/parallel_state.py:668] [rank_0_pp0_tp0_dp0_cp0] tp_groups: replica_groups.tp_groups=[[0]] [2025-07-26 00:33:52.837: I neuronx_distributed/paral +lel_layers/parallel_state.py:669] [rank_0_pp0_tp0_dp0_cp0] dp_groups: replica_groups.dp_groups=[[0]] [2025-07-26 00:33:52.837: I neuronx_distributed/paral +lel_layers/parallel_state.py:670] [rank_0_pp0_tp0_dp0_cp0] pp_groups: replica_groups.pp_groups=[[0]] [2025-07-26 00:33:52.837: I neuronx_distributed/paral +lel_layers/parallel_state.py:671] [rank_0_pp0_tp0_dp0_cp0] cp_groups: replica_groups.cp_groups=[[0]] [2025-07-26 00:33:52.837: I neuronx_distributed/paral +lel_layers/parallel_state.py:672] [rank_0_pp0_tp0_dp0_cp0] ep_model_groups: replica_groups.ep_model_groups=[[0]] [2025-07-26 00:33:52.837: I neuronx_distributed/paral +lel_layers/parallel_state.py:673] [rank_0_pp0_tp0_dp0_cp0] ep_data_groups: replica_groups.ep_data_groups=[[0]] INFO:Neuron:Generating 1 hlos for key: context_encodi +ng_model INFO:Neuron:Started loading module context_encoding_m +odel ERROR:**main**:Model compilation failed: NeuronLlama3 +Model.init_model() takes 1 positional argument but 2 were given ERROR:**main**:Compilation failed: NeuronLlama3Model. +init_model() takes 1 positional argument but 2 were given The init_model method is being called with an argument. Let me fix this: Accepted edits to modeling_llama3.py - Let's try compiling again: Command -python compile_model.py --checkpoint_path ./llama3_neuron_checkpoint --output_path ./llama3_compiled --batch_size 1 --seq_len 128 +python compile*model.py --checkpoint_path ./llama3_neuron_checkpoint --output_path ./llama3_compiled --batch_size 1 --seq_len 128 /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributedInference -/src/neuronx_distributed_inference/modules/attention/utils.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.custom_c -alls import neuron_cumsum /home/ec2-user/NeuronxSDK/NeuronxDistributedInference -/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference -/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference -/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_ -llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_ -llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_ -llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.attention_base import NeuronAttentionBase [2025-07-26 00:34:25.960: I neuronx_distributed/paral -lel_layers/parallel_state.py:628] > initializing tensor model parallel with size 1 [2025-07-26 00:34:25.960: I neuronx_distributed/paral -lel_layers/parallel_state.py:629] > initializing pipeline model parallel with size 1 [2025-07-26 00:34:25.960: I neuronx_distributed/paral -lel_layers/parallel_state.py:630] > initializing context model parallel with size 1 [2025-07-26 00:34:25.960: I neuronx_distributed/paral -lel_layers/parallel_state.py:631] > initializing data parallel with size 1 [2025-07-26 00:34:25.960: I neuronx_distributed/paral -lel_layers/parallel_state.py:632] > initializing world size to 1 2025-07-26 00:34:25.000966: 346389 INFO ||NEURON_CC -_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/neuronxcc-2.19.8089.0+8ab9f450/MODULE_13537832482499839909+e30acd3a/model.neff [2025-07-26 00:34:26.003: I neuronx_distributed/paral +/src/neuronx_distributed_inference/modules/attention/utils.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.custom_c +alls import neuron_cumsum /home/ec2-user/NeuronxSDK/NeuronxDistributedInference +/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference +/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference +/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx* +llama3/modeling*llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx* +llama3/modeling*llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx* +llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.attention_base import NeuronAttentionBase [2025-07-26 00:34:25.960: I neuronx_distributed/paral +lel_layers/parallel_state.py:628] > initializing tensor model parallel with size 1 [2025-07-26 00:34:25.960: I neuronx_distributed/paral +lel_layers/parallel_state.py:629] > initializing pipeline model parallel with size 1 [2025-07-26 00:34:25.960: I neuronx_distributed/paral +lel_layers/parallel_state.py:630] > initializing context model parallel with size 1 [2025-07-26 00:34:25.960: I neuronx_distributed/paral +lel_layers/parallel_state.py:631] > initializing data parallel with size 1 [2025-07-26 00:34:25.960: I neuronx_distributed/paral +lel_layers/parallel_state.py:632] > initializing world size to 1 2025-07-26 00:34:25.000966: 346389 INFO ||NEURON_CC +\_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/neuronxcc-2.19.8089.0+8ab9f450/MODULE_13537832482499839909+e30acd3a/model.neff [2025-07-26 00:34:26.003: I neuronx_distributed/paral lel_layers/parallel_state.py:379] [rank_0_pp-1_tp-1_dp-1_cp-1] Chosen Logic for replica groups ret_logic=, 'Ascending Ring PG Group')> [2025-07-26 00:34:26.004: I neuronx_distributed/paral -lel_layers/parallel_state.py:668] [rank_0_pp-1_tp-1_dp-1_cp-1] tp_groups: replica_groups.tp_groups=[[0]] [2025-07-26 00:34:26.004: I neuronx_distributed/paral -lel_layers/parallel_state.py:669] [rank_0_pp-1_tp-1_dp-1_cp-1] dp_groups: replica_groups.dp_groups=[[0]] [2025-07-26 00:34:26.004: I neuronx_distributed/paral -lel_layers/parallel_state.py:670] [rank_0_pp-1_tp-1_dp-1_cp-1] pp_groups: replica_groups.pp_groups=[[0]] [2025-07-26 00:34:26.004: I neuronx_distributed/paral -lel_layers/parallel_state.py:671] [rank_0_pp-1_tp-1_dp-1_cp-1] cp_groups: replica_groups.cp_groups=[[0]] [2025-07-26 00:34:26.004: I neuronx_distributed/paral -lel_layers/parallel_state.py:672] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_model_groups: replica_groups.ep_model_groups=[[0]] [2025-07-26 00:34:26.004: I neuronx_distributed/paral -lel_layers/parallel_state.py:673] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_data_groups: replica_groups.ep_data_groups=[[0]] WARNING:__main__:Unexpected keys in state dict: ['emb -ed_tokens.weight', 'norm.weight', 'lm_head.weight', 'layers.0.self_attn.q_proj.weight', 'layers.0.self_attn.k_proj.weight', 'layers.0.self_attn.v_proj.weight', 'layers.0.self_attn.o_proj.weight', 'layers.0.mlp.gate_proj.weight', 'layers.0.mlp.down_proj.weight', 'layers.0.mlp.up_proj.weight', 'layers.0.input_layernorm.weight', 'layers.0.post_attention_layernorm.weight', 'layers.1.self_attn.q_proj.weight', 'layers.1.self_attn.k_proj.weight', 'layers.1.self_attn.v_proj.weight', 'layers.1.self_attn.o_proj.weight', 'layers.1.mlp.gate_proj.weight', 'layers.1.mlp.down_proj.weight', 'layers.1.mlp.up_proj.weight', 'layers.1.input_layernorm.weight', 'layers.1.post_attention_layernorm.weight', 'layers.2.self_attn.q_proj.weight', 'layers.2.self_attn.k_proj.weight', 'layers.2.self_attn.v_proj.weight', 'layers.2.self_attn.o_proj.weight', 'layers.2.mlp.gate_proj.weight', 'layers.2.mlp.down_proj.weight', 'layers.2.mlp.up_proj.weight', 'layers.2.input_layernorm.weight', 'layers.2.post_attention_layernorm.weight', 'layers.3.self_attn.q_proj.weight', 'layers.3.self_attn.k_proj.weight', 'layers.3.self_attn.v_proj.weight', 'layers.3.self_attn.o_proj.weight', 'layers.3.mlp.gate_proj.weight', 'layers.3.mlp.down_proj.weight', 'layers.3.mlp.up_proj.weight', 'layers.3.input_layernorm.weight', 'layers.3.post_attention_layernorm.weight', 'layers.4.self_attn.q_proj.weight', 'layers.4.self_attn.k_proj.weight', 'layers.4.self_attn.v_proj.weight', 'layers.4.self_attn.o_proj.weight', 'layers.4.mlp.gate_proj.weight', 'layers.4.mlp.down_proj.weight', 'layers.4.mlp.up_proj.weight', 'layers.4.input_layernorm.weight', 'layers.4.post_attention_layernorm.weight', 'layers.5.self_attn.q_proj.weight', 'layers.5.self_attn.k_proj.weight', 'layers.5.self_attn.v_proj.weight', 'layers.5.self_attn.o_proj.weight', 'layers.5.mlp.gate_proj.weight', 'layers.5.mlp.down_proj.weight', 'layers.5.mlp.up_proj.weight', 'layers.5.input_layernorm.weight', 'layers.5.post_attention_layernorm.weight', 'layers.6.self_attn.q_proj.weight', 'layers.6.self_attn.k_proj.weight', 'layers.6.self_attn.v_proj.weight', 'layers.6.self_attn.o_proj.weight', 'layers.6.mlp.gate_proj.weight', 'layers.6.mlp.down_proj.weight', 'layers.6.mlp.up_proj.weight', 'layers.6.input_layernorm.weight', 'layers.6.post_attention_layernorm.weight', 'layers.7.self_attn.q_proj.weight', 'layers.7.self_attn.k_proj.weight', 'layers.7.self_attn.v_proj.weight', 'layers.7.self_attn.o_proj.weight', 'layers.7.mlp.gate_proj.weight', 'layers.7.mlp.down_proj.weight', 'layers.7.mlp.up_proj.weight', 'layers.7.input_layernorm.weight', 'layers.7.post_attention_layernorm.weight', 'layers.8.self_attn.q_proj.weight', 'layers.8.self_attn.k_proj.weight', 'layers.8.self_attn.v_proj.weight', 'layers.8.self_attn.o_proj.weight', 'layers.8.mlp.gate_proj.weight', 'layers.8.mlp.down_proj.weight', 'layers.8.mlp.up_proj.weight', 'layers.8.input_layernorm.weight', 'layers.8.post_attention_layernorm.weight', 'layers.9.self_attn.q_proj.weight', 'layers.9.self_attn.k_proj.weight', 'layers.9.self_attn.v_proj.weight', 'layers.9.self_attn.o_proj.weight', 'layers.9.mlp.gate_proj.weight', 'layers.9.mlp.down_proj.weight', 'layers.9.mlp.up_proj.weight', 'layers.9.input_layernorm.weight', 'layers.9.post_attention_layernorm.weight', 'layers.10.self_attn.q_proj.weight', 'layers.10.self_attn.k_proj.weight', 'layers.10.self_attn.v_proj.weight', 'layers.10.self_attn.o_proj.weight', 'layers.10.mlp.gate_proj.weight', 'layers.10.mlp.down_proj.weight', 'layers.10.mlp.up_proj.weight', 'layers.10.input_layernorm.weight', 'layers.10.post_attention_layernorm.weight', 'layers.11.self_attn.q_proj.weight', 'layers.11.self_attn.k_proj.weight', 'layers.11.self_attn.v_proj.weight', 'layers.11.self_attn.o_proj.weight', 'layers.11.mlp.gate_proj.weight', 'layers.11.mlp.down_proj.weight', 'layers.11.mlp.up_proj.weight', 'layers.11.input_layernorm.weight', 'layers.11.post_attention_layernorm.weight', 'layers.12.self_attn.q_proj.weight', 'layers.12.self_attn.k_proj.weight', 'layers.12.self_attn.v_proj.weight', 'layers.12.self_attn.o_proj.weight', 'layers.12.mlp.gate_proj.weight', 'layers.12.mlp.down_proj.weight', 'layers.12.mlp.up_proj.weight', 'layers.12.input_layernorm.weight', 'layers.12.post_attention_layernorm.weight', 'layers.13.self_attn.q_proj.weight', 'layers.13.self_attn.k_proj.weight', 'layers.13.self_attn.v_proj.weight', 'layers.13.self_attn.o_proj.weight', 'layers.13.mlp.gate_proj.weight', 'layers.13.mlp.down_proj.weight', 'layers.13.mlp.up_proj.weight', 'layers.13.input_layernorm.weight', 'layers.13.post_attention_layernorm.weight', 'layers.14.self_attn.q_proj.weight', 'layers.14.self_attn.k_proj.weight', 'layers.14.self_attn.v_proj.weight', 'layers.14.self_attn.o_proj.weight', 'layers.14.mlp.gate_proj.weight', 'layers.14.mlp.down_proj.weight', 'layers.14.mlp.up_proj.weight', 'layers.14.input_layernorm.weight', 'layers.14.post_attention_layernorm.weight', 'layers.15.self_attn.q_proj.weight', 'layers.15.self_attn.k_proj.weight', 'layers.15.self_attn.v_proj.weight', 'layers.15.self_attn.o_proj.weight', 'layers.15.mlp.gate_proj.weight', 'layers.15.mlp.down_proj.weight', 'layers.15.mlp.up_proj.weight', 'layers.15.input_layernorm.weight', 'layers.15.post_attention_layernorm.weight', 'layers.0.self_attn.rank_util.rank', 'layers.1.self_attn.rank_util.rank', 'layers.2.self_attn.rank_util.rank', 'layers.3.self_attn.rank_util.rank', 'layers.4.self_attn.rank_util.rank', 'layers.5.self_attn.rank_util.rank', 'layers.6.self_attn.rank_util.rank', 'layers.7.self_attn.rank_util.rank', 'layers.8.self_attn.rank_util.rank', 'layers.9.self_attn.rank_util.rank', 'layers.10.self_attn.rank_util.rank', 'layers.11.self_attn.rank_util.rank', 'layers.12.self_attn.rank_util.rank', 'layers.13.self_attn.rank_util.rank', 'layers.14.self_attn.rank_util.rank', 'layers.15.self_attn.rank_util.rank', 'rank_util.rank'] INFO:Neuron:Saving the neuron_config to llama3_compil -ed/ INFO:Neuron:Generating HLOs for the following models: - ['context_encoding_model', 'token_generation_model'][2025-07-26 00:34:27.425: I neuronx_distributed/paral -lel_layers/parallel_state.py:628] > initializing tensor model parallel with size 1 [2025-07-26 00:34:27.425: I neuronx_distributed/paral -lel_layers/parallel_state.py:629] > initializing pipeline model parallel with size 1 [2025-07-26 00:34:27.425: I neuronx_distributed/paral -lel_layers/parallel_state.py:630] > initializing context model parallel with size 1 [2025-07-26 00:34:27.425: I neuronx_distributed/paral -lel_layers/parallel_state.py:631] > initializing data parallel with size 1 [2025-07-26 00:34:27.425: I neuronx_distributed/paral -lel_layers/parallel_state.py:632] > initializing world size to 1 [2025-07-26 00:34:27.426: I neuronx_distributed/paral -lel_layers/parallel_state.py:379] [rank_0_pp0_tp0_dp0_cp0] Chosen Logic for replica groups ret_logic=, 'Ascending Ring PG Group')> [2025-07-26 00:34:27.426: I neuronx_distributed/paral -lel_layers/parallel_state.py:668] [rank_0_pp0_tp0_dp0_cp0] tp_groups: replica_groups.tp_groups=[[0]] [2025-07-26 00:34:27.427: I neuronx_distributed/paral -lel_layers/parallel_state.py:669] [rank_0_pp0_tp0_dp0_cp0] dp_groups: replica_groups.dp_groups=[[0]] [2025-07-26 00:34:27.427: I neuronx_distributed/paral -lel_layers/parallel_state.py:670] [rank_0_pp0_tp0_dp0_cp0] pp_groups: replica_groups.pp_groups=[[0]] [2025-07-26 00:34:27.427: I neuronx_distributed/paral -lel_layers/parallel_state.py:671] [rank_0_pp0_tp0_dp0_cp0] cp_groups: replica_groups.cp_groups=[[0]] [2025-07-26 00:34:27.427: I neuronx_distributed/paral -lel_layers/parallel_state.py:672] [rank_0_pp0_tp0_dp0_cp0] ep_model_groups: replica_groups.ep_model_groups=[[0]] [2025-07-26 00:34:27.427: I neuronx_distributed/paral -lel_layers/parallel_state.py:673] [rank_0_pp0_tp0_dp0_cp0] ep_data_groups: replica_groups.ep_data_groups=[[0]] INFO:Neuron:Generating 1 hlos for key: context_encodi -ng_model INFO:Neuron:Started loading module context_encoding_m -odel WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! ERROR:__main__:Model compilation failed: 'NeuronLlama -3Model' object has no attribute 'on_device_sampling' ERROR:__main__:Compilation failed: 'NeuronLlama3Model -' object has no attribute 'on_device_sampling' +lel_layers/parallel_state.py:668] [rank_0_pp-1_tp-1_dp-1_cp-1] tp_groups: replica_groups.tp_groups=[[0]] [2025-07-26 00:34:26.004: I neuronx_distributed/paral +lel_layers/parallel_state.py:669] [rank_0_pp-1_tp-1_dp-1_cp-1] dp_groups: replica_groups.dp_groups=[[0]] [2025-07-26 00:34:26.004: I neuronx_distributed/paral +lel_layers/parallel_state.py:670] [rank_0_pp-1_tp-1_dp-1_cp-1] pp_groups: replica_groups.pp_groups=[[0]] [2025-07-26 00:34:26.004: I neuronx_distributed/paral +lel_layers/parallel_state.py:671] [rank_0_pp-1_tp-1_dp-1_cp-1] cp_groups: replica_groups.cp_groups=[[0]] [2025-07-26 00:34:26.004: I neuronx_distributed/paral +lel_layers/parallel_state.py:672] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_model_groups: replica_groups.ep_model_groups=[[0]] [2025-07-26 00:34:26.004: I neuronx_distributed/paral +lel_layers/parallel_state.py:673] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_data_groups: replica_groups.ep_data_groups=[[0]] WARNING:**main**:Unexpected keys in state dict: ['emb +ed_tokens.weight', 'norm.weight', 'lm_head.weight', 'layers.0.self_attn.q_proj.weight', 'layers.0.self_attn.k_proj.weight', 'layers.0.self_attn.v_proj.weight', 'layers.0.self_attn.o_proj.weight', 'layers.0.mlp.gate_proj.weight', 'layers.0.mlp.down_proj.weight', 'layers.0.mlp.up_proj.weight', 'layers.0.input_layernorm.weight', 'layers.0.post_attention_layernorm.weight', 'layers.1.self_attn.q_proj.weight', 'layers.1.self_attn.k_proj.weight', 'layers.1.self_attn.v_proj.weight', 'layers.1.self_attn.o_proj.weight', 'layers.1.mlp.gate_proj.weight', 'layers.1.mlp.down_proj.weight', 'layers.1.mlp.up_proj.weight', 'layers.1.input_layernorm.weight', 'layers.1.post_attention_layernorm.weight', 'layers.2.self_attn.q_proj.weight', 'layers.2.self_attn.k_proj.weight', 'layers.2.self_attn.v_proj.weight', 'layers.2.self_attn.o_proj.weight', 'layers.2.mlp.gate_proj.weight', 'layers.2.mlp.down_proj.weight', 'layers.2.mlp.up_proj.weight', 'layers.2.input_layernorm.weight', 'layers.2.post_attention_layernorm.weight', 'layers.3.self_attn.q_proj.weight', 'layers.3.self_attn.k_proj.weight', 'layers.3.self_attn.v_proj.weight', 'layers.3.self_attn.o_proj.weight', 'layers.3.mlp.gate_proj.weight', 'layers.3.mlp.down_proj.weight', 'layers.3.mlp.up_proj.weight', 'layers.3.input_layernorm.weight', 'layers.3.post_attention_layernorm.weight', 'layers.4.self_attn.q_proj.weight', 'layers.4.self_attn.k_proj.weight', 'layers.4.self_attn.v_proj.weight', 'layers.4.self_attn.o_proj.weight', 'layers.4.mlp.gate_proj.weight', 'layers.4.mlp.down_proj.weight', 'layers.4.mlp.up_proj.weight', 'layers.4.input_layernorm.weight', 'layers.4.post_attention_layernorm.weight', 'layers.5.self_attn.q_proj.weight', 'layers.5.self_attn.k_proj.weight', 'layers.5.self_attn.v_proj.weight', 'layers.5.self_attn.o_proj.weight', 'layers.5.mlp.gate_proj.weight', 'layers.5.mlp.down_proj.weight', 'layers.5.mlp.up_proj.weight', 'layers.5.input_layernorm.weight', 'layers.5.post_attention_layernorm.weight', 'layers.6.self_attn.q_proj.weight', 'layers.6.self_attn.k_proj.weight', 'layers.6.self_attn.v_proj.weight', 'layers.6.self_attn.o_proj.weight', 'layers.6.mlp.gate_proj.weight', 'layers.6.mlp.down_proj.weight', 'layers.6.mlp.up_proj.weight', 'layers.6.input_layernorm.weight', 'layers.6.post_attention_layernorm.weight', 'layers.7.self_attn.q_proj.weight', 'layers.7.self_attn.k_proj.weight', 'layers.7.self_attn.v_proj.weight', 'layers.7.self_attn.o_proj.weight', 'layers.7.mlp.gate_proj.weight', 'layers.7.mlp.down_proj.weight', 'layers.7.mlp.up_proj.weight', 'layers.7.input_layernorm.weight', 'layers.7.post_attention_layernorm.weight', 'layers.8.self_attn.q_proj.weight', 'layers.8.self_attn.k_proj.weight', 'layers.8.self_attn.v_proj.weight', 'layers.8.self_attn.o_proj.weight', 'layers.8.mlp.gate_proj.weight', 'layers.8.mlp.down_proj.weight', 'layers.8.mlp.up_proj.weight', 'layers.8.input_layernorm.weight', 'layers.8.post_attention_layernorm.weight', 'layers.9.self_attn.q_proj.weight', 'layers.9.self_attn.k_proj.weight', 'layers.9.self_attn.v_proj.weight', 'layers.9.self_attn.o_proj.weight', 'layers.9.mlp.gate_proj.weight', 'layers.9.mlp.down_proj.weight', 'layers.9.mlp.up_proj.weight', 'layers.9.input_layernorm.weight', 'layers.9.post_attention_layernorm.weight', 'layers.10.self_attn.q_proj.weight', 'layers.10.self_attn.k_proj.weight', 'layers.10.self_attn.v_proj.weight', 'layers.10.self_attn.o_proj.weight', 'layers.10.mlp.gate_proj.weight', 'layers.10.mlp.down_proj.weight', 'layers.10.mlp.up_proj.weight', 'layers.10.input_layernorm.weight', 'layers.10.post_attention_layernorm.weight', 'layers.11.self_attn.q_proj.weight', 'layers.11.self_attn.k_proj.weight', 'layers.11.self_attn.v_proj.weight', 'layers.11.self_attn.o_proj.weight', 'layers.11.mlp.gate_proj.weight', 'layers.11.mlp.down_proj.weight', 'layers.11.mlp.up_proj.weight', 'layers.11.input_layernorm.weight', 'layers.11.post_attention_layernorm.weight', 'layers.12.self_attn.q_proj.weight', 'layers.12.self_attn.k_proj.weight', 'layers.12.self_attn.v_proj.weight', 'layers.12.self_attn.o_proj.weight', 'layers.12.mlp.gate_proj.weight', 'layers.12.mlp.down_proj.weight', 'layers.12.mlp.up_proj.weight', 'layers.12.input_layernorm.weight', 'layers.12.post_attention_layernorm.weight', 'layers.13.self_attn.q_proj.weight', 'layers.13.self_attn.k_proj.weight', 'layers.13.self_attn.v_proj.weight', 'layers.13.self_attn.o_proj.weight', 'layers.13.mlp.gate_proj.weight', 'layers.13.mlp.down_proj.weight', 'layers.13.mlp.up_proj.weight', 'layers.13.input_layernorm.weight', 'layers.13.post_attention_layernorm.weight', 'layers.14.self_attn.q_proj.weight', 'layers.14.self_attn.k_proj.weight', 'layers.14.self_attn.v_proj.weight', 'layers.14.self_attn.o_proj.weight', 'layers.14.mlp.gate_proj.weight', 'layers.14.mlp.down_proj.weight', 'layers.14.mlp.up_proj.weight', 'layers.14.input_layernorm.weight', 'layers.14.post_attention_layernorm.weight', 'layers.15.self_attn.q_proj.weight', 'layers.15.self_attn.k_proj.weight', 'layers.15.self_attn.v_proj.weight', 'layers.15.self_attn.o_proj.weight', 'layers.15.mlp.gate_proj.weight', 'layers.15.mlp.down_proj.weight', 'layers.15.mlp.up_proj.weight', 'layers.15.input_layernorm.weight', 'layers.15.post_attention_layernorm.weight', 'layers.0.self_attn.rank_util.rank', 'layers.1.self_attn.rank_util.rank', 'layers.2.self_attn.rank_util.rank', 'layers.3.self_attn.rank_util.rank', 'layers.4.self_attn.rank_util.rank', 'layers.5.self_attn.rank_util.rank', 'layers.6.self_attn.rank_util.rank', 'layers.7.self_attn.rank_util.rank', 'layers.8.self_attn.rank_util.rank', 'layers.9.self_attn.rank_util.rank', 'layers.10.self_attn.rank_util.rank', 'layers.11.self_attn.rank_util.rank', 'layers.12.self_attn.rank_util.rank', 'layers.13.self_attn.rank_util.rank', 'layers.14.self_attn.rank_util.rank', 'layers.15.self_attn.rank_util.rank', 'rank_util.rank'] INFO:Neuron:Saving the neuron_config to llama3_compil +ed/ INFO:Neuron:Generating HLOs for the following models: +['context_encoding_model', 'token_generation_model'][2025-07-26 00:34:27.425: I neuronx_distributed/paral lel_layers/parallel_state.py:628] > initializing tensor model parallel with size 1 [2025-07-26 00:34:27.425: I neuronx_distributed/paral +lel_layers/parallel_state.py:629] > initializing pipeline model parallel with size 1 [2025-07-26 00:34:27.425: I neuronx_distributed/paral +lel_layers/parallel_state.py:630] > initializing context model parallel with size 1 [2025-07-26 00:34:27.425: I neuronx_distributed/paral +lel_layers/parallel_state.py:631] > initializing data parallel with size 1 [2025-07-26 00:34:27.425: I neuronx_distributed/paral +lel_layers/parallel_state.py:632] > initializing world size to 1 [2025-07-26 00:34:27.426: I neuronx_distributed/paral +lel_layers/parallel_state.py:379] [rank_0_pp0_tp0_dp0_cp0] Chosen Logic for replica groups ret_logic=, 'Ascending Ring PG Group')> [2025-07-26 00:34:27.426: I neuronx_distributed/paral +lel_layers/parallel_state.py:668] [rank_0_pp0_tp0_dp0_cp0] tp_groups: replica_groups.tp_groups=[[0]] [2025-07-26 00:34:27.427: I neuronx_distributed/paral +lel_layers/parallel_state.py:669] [rank_0_pp0_tp0_dp0_cp0] dp_groups: replica_groups.dp_groups=[[0]] [2025-07-26 00:34:27.427: I neuronx_distributed/paral +lel_layers/parallel_state.py:670] [rank_0_pp0_tp0_dp0_cp0] pp_groups: replica_groups.pp_groups=[[0]] [2025-07-26 00:34:27.427: I neuronx_distributed/paral +lel_layers/parallel_state.py:671] [rank_0_pp0_tp0_dp0_cp0] cp_groups: replica_groups.cp_groups=[[0]] [2025-07-26 00:34:27.427: I neuronx_distributed/paral +lel_layers/parallel_state.py:672] [rank_0_pp0_tp0_dp0_cp0] ep_model_groups: replica_groups.ep_model_groups=[[0]] [2025-07-26 00:34:27.427: I neuronx_distributed/paral +lel_layers/parallel_state.py:673] [rank_0_pp0_tp0_dp0_cp0] ep_data_groups: replica_groups.ep_data_groups=[[0]] INFO:Neuron:Generating 1 hlos for key: context_encodi +ng_model INFO:Neuron:Started loading module context_encoding_m +odel WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! ERROR:**main**:Model compilation failed: 'NeuronLlama +3Model' object has no attribute 'on_device_sampling' ERROR:**main**:Compilation failed: 'NeuronLlama3Model +' object has no attribute 'on_device_sampling' Great progress! The model is loading and the GQA is being handled correctly (converting to MHA for TP=1). I need to add the on_device_sampling attribute. Let me add this: Accepted edits to modeling_llama3.py - Let's try compiling again: Command -python compile_model.py --checkpoint_path ./llama3_neuron_checkpoint --output_path ./llama3_compiled --batch_size 1 --seq_len 128 +python compile*model.py --checkpoint_path ./llama3_neuron_checkpoint --output_path ./llama3_compiled --batch_size 1 --seq_len 128 /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributedInference -/src/neuronx_distributed_inference/modules/attention/utils.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.custom_c -alls import neuron_cumsum /home/ec2-user/NeuronxSDK/NeuronxDistributedInference -/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference -/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference -/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_ -llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_ -llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_ -llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.attention_base import NeuronAttentionBase [2025-07-26 00:35:17.362: I neuronx_distributed/paral -lel_layers/parallel_state.py:628] > initializing tensor model parallel with size 1 [2025-07-26 00:35:17.362: I neuronx_distributed/paral -lel_layers/parallel_state.py:629] > initializing pipeline model parallel with size 1 [2025-07-26 00:35:17.362: I neuronx_distributed/paral -lel_layers/parallel_state.py:630] > initializing context model parallel with size 1 [2025-07-26 00:35:17.362: I neuronx_distributed/paral -lel_layers/parallel_state.py:631] > initializing data parallel with size 1 [2025-07-26 00:35:17.362: I neuronx_distributed/paral -lel_layers/parallel_state.py:632] > initializing world size to 1 2025-07-26 00:35:17.000367: 347802 INFO ||NEURON_CC -_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/neuronxcc-2.19.8089.0+8ab9f450/MODULE_13537832482499839909+e30acd3a/model.neff [2025-07-26 00:35:17.405: I neuronx_distributed/paral +/src/neuronx_distributed_inference/modules/attention/utils.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.custom_c +alls import neuron_cumsum /home/ec2-user/NeuronxSDK/NeuronxDistributedInference +/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference +/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference +/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx* +llama3/modeling*llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx* +llama3/modeling*llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx* +llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.attention_base import NeuronAttentionBase [2025-07-26 00:35:17.362: I neuronx_distributed/paral +lel_layers/parallel_state.py:628] > initializing tensor model parallel with size 1 [2025-07-26 00:35:17.362: I neuronx_distributed/paral +lel_layers/parallel_state.py:629] > initializing pipeline model parallel with size 1 [2025-07-26 00:35:17.362: I neuronx_distributed/paral +lel_layers/parallel_state.py:630] > initializing context model parallel with size 1 [2025-07-26 00:35:17.362: I neuronx_distributed/paral +lel_layers/parallel_state.py:631] > initializing data parallel with size 1 [2025-07-26 00:35:17.362: I neuronx_distributed/paral +lel_layers/parallel_state.py:632] > initializing world size to 1 2025-07-26 00:35:17.000367: 347802 INFO ||NEURON_CC +\_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/neuronxcc-2.19.8089.0+8ab9f450/MODULE_13537832482499839909+e30acd3a/model.neff [2025-07-26 00:35:17.405: I neuronx_distributed/paral lel_layers/parallel_state.py:379] [rank_0_pp-1_tp-1_dp-1_cp-1] Chosen Logic for replica groups ret_logic=, 'Ascending Ring PG Group')> [2025-07-26 00:35:17.406: I neuronx_distributed/paral -lel_layers/parallel_state.py:668] [rank_0_pp-1_tp-1_dp-1_cp-1] tp_groups: replica_groups.tp_groups=[[0]] [2025-07-26 00:35:17.406: I neuronx_distributed/paral -lel_layers/parallel_state.py:669] [rank_0_pp-1_tp-1_dp-1_cp-1] dp_groups: replica_groups.dp_groups=[[0]] [2025-07-26 00:35:17.406: I neuronx_distributed/paral -lel_layers/parallel_state.py:670] [rank_0_pp-1_tp-1_dp-1_cp-1] pp_groups: replica_groups.pp_groups=[[0]] [2025-07-26 00:35:17.406: I neuronx_distributed/paral -lel_layers/parallel_state.py:671] [rank_0_pp-1_tp-1_dp-1_cp-1] cp_groups: replica_groups.cp_groups=[[0]] [2025-07-26 00:35:17.406: I neuronx_distributed/paral -lel_layers/parallel_state.py:672] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_model_groups: replica_groups.ep_model_groups=[[0]] [2025-07-26 00:35:17.406: I neuronx_distributed/paral -lel_layers/parallel_state.py:673] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_data_groups: replica_groups.ep_data_groups=[[0]] WARNING:__main__:Unexpected keys in state dict: ['emb -ed_tokens.weight', 'norm.weight', 'lm_head.weight', 'layers.0.self_attn.q_proj.weight', 'layers.0.self_attn.k_proj.weight', 'layers.0.self_attn.v_proj.weight', 'layers.0.self_attn.o_proj.weight', 'layers.0.mlp.gate_proj.weight', 'layers.0.mlp.down_proj.weight', 'layers.0.mlp.up_proj.weight', 'layers.0.input_layernorm.weight', 'layers.0.post_attention_layernorm.weight', 'layers.1.self_attn.q_proj.weight', 'layers.1.self_attn.k_proj.weight', 'layers.1.self_attn.v_proj.weight', 'layers.1.self_attn.o_proj.weight', 'layers.1.mlp.gate_proj.weight', 'layers.1.mlp.down_proj.weight', 'layers.1.mlp.up_proj.weight', 'layers.1.input_layernorm.weight', 'layers.1.post_attention_layernorm.weight', 'layers.2.self_attn.q_proj.weight', 'layers.2.self_attn.k_proj.weight', 'layers.2.self_attn.v_proj.weight', 'layers.2.self_attn.o_proj.weight', 'layers.2.mlp.gate_proj.weight', 'layers.2.mlp.down_proj.weight', 'layers.2.mlp.up_proj.weight', 'layers.2.input_layernorm.weight', 'layers.2.post_attention_layernorm.weight', 'layers.3.self_attn.q_proj.weight', 'layers.3.self_attn.k_proj.weight', 'layers.3.self_attn.v_proj.weight', 'layers.3.self_attn.o_proj.weight', 'layers.3.mlp.gate_proj.weight', 'layers.3.mlp.down_proj.weight', 'layers.3.mlp.up_proj.weight', 'layers.3.input_layernorm.weight', 'layers.3.post_attention_layernorm.weight', 'layers.4.self_attn.q_proj.weight', 'layers.4.self_attn.k_proj.weight', 'layers.4.self_attn.v_proj.weight', 'layers.4.self_attn.o_proj.weight', 'layers.4.mlp.gate_proj.weight', 'layers.4.mlp.down_proj.weight', 'layers.4.mlp.up_proj.weight', 'layers.4.input_layernorm.weight', 'layers.4.post_attention_layernorm.weight', 'layers.5.self_attn.q_proj.weight', 'layers.5.self_attn.k_proj.weight', 'layers.5.self_attn.v_proj.weight', 'layers.5.self_attn.o_proj.weight', 'layers.5.mlp.gate_proj.weight', 'layers.5.mlp.down_proj.weight', 'layers.5.mlp.up_proj.weight', 'layers.5.input_layernorm.weight', 'layers.5.post_attention_layernorm.weight', 'layers.6.self_attn.q_proj.weight', 'layers.6.self_attn.k_proj.weight', 'layers.6.self_attn.v_proj.weight', 'layers.6.self_attn.o_proj.weight', 'layers.6.mlp.gate_proj.weight', 'layers.6.mlp.down_proj.weight', 'layers.6.mlp.up_proj.weight', 'layers.6.input_layernorm.weight', 'layers.6.post_attention_layernorm.weight', 'layers.7.self_attn.q_proj.weight', 'layers.7.self_attn.k_proj.weight', 'layers.7.self_attn.v_proj.weight', 'layers.7.self_attn.o_proj.weight', 'layers.7.mlp.gate_proj.weight', 'layers.7.mlp.down_proj.weight', 'layers.7.mlp.up_proj.weight', 'layers.7.input_layernorm.weight', 'layers.7.post_attention_layernorm.weight', 'layers.8.self_attn.q_proj.weight', 'layers.8.self_attn.k_proj.weight', 'layers.8.self_attn.v_proj.weight', 'layers.8.self_attn.o_proj.weight', 'layers.8.mlp.gate_proj.weight', 'layers.8.mlp.down_proj.weight', 'layers.8.mlp.up_proj.weight', 'layers.8.input_layernorm.weight', 'layers.8.post_attention_layernorm.weight', 'layers.9.self_attn.q_proj.weight', 'layers.9.self_attn.k_proj.weight', 'layers.9.self_attn.v_proj.weight', 'layers.9.self_attn.o_proj.weight', 'layers.9.mlp.gate_proj.weight', 'layers.9.mlp.down_proj.weight', 'layers.9.mlp.up_proj.weight', 'layers.9.input_layernorm.weight', 'layers.9.post_attention_layernorm.weight', 'layers.10.self_attn.q_proj.weight', 'layers.10.self_attn.k_proj.weight', 'layers.10.self_attn.v_proj.weight', 'layers.10.self_attn.o_proj.weight', 'layers.10.mlp.gate_proj.weight', 'layers.10.mlp.down_proj.weight', 'layers.10.mlp.up_proj.weight', 'layers.10.input_layernorm.weight', 'layers.10.post_attention_layernorm.weight', 'layers.11.self_attn.q_proj.weight', 'layers.11.self_attn.k_proj.weight', 'layers.11.self_attn.v_proj.weight', 'layers.11.self_attn.o_proj.weight', 'layers.11.mlp.gate_proj.weight', 'layers.11.mlp.down_proj.weight', 'layers.11.mlp.up_proj.weight', 'layers.11.input_layernorm.weight', 'layers.11.post_attention_layernorm.weight', 'layers.12.self_attn.q_proj.weight', 'layers.12.self_attn.k_proj.weight', 'layers.12.self_attn.v_proj.weight', 'layers.12.self_attn.o_proj.weight', 'layers.12.mlp.gate_proj.weight', 'layers.12.mlp.down_proj.weight', 'layers.12.mlp.up_proj.weight', 'layers.12.input_layernorm.weight', 'layers.12.post_attention_layernorm.weight', 'layers.13.self_attn.q_proj.weight', 'layers.13.self_attn.k_proj.weight', 'layers.13.self_attn.v_proj.weight', 'layers.13.self_attn.o_proj.weight', 'layers.13.mlp.gate_proj.weight', 'layers.13.mlp.down_proj.weight', 'layers.13.mlp.up_proj.weight', 'layers.13.input_layernorm.weight', 'layers.13.post_attention_layernorm.weight', 'layers.14.self_attn.q_proj.weight', 'layers.14.self_attn.k_proj.weight', 'layers.14.self_attn.v_proj.weight', 'layers.14.self_attn.o_proj.weight', 'layers.14.mlp.gate_proj.weight', 'layers.14.mlp.down_proj.weight', 'layers.14.mlp.up_proj.weight', 'layers.14.input_layernorm.weight', 'layers.14.post_attention_layernorm.weight', 'layers.15.self_attn.q_proj.weight', 'layers.15.self_attn.k_proj.weight', 'layers.15.self_attn.v_proj.weight', 'layers.15.self_attn.o_proj.weight', 'layers.15.mlp.gate_proj.weight', 'layers.15.mlp.down_proj.weight', 'layers.15.mlp.up_proj.weight', 'layers.15.input_layernorm.weight', 'layers.15.post_attention_layernorm.weight', 'layers.0.self_attn.rank_util.rank', 'layers.1.self_attn.rank_util.rank', 'layers.2.self_attn.rank_util.rank', 'layers.3.self_attn.rank_util.rank', 'layers.4.self_attn.rank_util.rank', 'layers.5.self_attn.rank_util.rank', 'layers.6.self_attn.rank_util.rank', 'layers.7.self_attn.rank_util.rank', 'layers.8.self_attn.rank_util.rank', 'layers.9.self_attn.rank_util.rank', 'layers.10.self_attn.rank_util.rank', 'layers.11.self_attn.rank_util.rank', 'layers.12.self_attn.rank_util.rank', 'layers.13.self_attn.rank_util.rank', 'layers.14.self_attn.rank_util.rank', 'layers.15.self_attn.rank_util.rank', 'rank_util.rank'] INFO:Neuron:Saving the neuron_config to llama3_compil -ed/ INFO:Neuron:Generating HLOs for the following models: - ['context_encoding_model', 'token_generation_model'][2025-07-26 00:35:18.826: I neuronx_distributed/paral -lel_layers/parallel_state.py:628] > initializing tensor model parallel with size 1 [2025-07-26 00:35:18.826: I neuronx_distributed/paral -lel_layers/parallel_state.py:629] > initializing pipeline model parallel with size 1 [2025-07-26 00:35:18.826: I neuronx_distributed/paral -lel_layers/parallel_state.py:630] > initializing context model parallel with size 1 [2025-07-26 00:35:18.826: I neuronx_distributed/paral -lel_layers/parallel_state.py:631] > initializing data parallel with size 1 [2025-07-26 00:35:18.826: I neuronx_distributed/paral -lel_layers/parallel_state.py:632] > initializing world size to 1 [2025-07-26 00:35:18.827: I neuronx_distributed/paral -lel_layers/parallel_state.py:379] [rank_0_pp0_tp0_dp0_cp0] Chosen Logic for replica groups ret_logic=, 'Ascending Ring PG Group')> [2025-07-26 00:35:18.827: I neuronx_distributed/paral -lel_layers/parallel_state.py:668] [rank_0_pp0_tp0_dp0_cp0] tp_groups: replica_groups.tp_groups=[[0]] [2025-07-26 00:35:18.827: I neuronx_distributed/paral -lel_layers/parallel_state.py:669] [rank_0_pp0_tp0_dp0_cp0] dp_groups: replica_groups.dp_groups=[[0]] [2025-07-26 00:35:18.827: I neuronx_distributed/paral -lel_layers/parallel_state.py:670] [rank_0_pp0_tp0_dp0_cp0] pp_groups: replica_groups.pp_groups=[[0]] [2025-07-26 00:35:18.827: I neuronx_distributed/paral -lel_layers/parallel_state.py:671] [rank_0_pp0_tp0_dp0_cp0] cp_groups: replica_groups.cp_groups=[[0]] [2025-07-26 00:35:18.827: I neuronx_distributed/paral -lel_layers/parallel_state.py:672] [rank_0_pp0_tp0_dp0_cp0] ep_model_groups: replica_groups.ep_model_groups=[[0]] [2025-07-26 00:35:18.827: I neuronx_distributed/paral -lel_layers/parallel_state.py:673] [rank_0_pp0_tp0_dp0_cp0] ep_data_groups: replica_groups.ep_data_groups=[[0]] INFO:Neuron:Generating 1 hlos for key: context_encodi -ng_model INFO:Neuron:Started loading module context_encoding_m -odel WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! ERROR:__main__:Model compilation failed: 'NeuronLlama -3Model' object has no attribute 'post_init' ERROR:__main__:Compilation failed: 'NeuronLlama3Model -' object has no attribute 'post_init' +lel_layers/parallel_state.py:668] [rank_0_pp-1_tp-1_dp-1_cp-1] tp_groups: replica_groups.tp_groups=[[0]] [2025-07-26 00:35:17.406: I neuronx_distributed/paral +lel_layers/parallel_state.py:669] [rank_0_pp-1_tp-1_dp-1_cp-1] dp_groups: replica_groups.dp_groups=[[0]] [2025-07-26 00:35:17.406: I neuronx_distributed/paral +lel_layers/parallel_state.py:670] [rank_0_pp-1_tp-1_dp-1_cp-1] pp_groups: replica_groups.pp_groups=[[0]] [2025-07-26 00:35:17.406: I neuronx_distributed/paral +lel_layers/parallel_state.py:671] [rank_0_pp-1_tp-1_dp-1_cp-1] cp_groups: replica_groups.cp_groups=[[0]] [2025-07-26 00:35:17.406: I neuronx_distributed/paral +lel_layers/parallel_state.py:672] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_model_groups: replica_groups.ep_model_groups=[[0]] [2025-07-26 00:35:17.406: I neuronx_distributed/paral +lel_layers/parallel_state.py:673] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_data_groups: replica_groups.ep_data_groups=[[0]] WARNING:**main**:Unexpected keys in state dict: ['emb +ed_tokens.weight', 'norm.weight', 'lm_head.weight', 'layers.0.self_attn.q_proj.weight', 'layers.0.self_attn.k_proj.weight', 'layers.0.self_attn.v_proj.weight', 'layers.0.self_attn.o_proj.weight', 'layers.0.mlp.gate_proj.weight', 'layers.0.mlp.down_proj.weight', 'layers.0.mlp.up_proj.weight', 'layers.0.input_layernorm.weight', 'layers.0.post_attention_layernorm.weight', 'layers.1.self_attn.q_proj.weight', 'layers.1.self_attn.k_proj.weight', 'layers.1.self_attn.v_proj.weight', 'layers.1.self_attn.o_proj.weight', 'layers.1.mlp.gate_proj.weight', 'layers.1.mlp.down_proj.weight', 'layers.1.mlp.up_proj.weight', 'layers.1.input_layernorm.weight', 'layers.1.post_attention_layernorm.weight', 'layers.2.self_attn.q_proj.weight', 'layers.2.self_attn.k_proj.weight', 'layers.2.self_attn.v_proj.weight', 'layers.2.self_attn.o_proj.weight', 'layers.2.mlp.gate_proj.weight', 'layers.2.mlp.down_proj.weight', 'layers.2.mlp.up_proj.weight', 'layers.2.input_layernorm.weight', 'layers.2.post_attention_layernorm.weight', 'layers.3.self_attn.q_proj.weight', 'layers.3.self_attn.k_proj.weight', 'layers.3.self_attn.v_proj.weight', 'layers.3.self_attn.o_proj.weight', 'layers.3.mlp.gate_proj.weight', 'layers.3.mlp.down_proj.weight', 'layers.3.mlp.up_proj.weight', 'layers.3.input_layernorm.weight', 'layers.3.post_attention_layernorm.weight', 'layers.4.self_attn.q_proj.weight', 'layers.4.self_attn.k_proj.weight', 'layers.4.self_attn.v_proj.weight', 'layers.4.self_attn.o_proj.weight', 'layers.4.mlp.gate_proj.weight', 'layers.4.mlp.down_proj.weight', 'layers.4.mlp.up_proj.weight', 'layers.4.input_layernorm.weight', 'layers.4.post_attention_layernorm.weight', 'layers.5.self_attn.q_proj.weight', 'layers.5.self_attn.k_proj.weight', 'layers.5.self_attn.v_proj.weight', 'layers.5.self_attn.o_proj.weight', 'layers.5.mlp.gate_proj.weight', 'layers.5.mlp.down_proj.weight', 'layers.5.mlp.up_proj.weight', 'layers.5.input_layernorm.weight', 'layers.5.post_attention_layernorm.weight', 'layers.6.self_attn.q_proj.weight', 'layers.6.self_attn.k_proj.weight', 'layers.6.self_attn.v_proj.weight', 'layers.6.self_attn.o_proj.weight', 'layers.6.mlp.gate_proj.weight', 'layers.6.mlp.down_proj.weight', 'layers.6.mlp.up_proj.weight', 'layers.6.input_layernorm.weight', 'layers.6.post_attention_layernorm.weight', 'layers.7.self_attn.q_proj.weight', 'layers.7.self_attn.k_proj.weight', 'layers.7.self_attn.v_proj.weight', 'layers.7.self_attn.o_proj.weight', 'layers.7.mlp.gate_proj.weight', 'layers.7.mlp.down_proj.weight', 'layers.7.mlp.up_proj.weight', 'layers.7.input_layernorm.weight', 'layers.7.post_attention_layernorm.weight', 'layers.8.self_attn.q_proj.weight', 'layers.8.self_attn.k_proj.weight', 'layers.8.self_attn.v_proj.weight', 'layers.8.self_attn.o_proj.weight', 'layers.8.mlp.gate_proj.weight', 'layers.8.mlp.down_proj.weight', 'layers.8.mlp.up_proj.weight', 'layers.8.input_layernorm.weight', 'layers.8.post_attention_layernorm.weight', 'layers.9.self_attn.q_proj.weight', 'layers.9.self_attn.k_proj.weight', 'layers.9.self_attn.v_proj.weight', 'layers.9.self_attn.o_proj.weight', 'layers.9.mlp.gate_proj.weight', 'layers.9.mlp.down_proj.weight', 'layers.9.mlp.up_proj.weight', 'layers.9.input_layernorm.weight', 'layers.9.post_attention_layernorm.weight', 'layers.10.self_attn.q_proj.weight', 'layers.10.self_attn.k_proj.weight', 'layers.10.self_attn.v_proj.weight', 'layers.10.self_attn.o_proj.weight', 'layers.10.mlp.gate_proj.weight', 'layers.10.mlp.down_proj.weight', 'layers.10.mlp.up_proj.weight', 'layers.10.input_layernorm.weight', 'layers.10.post_attention_layernorm.weight', 'layers.11.self_attn.q_proj.weight', 'layers.11.self_attn.k_proj.weight', 'layers.11.self_attn.v_proj.weight', 'layers.11.self_attn.o_proj.weight', 'layers.11.mlp.gate_proj.weight', 'layers.11.mlp.down_proj.weight', 'layers.11.mlp.up_proj.weight', 'layers.11.input_layernorm.weight', 'layers.11.post_attention_layernorm.weight', 'layers.12.self_attn.q_proj.weight', 'layers.12.self_attn.k_proj.weight', 'layers.12.self_attn.v_proj.weight', 'layers.12.self_attn.o_proj.weight', 'layers.12.mlp.gate_proj.weight', 'layers.12.mlp.down_proj.weight', 'layers.12.mlp.up_proj.weight', 'layers.12.input_layernorm.weight', 'layers.12.post_attention_layernorm.weight', 'layers.13.self_attn.q_proj.weight', 'layers.13.self_attn.k_proj.weight', 'layers.13.self_attn.v_proj.weight', 'layers.13.self_attn.o_proj.weight', 'layers.13.mlp.gate_proj.weight', 'layers.13.mlp.down_proj.weight', 'layers.13.mlp.up_proj.weight', 'layers.13.input_layernorm.weight', 'layers.13.post_attention_layernorm.weight', 'layers.14.self_attn.q_proj.weight', 'layers.14.self_attn.k_proj.weight', 'layers.14.self_attn.v_proj.weight', 'layers.14.self_attn.o_proj.weight', 'layers.14.mlp.gate_proj.weight', 'layers.14.mlp.down_proj.weight', 'layers.14.mlp.up_proj.weight', 'layers.14.input_layernorm.weight', 'layers.14.post_attention_layernorm.weight', 'layers.15.self_attn.q_proj.weight', 'layers.15.self_attn.k_proj.weight', 'layers.15.self_attn.v_proj.weight', 'layers.15.self_attn.o_proj.weight', 'layers.15.mlp.gate_proj.weight', 'layers.15.mlp.down_proj.weight', 'layers.15.mlp.up_proj.weight', 'layers.15.input_layernorm.weight', 'layers.15.post_attention_layernorm.weight', 'layers.0.self_attn.rank_util.rank', 'layers.1.self_attn.rank_util.rank', 'layers.2.self_attn.rank_util.rank', 'layers.3.self_attn.rank_util.rank', 'layers.4.self_attn.rank_util.rank', 'layers.5.self_attn.rank_util.rank', 'layers.6.self_attn.rank_util.rank', 'layers.7.self_attn.rank_util.rank', 'layers.8.self_attn.rank_util.rank', 'layers.9.self_attn.rank_util.rank', 'layers.10.self_attn.rank_util.rank', 'layers.11.self_attn.rank_util.rank', 'layers.12.self_attn.rank_util.rank', 'layers.13.self_attn.rank_util.rank', 'layers.14.self_attn.rank_util.rank', 'layers.15.self_attn.rank_util.rank', 'rank_util.rank'] INFO:Neuron:Saving the neuron_config to llama3_compil +ed/ INFO:Neuron:Generating HLOs for the following models: +['context_encoding_model', 'token_generation_model'][2025-07-26 00:35:18.826: I neuronx_distributed/paral lel_layers/parallel_state.py:628] > initializing tensor model parallel with size 1 [2025-07-26 00:35:18.826: I neuronx_distributed/paral +lel_layers/parallel_state.py:629] > initializing pipeline model parallel with size 1 [2025-07-26 00:35:18.826: I neuronx_distributed/paral +lel_layers/parallel_state.py:630] > initializing context model parallel with size 1 [2025-07-26 00:35:18.826: I neuronx_distributed/paral +lel_layers/parallel_state.py:631] > initializing data parallel with size 1 [2025-07-26 00:35:18.826: I neuronx_distributed/paral +lel_layers/parallel_state.py:632] > initializing world size to 1 [2025-07-26 00:35:18.827: I neuronx_distributed/paral +lel_layers/parallel_state.py:379] [rank_0_pp0_tp0_dp0_cp0] Chosen Logic for replica groups ret_logic=, 'Ascending Ring PG Group')> [2025-07-26 00:35:18.827: I neuronx_distributed/paral +lel_layers/parallel_state.py:668] [rank_0_pp0_tp0_dp0_cp0] tp_groups: replica_groups.tp_groups=[[0]] [2025-07-26 00:35:18.827: I neuronx_distributed/paral +lel_layers/parallel_state.py:669] [rank_0_pp0_tp0_dp0_cp0] dp_groups: replica_groups.dp_groups=[[0]] [2025-07-26 00:35:18.827: I neuronx_distributed/paral +lel_layers/parallel_state.py:670] [rank_0_pp0_tp0_dp0_cp0] pp_groups: replica_groups.pp_groups=[[0]] [2025-07-26 00:35:18.827: I neuronx_distributed/paral +lel_layers/parallel_state.py:671] [rank_0_pp0_tp0_dp0_cp0] cp_groups: replica_groups.cp_groups=[[0]] [2025-07-26 00:35:18.827: I neuronx_distributed/paral +lel_layers/parallel_state.py:672] [rank_0_pp0_tp0_dp0_cp0] ep_model_groups: replica_groups.ep_model_groups=[[0]] [2025-07-26 00:35:18.827: I neuronx_distributed/paral +lel_layers/parallel_state.py:673] [rank_0_pp0_tp0_dp0_cp0] ep_data_groups: replica_groups.ep_data_groups=[[0]] INFO:Neuron:Generating 1 hlos for key: context_encodi +ng_model INFO:Neuron:Started loading module context_encoding_m +odel WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! ERROR:**main**:Model compilation failed: 'NeuronLlama +3Model' object has no attribute 'post_init' ERROR:**main**:Compilation failed: 'NeuronLlama3Model +' object has no attribute 'post_init' I need to remove the post_init() call since it's not needed for the compilation model. Let me fix this: Error(s) while editing @@ -1705,7 +1692,6 @@ Accepted edits to modeling_llama3.py - Let me also remove the post_init call from the init_model method: Read file(s) @@ -1715,143 +1701,141 @@ Accepted edits to modeling_llama3.py - Now let's try compiling again: Command -python compile_model.py --checkpoint_path ./llama3_neuron_checkpoint --output_path ./llama3_compiled --batch_size 1 --seq_len 128 +python compile*model.py --checkpoint_path ./llama3_neuron_checkpoint --output_path ./llama3_compiled --batch_size 1 --seq_len 128 /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +onx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neur -onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +onx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributedInference -/src/neuronx_distributed_inference/modules/attention/utils.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.custom_c -alls import neuron_cumsum /home/ec2-user/NeuronxSDK/NeuronxDistributedInference -/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference -/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference -/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_ -llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_ -llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_ -llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio -n.attention_base import NeuronAttentionBase [2025-07-26 00:36:07.602: I neuronx_distributed/paral -lel_layers/parallel_state.py:628] > initializing tensor model parallel with size 1 [2025-07-26 00:36:07.602: I neuronx_distributed/paral -lel_layers/parallel_state.py:629] > initializing pipeline model parallel with size 1 [2025-07-26 00:36:07.603: I neuronx_distributed/paral -lel_layers/parallel_state.py:630] > initializing context model parallel with size 1 [2025-07-26 00:36:07.603: I neuronx_distributed/paral -lel_layers/parallel_state.py:631] > initializing data parallel with size 1 [2025-07-26 00:36:07.603: I neuronx_distributed/paral -lel_layers/parallel_state.py:632] > initializing world size to 1 2025-07-26 00:36:07.000608: 349190 INFO ||NEURON_CC -_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/neuronxcc-2.19.8089.0+8ab9f450/MODULE_13537832482499839909+e30acd3a/model.neff [2025-07-26 00:36:07.649: I neuronx_distributed/paral +/src/neuronx_distributed_inference/modules/attention/utils.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.custom_c +alls import neuron_cumsum /home/ec2-user/NeuronxSDK/NeuronxDistributedInference +/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference +/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference +/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx* +llama3/modeling*llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx* +llama3/modeling*llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx* +llama3/modeling*llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attentio +n.attention_base import NeuronAttentionBase [2025-07-26 00:36:07.602: I neuronx_distributed/paral +lel_layers/parallel_state.py:628] > initializing tensor model parallel with size 1 [2025-07-26 00:36:07.602: I neuronx_distributed/paral +lel_layers/parallel_state.py:629] > initializing pipeline model parallel with size 1 [2025-07-26 00:36:07.603: I neuronx_distributed/paral +lel_layers/parallel_state.py:630] > initializing context model parallel with size 1 [2025-07-26 00:36:07.603: I neuronx_distributed/paral +lel_layers/parallel_state.py:631] > initializing data parallel with size 1 [2025-07-26 00:36:07.603: I neuronx_distributed/paral +lel_layers/parallel_state.py:632] > initializing world size to 1 2025-07-26 00:36:07.000608: 349190 INFO ||NEURON_CC +\_WRAPPER||: Using a cached neff at /var/tmp/neuron-compile-cache/neuronxcc-2.19.8089.0+8ab9f450/MODULE_13537832482499839909+e30acd3a/model.neff [2025-07-26 00:36:07.649: I neuronx_distributed/paral lel_layers/parallel_state.py:379] [rank_0_pp-1_tp-1_dp-1_cp-1] Chosen Logic for replica groups ret_logic=, 'Ascending Ring PG Group')> [2025-07-26 00:36:07.650: I neuronx_distributed/paral -lel_layers/parallel_state.py:668] [rank_0_pp-1_tp-1_dp-1_cp-1] tp_groups: replica_groups.tp_groups=[[0]] [2025-07-26 00:36:07.650: I neuronx_distributed/paral -lel_layers/parallel_state.py:669] [rank_0_pp-1_tp-1_dp-1_cp-1] dp_groups: replica_groups.dp_groups=[[0]] [2025-07-26 00:36:07.650: I neuronx_distributed/paral -lel_layers/parallel_state.py:670] [rank_0_pp-1_tp-1_dp-1_cp-1] pp_groups: replica_groups.pp_groups=[[0]] [2025-07-26 00:36:07.650: I neuronx_distributed/paral -lel_layers/parallel_state.py:671] [rank_0_pp-1_tp-1_dp-1_cp-1] cp_groups: replica_groups.cp_groups=[[0]] [2025-07-26 00:36:07.650: I neuronx_distributed/paral -lel_layers/parallel_state.py:672] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_model_groups: replica_groups.ep_model_groups=[[0]] [2025-07-26 00:36:07.650: I neuronx_distributed/paral -lel_layers/parallel_state.py:673] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_data_groups: replica_groups.ep_data_groups=[[0]] WARNING:__main__:Unexpected keys in state dict: ['emb -ed_tokens.weight', 'norm.weight', 'lm_head.weight', 'layers.0.self_attn.q_proj.weight', 'layers.0.self_attn.k_proj.weight', 'layers.0.self_attn.v_proj.weight', 'layers.0.self_attn.o_proj.weight', 'layers.0.mlp.gate_proj.weight', 'layers.0.mlp.down_proj.weight', 'layers.0.mlp.up_proj.weight', 'layers.0.input_layernorm.weight', 'layers.0.post_attention_layernorm.weight', 'layers.1.self_attn.q_proj.weight', 'layers.1.self_attn.k_proj.weight', 'layers.1.self_attn.v_proj.weight', 'layers.1.self_attn.o_proj.weight', 'layers.1.mlp.gate_proj.weight', 'layers.1.mlp.down_proj.weight', 'layers.1.mlp.up_proj.weight', 'layers.1.input_layernorm.weight', 'layers.1.post_attention_layernorm.weight', 'layers.2.self_attn.q_proj.weight', 'layers.2.self_attn.k_proj.weight', 'layers.2.self_attn.v_proj.weight', 'layers.2.self_attn.o_proj.weight', 'layers.2.mlp.gate_proj.weight', 'layers.2.mlp.down_proj.weight', 'layers.2.mlp.up_proj.weight', 'layers.2.input_layernorm.weight', 'layers.2.post_attention_layernorm.weight', 'layers.3.self_attn.q_proj.weight', 'layers.3.self_attn.k_proj.weight', 'layers.3.self_attn.v_proj.weight', 'layers.3.self_attn.o_proj.weight', 'layers.3.mlp.gate_proj.weight', 'layers.3.mlp.down_proj.weight', 'layers.3.mlp.up_proj.weight', 'layers.3.input_layernorm.weight', 'layers.3.post_attention_layernorm.weight', 'layers.4.self_attn.q_proj.weight', 'layers.4.self_attn.k_proj.weight', 'layers.4.self_attn.v_proj.weight', 'layers.4.self_attn.o_proj.weight', 'layers.4.mlp.gate_proj.weight', 'layers.4.mlp.down_proj.weight', 'layers.4.mlp.up_proj.weight', 'layers.4.input_layernorm.weight', 'layers.4.post_attention_layernorm.weight', 'layers.5.self_attn.q_proj.weight', 'layers.5.self_attn.k_proj.weight', 'layers.5.self_attn.v_proj.weight', 'layers.5.self_attn.o_proj.weight', 'layers.5.mlp.gate_proj.weight', 'layers.5.mlp.down_proj.weight', 'layers.5.mlp.up_proj.weight', 'layers.5.input_layernorm.weight', 'layers.5.post_attention_layernorm.weight', 'layers.6.self_attn.q_proj.weight', 'layers.6.self_attn.k_proj.weight', 'layers.6.self_attn.v_proj.weight', 'layers.6.self_attn.o_proj.weight', 'layers.6.mlp.gate_proj.weight', 'layers.6.mlp.down_proj.weight', 'layers.6.mlp.up_proj.weight', 'layers.6.input_layernorm.weight', 'layers.6.post_attention_layernorm.weight', 'layers.7.self_attn.q_proj.weight', 'layers.7.self_attn.k_proj.weight', 'layers.7.self_attn.v_proj.weight', 'layers.7.self_attn.o_proj.weight', 'layers.7.mlp.gate_proj.weight', 'layers.7.mlp.down_proj.weight', 'layers.7.mlp.up_proj.weight', 'layers.7.input_layernorm.weight', 'layers.7.post_attention_layernorm.weight', 'layers.8.self_attn.q_proj.weight', 'layers.8.self_attn.k_proj.weight', 'layers.8.self_attn.v_proj.weight', 'layers.8.self_attn.o_proj.weight', 'layers.8.mlp.gate_proj.weight', 'layers.8.mlp.down_proj.weight', 'layers.8.mlp.up_proj.weight', 'layers.8.input_layernorm.weight', 'layers.8.post_attention_layernorm.weight', 'layers.9.self_attn.q_proj.weight', 'layers.9.self_attn.k_proj.weight', 'layers.9.self_attn.v_proj.weight', 'layers.9.self_attn.o_proj.weight', 'layers.9.mlp.gate_proj.weight', 'layers.9.mlp.down_proj.weight', 'layers.9.mlp.up_proj.weight', 'layers.9.input_layernorm.weight', 'layers.9.post_attention_layernorm.weight', 'layers.10.self_attn.q_proj.weight', 'layers.10.self_attn.k_proj.weight', 'layers.10.self_attn.v_proj.weight', 'layers.10.self_attn.o_proj.weight', 'layers.10.mlp.gate_proj.weight', 'layers.10.mlp.down_proj.weight', 'layers.10.mlp.up_proj.weight', 'layers.10.input_layernorm.weight', 'layers.10.post_attention_layernorm.weight', 'layers.11.self_attn.q_proj.weight', 'layers.11.self_attn.k_proj.weight', 'layers.11.self_attn.v_proj.weight', 'layers.11.self_attn.o_proj.weight', 'layers.11.mlp.gate_proj.weight', 'layers.11.mlp.down_proj.weight', 'layers.11.mlp.up_proj.weight', 'layers.11.input_layernorm.weight', 'layers.11.post_attention_layernorm.weight', 'layers.12.self_attn.q_proj.weight', 'layers.12.self_attn.k_proj.weight', 'layers.12.self_attn.v_proj.weight', 'layers.12.self_attn.o_proj.weight', 'layers.12.mlp.gate_proj.weight', 'layers.12.mlp.down_proj.weight', 'layers.12.mlp.up_proj.weight', 'layers.12.input_layernorm.weight', 'layers.12.post_attention_layernorm.weight', 'layers.13.self_attn.q_proj.weight', 'layers.13.self_attn.k_proj.weight', 'layers.13.self_attn.v_proj.weight', 'layers.13.self_attn.o_proj.weight', 'layers.13.mlp.gate_proj.weight', 'layers.13.mlp.down_proj.weight', 'layers.13.mlp.up_proj.weight', 'layers.13.input_layernorm.weight', 'layers.13.post_attention_layernorm.weight', 'layers.14.self_attn.q_proj.weight', 'layers.14.self_attn.k_proj.weight', 'layers.14.self_attn.v_proj.weight', 'layers.14.self_attn.o_proj.weight', 'layers.14.mlp.gate_proj.weight', 'layers.14.mlp.down_proj.weight', 'layers.14.mlp.up_proj.weight', 'layers.14.input_layernorm.weight', 'layers.14.post_attention_layernorm.weight', 'layers.15.self_attn.q_proj.weight', 'layers.15.self_attn.k_proj.weight', 'layers.15.self_attn.v_proj.weight', 'layers.15.self_attn.o_proj.weight', 'layers.15.mlp.gate_proj.weight', 'layers.15.mlp.down_proj.weight', 'layers.15.mlp.up_proj.weight', 'layers.15.input_layernorm.weight', 'layers.15.post_attention_layernorm.weight', 'layers.0.self_attn.rank_util.rank', 'layers.1.self_attn.rank_util.rank', 'layers.2.self_attn.rank_util.rank', 'layers.3.self_attn.rank_util.rank', 'layers.4.self_attn.rank_util.rank', 'layers.5.self_attn.rank_util.rank', 'layers.6.self_attn.rank_util.rank', 'layers.7.self_attn.rank_util.rank', 'layers.8.self_attn.rank_util.rank', 'layers.9.self_attn.rank_util.rank', 'layers.10.self_attn.rank_util.rank', 'layers.11.self_attn.rank_util.rank', 'layers.12.self_attn.rank_util.rank', 'layers.13.self_attn.rank_util.rank', 'layers.14.self_attn.rank_util.rank', 'layers.15.self_attn.rank_util.rank', 'rank_util.rank'] INFO:Neuron:Saving the neuron_config to llama3_compil -ed/ INFO:Neuron:Generating HLOs for the following models: - ['context_encoding_model', 'token_generation_model'][2025-07-26 00:36:09.073: I neuronx_distributed/paral -lel_layers/parallel_state.py:628] > initializing tensor model parallel with size 1 [2025-07-26 00:36:09.073: I neuronx_distributed/paral -lel_layers/parallel_state.py:629] > initializing pipeline model parallel with size 1 [2025-07-26 00:36:09.073: I neuronx_distributed/paral -lel_layers/parallel_state.py:630] > initializing context model parallel with size 1 [2025-07-26 00:36:09.073: I neuronx_distributed/paral -lel_layers/parallel_state.py:631] > initializing data parallel with size 1 [2025-07-26 00:36:09.073: I neuronx_distributed/paral -lel_layers/parallel_state.py:632] > initializing world size to 1 [2025-07-26 00:36:09.074: I neuronx_distributed/paral -lel_layers/parallel_state.py:379] [rank_0_pp0_tp0_dp0_cp0] Chosen Logic for replica groups ret_logic=, 'Ascending Ring PG Group')> [2025-07-26 00:36:09.074: I neuronx_distributed/paral -lel_layers/parallel_state.py:668] [rank_0_pp0_tp0_dp0_cp0] tp_groups: replica_groups.tp_groups=[[0]] [2025-07-26 00:36:09.074: I neuronx_distributed/paral -lel_layers/parallel_state.py:669] [rank_0_pp0_tp0_dp0_cp0] dp_groups: replica_groups.dp_groups=[[0]] [2025-07-26 00:36:09.074: I neuronx_distributed/paral -lel_layers/parallel_state.py:670] [rank_0_pp0_tp0_dp0_cp0] pp_groups: replica_groups.pp_groups=[[0]] [2025-07-26 00:36:09.074: I neuronx_distributed/paral -lel_layers/parallel_state.py:671] [rank_0_pp0_tp0_dp0_cp0] cp_groups: replica_groups.cp_groups=[[0]] [2025-07-26 00:36:09.074: I neuronx_distributed/paral -lel_layers/parallel_state.py:672] [rank_0_pp0_tp0_dp0_cp0] ep_model_groups: replica_groups.ep_model_groups=[[0]] [2025-07-26 00:36:09.075: I neuronx_distributed/paral -lel_layers/parallel_state.py:673] [rank_0_pp0_tp0_dp0_cp0] ep_data_groups: replica_groups.ep_data_groups=[[0]] INFO:Neuron:Generating 1 hlos for key: context_encodi -ng_model INFO:Neuron:Started loading module context_encoding_m -odel WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not - divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! INFO:Neuron:Finished loading module context_encoding_ -model in 0.053829193115234375 seconds INFO:Neuron:generating HLO: context_encoding_model, i -nput example shape = torch.Size([1, 128]) ERROR:__main__:Model compilation failed: You cannot s -pecify both input_ids and inputs_embeds at the same time ERROR:__main__:Compilation failed: You cannot specify - both input_ids and inputs_embeds at the same time +lel_layers/parallel_state.py:668] [rank_0_pp-1_tp-1_dp-1_cp-1] tp_groups: replica_groups.tp_groups=[[0]] [2025-07-26 00:36:07.650: I neuronx_distributed/paral +lel_layers/parallel_state.py:669] [rank_0_pp-1_tp-1_dp-1_cp-1] dp_groups: replica_groups.dp_groups=[[0]] [2025-07-26 00:36:07.650: I neuronx_distributed/paral +lel_layers/parallel_state.py:670] [rank_0_pp-1_tp-1_dp-1_cp-1] pp_groups: replica_groups.pp_groups=[[0]] [2025-07-26 00:36:07.650: I neuronx_distributed/paral +lel_layers/parallel_state.py:671] [rank_0_pp-1_tp-1_dp-1_cp-1] cp_groups: replica_groups.cp_groups=[[0]] [2025-07-26 00:36:07.650: I neuronx_distributed/paral +lel_layers/parallel_state.py:672] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_model_groups: replica_groups.ep_model_groups=[[0]] [2025-07-26 00:36:07.650: I neuronx_distributed/paral +lel_layers/parallel_state.py:673] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_data_groups: replica_groups.ep_data_groups=[[0]] WARNING:**main**:Unexpected keys in state dict: ['emb +ed_tokens.weight', 'norm.weight', 'lm_head.weight', 'layers.0.self_attn.q_proj.weight', 'layers.0.self_attn.k_proj.weight', 'layers.0.self_attn.v_proj.weight', 'layers.0.self_attn.o_proj.weight', 'layers.0.mlp.gate_proj.weight', 'layers.0.mlp.down_proj.weight', 'layers.0.mlp.up_proj.weight', 'layers.0.input_layernorm.weight', 'layers.0.post_attention_layernorm.weight', 'layers.1.self_attn.q_proj.weight', 'layers.1.self_attn.k_proj.weight', 'layers.1.self_attn.v_proj.weight', 'layers.1.self_attn.o_proj.weight', 'layers.1.mlp.gate_proj.weight', 'layers.1.mlp.down_proj.weight', 'layers.1.mlp.up_proj.weight', 'layers.1.input_layernorm.weight', 'layers.1.post_attention_layernorm.weight', 'layers.2.self_attn.q_proj.weight', 'layers.2.self_attn.k_proj.weight', 'layers.2.self_attn.v_proj.weight', 'layers.2.self_attn.o_proj.weight', 'layers.2.mlp.gate_proj.weight', 'layers.2.mlp.down_proj.weight', 'layers.2.mlp.up_proj.weight', 'layers.2.input_layernorm.weight', 'layers.2.post_attention_layernorm.weight', 'layers.3.self_attn.q_proj.weight', 'layers.3.self_attn.k_proj.weight', 'layers.3.self_attn.v_proj.weight', 'layers.3.self_attn.o_proj.weight', 'layers.3.mlp.gate_proj.weight', 'layers.3.mlp.down_proj.weight', 'layers.3.mlp.up_proj.weight', 'layers.3.input_layernorm.weight', 'layers.3.post_attention_layernorm.weight', 'layers.4.self_attn.q_proj.weight', 'layers.4.self_attn.k_proj.weight', 'layers.4.self_attn.v_proj.weight', 'layers.4.self_attn.o_proj.weight', 'layers.4.mlp.gate_proj.weight', 'layers.4.mlp.down_proj.weight', 'layers.4.mlp.up_proj.weight', 'layers.4.input_layernorm.weight', 'layers.4.post_attention_layernorm.weight', 'layers.5.self_attn.q_proj.weight', 'layers.5.self_attn.k_proj.weight', 'layers.5.self_attn.v_proj.weight', 'layers.5.self_attn.o_proj.weight', 'layers.5.mlp.gate_proj.weight', 'layers.5.mlp.down_proj.weight', 'layers.5.mlp.up_proj.weight', 'layers.5.input_layernorm.weight', 'layers.5.post_attention_layernorm.weight', 'layers.6.self_attn.q_proj.weight', 'layers.6.self_attn.k_proj.weight', 'layers.6.self_attn.v_proj.weight', 'layers.6.self_attn.o_proj.weight', 'layers.6.mlp.gate_proj.weight', 'layers.6.mlp.down_proj.weight', 'layers.6.mlp.up_proj.weight', 'layers.6.input_layernorm.weight', 'layers.6.post_attention_layernorm.weight', 'layers.7.self_attn.q_proj.weight', 'layers.7.self_attn.k_proj.weight', 'layers.7.self_attn.v_proj.weight', 'layers.7.self_attn.o_proj.weight', 'layers.7.mlp.gate_proj.weight', 'layers.7.mlp.down_proj.weight', 'layers.7.mlp.up_proj.weight', 'layers.7.input_layernorm.weight', 'layers.7.post_attention_layernorm.weight', 'layers.8.self_attn.q_proj.weight', 'layers.8.self_attn.k_proj.weight', 'layers.8.self_attn.v_proj.weight', 'layers.8.self_attn.o_proj.weight', 'layers.8.mlp.gate_proj.weight', 'layers.8.mlp.down_proj.weight', 'layers.8.mlp.up_proj.weight', 'layers.8.input_layernorm.weight', 'layers.8.post_attention_layernorm.weight', 'layers.9.self_attn.q_proj.weight', 'layers.9.self_attn.k_proj.weight', 'layers.9.self_attn.v_proj.weight', 'layers.9.self_attn.o_proj.weight', 'layers.9.mlp.gate_proj.weight', 'layers.9.mlp.down_proj.weight', 'layers.9.mlp.up_proj.weight', 'layers.9.input_layernorm.weight', 'layers.9.post_attention_layernorm.weight', 'layers.10.self_attn.q_proj.weight', 'layers.10.self_attn.k_proj.weight', 'layers.10.self_attn.v_proj.weight', 'layers.10.self_attn.o_proj.weight', 'layers.10.mlp.gate_proj.weight', 'layers.10.mlp.down_proj.weight', 'layers.10.mlp.up_proj.weight', 'layers.10.input_layernorm.weight', 'layers.10.post_attention_layernorm.weight', 'layers.11.self_attn.q_proj.weight', 'layers.11.self_attn.k_proj.weight', 'layers.11.self_attn.v_proj.weight', 'layers.11.self_attn.o_proj.weight', 'layers.11.mlp.gate_proj.weight', 'layers.11.mlp.down_proj.weight', 'layers.11.mlp.up_proj.weight', 'layers.11.input_layernorm.weight', 'layers.11.post_attention_layernorm.weight', 'layers.12.self_attn.q_proj.weight', 'layers.12.self_attn.k_proj.weight', 'layers.12.self_attn.v_proj.weight', 'layers.12.self_attn.o_proj.weight', 'layers.12.mlp.gate_proj.weight', 'layers.12.mlp.down_proj.weight', 'layers.12.mlp.up_proj.weight', 'layers.12.input_layernorm.weight', 'layers.12.post_attention_layernorm.weight', 'layers.13.self_attn.q_proj.weight', 'layers.13.self_attn.k_proj.weight', 'layers.13.self_attn.v_proj.weight', 'layers.13.self_attn.o_proj.weight', 'layers.13.mlp.gate_proj.weight', 'layers.13.mlp.down_proj.weight', 'layers.13.mlp.up_proj.weight', 'layers.13.input_layernorm.weight', 'layers.13.post_attention_layernorm.weight', 'layers.14.self_attn.q_proj.weight', 'layers.14.self_attn.k_proj.weight', 'layers.14.self_attn.v_proj.weight', 'layers.14.self_attn.o_proj.weight', 'layers.14.mlp.gate_proj.weight', 'layers.14.mlp.down_proj.weight', 'layers.14.mlp.up_proj.weight', 'layers.14.input_layernorm.weight', 'layers.14.post_attention_layernorm.weight', 'layers.15.self_attn.q_proj.weight', 'layers.15.self_attn.k_proj.weight', 'layers.15.self_attn.v_proj.weight', 'layers.15.self_attn.o_proj.weight', 'layers.15.mlp.gate_proj.weight', 'layers.15.mlp.down_proj.weight', 'layers.15.mlp.up_proj.weight', 'layers.15.input_layernorm.weight', 'layers.15.post_attention_layernorm.weight', 'layers.0.self_attn.rank_util.rank', 'layers.1.self_attn.rank_util.rank', 'layers.2.self_attn.rank_util.rank', 'layers.3.self_attn.rank_util.rank', 'layers.4.self_attn.rank_util.rank', 'layers.5.self_attn.rank_util.rank', 'layers.6.self_attn.rank_util.rank', 'layers.7.self_attn.rank_util.rank', 'layers.8.self_attn.rank_util.rank', 'layers.9.self_attn.rank_util.rank', 'layers.10.self_attn.rank_util.rank', 'layers.11.self_attn.rank_util.rank', 'layers.12.self_attn.rank_util.rank', 'layers.13.self_attn.rank_util.rank', 'layers.14.self_attn.rank_util.rank', 'layers.15.self_attn.rank_util.rank', 'rank_util.rank'] INFO:Neuron:Saving the neuron_config to llama3_compil +ed/ INFO:Neuron:Generating HLOs for the following models: +['context_encoding_model', 'token_generation_model'][2025-07-26 00:36:09.073: I neuronx_distributed/paral lel_layers/parallel_state.py:628] > initializing tensor model parallel with size 1 [2025-07-26 00:36:09.073: I neuronx_distributed/paral +lel_layers/parallel_state.py:629] > initializing pipeline model parallel with size 1 [2025-07-26 00:36:09.073: I neuronx_distributed/paral +lel_layers/parallel_state.py:630] > initializing context model parallel with size 1 [2025-07-26 00:36:09.073: I neuronx_distributed/paral +lel_layers/parallel_state.py:631] > initializing data parallel with size 1 [2025-07-26 00:36:09.073: I neuronx_distributed/paral +lel_layers/parallel_state.py:632] > initializing world size to 1 [2025-07-26 00:36:09.074: I neuronx_distributed/paral +lel_layers/parallel_state.py:379] [rank_0_pp0_tp0_dp0_cp0] Chosen Logic for replica groups ret_logic=, 'Ascending Ring PG Group')> [2025-07-26 00:36:09.074: I neuronx_distributed/paral +lel_layers/parallel_state.py:668] [rank_0_pp0_tp0_dp0_cp0] tp_groups: replica_groups.tp_groups=[[0]] [2025-07-26 00:36:09.074: I neuronx_distributed/paral +lel_layers/parallel_state.py:669] [rank_0_pp0_tp0_dp0_cp0] dp_groups: replica_groups.dp_groups=[[0]] [2025-07-26 00:36:09.074: I neuronx_distributed/paral +lel_layers/parallel_state.py:670] [rank_0_pp0_tp0_dp0_cp0] pp_groups: replica_groups.pp_groups=[[0]] [2025-07-26 00:36:09.074: I neuronx_distributed/paral +lel_layers/parallel_state.py:671] [rank_0_pp0_tp0_dp0_cp0] cp_groups: replica_groups.cp_groups=[[0]] [2025-07-26 00:36:09.074: I neuronx_distributed/paral +lel_layers/parallel_state.py:672] [rank_0_pp0_tp0_dp0_cp0] ep_model_groups: replica_groups.ep_model_groups=[[0]] [2025-07-26 00:36:09.075: I neuronx_distributed/paral +lel_layers/parallel_state.py:673] [rank_0_pp0_tp0_dp0_cp0] ep_data_groups: replica_groups.ep_data_groups=[[0]] INFO:Neuron:Generating 1 hlos for key: context_encodi +ng_model INFO:Neuron:Started loading module context_encoding_m +odel WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not +divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! INFO:Neuron:Finished loading module context_encoding* +model in 0.053829193115234375 seconds INFO:Neuron:generating HLO: context_encoding_model, i +nput example shape = torch.Size([1, 128]) ERROR:**main**:Model compilation failed: You cannot s +pecify both input_ids and inputs_embeds at the same time ERROR:**main**:Compilation failed: You cannot specify +both input_ids and inputs_embeds at the same time Great progress! The model is loading successfully and the GQA conversion is working. The error is about conflicting input parameters. This is likely in the forward method. Let me check and fix the forward method signature: Read file(s) @@ -1878,16 +1862,9 @@ Session Too Long Continue Summarize and continue in a new session. - - - - - - Claude Sonnet 4.0 Autopilot - Session #2 New Session @@ -1895,10 +1872,6 @@ New Session (Continu... New Session (Continu... - - - - Kiro Understood. I have the context from our previous conversation and am ready to continue helping you. @@ -1911,7 +1884,6 @@ Created INFERENCE_IMPLEMENTATION_SUMMARY.md - I've created a comprehensive summary document that captures the entire journey from the previous context. The document includes: Overview of the implementation goal @@ -1926,13 +1898,13 @@ The summary shows that we achieved complete success - the model loads properly, here is the error: [1/1] Prompt: Hello, how are you? ----------------------------------------- +--- -ERROR:__main__:Tokenizer is required for text generation +ERROR:**main**:Tokenizer is required for text generation Generated: Error: No tokenizer available --------------------------------------- +--- Kiro I see the issue. The inference script is still failing because it can't load the tokenizer properly. Let me check the current state of the inference script and fix the tokenizer loading issue. @@ -1958,59 +1930,58 @@ Accepted edits to run_inference.py - Now let's test the inference again to see if the tokenizer loads properly: Command python run_inference.py --model_path ./llama3_compiled --prompt "Hello, how are you?" --max_new_tokens 5 /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuro -nx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +nx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuro -nx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +nx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuro -nx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +nx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuron -x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuron -x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuron -x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuron -x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuron -x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuron -x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inference/modules/attention/utils.py:14: Deprecation -Warning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.custom_calls import neuron_cumsum +Warning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.custom_calls import neuron_cumsum /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: Dep -recationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttention_QKV +recationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: Dep -recationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttention_QKV +recationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: Dep -recationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttention_QKV +recationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is de -precated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base import NeuronAttentionBase +precated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is de -precated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base import NeuronAttentionBase +precated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is de -precated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base import NeuronAttentionBase +precated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base import NeuronAttentionBase WARNING:root:NeuronConfig init: Unexpected keyword arguments: {'apply_seq_ids_mask': False, 'enable_long_context_mode': False, 'en -able_output_completion_notifications': False, 'enable_token_tree': False, 'is_chunked_prefill': False, 'is_prefill_stage': None, 'kv_cache_tiling': False, 'scratchpad_page_size': None, 'skip_warmup': False, 'tile_cc': False, 'weights_to_skip_layout_optimization': []} INFO:Neuron:Sharding weights on load... +able_output_completion_notifications': False, 'enable_token_tree': False, 'is_chunked_prefill': False, 'is_prefill_stage': None, 'kv_cache_tiling': False, 'scratchpad_page_size': None, 'skip_warmup': False, 'tile_cc': False, 'weights_to_skip_layout_optimization': []} INFO:Neuron:Sharding weights on load... INFO:Neuron:Sharding Weights for ranks: 0...0 [2025-07-26 16:57:46.225: I neuronx_distributed/parallel_layers/parallel_state.py:628] > initializing tensor model parallel with s -ize 1 [2025-07-26 16:57:46.225: I neuronx_distributed/parallel_layers/parallel_state.py:629] > initializing pipeline model parallel with - size 1 [2025-07-26 16:57:46.225: I neuronx_distributed/parallel_layers/parallel_state.py:630] > initializing context model parallel with -size 1 [2025-07-26 16:57:46.225: I neuronx_distributed/parallel_layers/parallel_state.py:631] > initializing data parallel with size 1 +ize 1 [2025-07-26 16:57:46.225: I neuronx_distributed/parallel_layers/parallel_state.py:629] > initializing pipeline model parallel with +size 1 [2025-07-26 16:57:46.225: I neuronx_distributed/parallel_layers/parallel_state.py:630] > initializing context model parallel with +size 1 [2025-07-26 16:57:46.225: I neuronx_distributed/parallel_layers/parallel_state.py:631] > initializing data parallel with size 1 [2025-07-26 16:57:46.225: I neuronx_distributed/parallel_layers/parallel_state.py:632] > initializing world size to 1 [2025-07-26 16:57:46.226: I neuronx_distributed/parallel_layers/parallel_state.py:379] [rank_0_pp-1_tp-1_dp-1_cp-1] Chosen Logic f -or replica groups ret_logic=, 'Ascending Ring PG Group')> [2025-07-26 16:57:46.226: I neuronx_distributed/parallel_layers/parallel_state.py:668] [rank_0_pp-1_tp-1_dp-1_cp-1] tp_groups: rep -lica_groups.tp_groups=[[0]] [2025-07-26 16:57:46.226: I neuronx_distributed/parallel_layers/parallel_state.py:669] [rank_0_pp-1_tp-1_dp-1_cp-1] dp_groups: rep -lica_groups.dp_groups=[[0]] [2025-07-26 16:57:46.227: I neuronx_distributed/parallel_layers/parallel_state.py:670] [rank_0_pp-1_tp-1_dp-1_cp-1] pp_groups: rep -lica_groups.pp_groups=[[0]] [2025-07-26 16:57:46.227: I neuronx_distributed/parallel_layers/parallel_state.py:671] [rank_0_pp-1_tp-1_dp-1_cp-1] cp_groups: rep -lica_groups.cp_groups=[[0]] [2025-07-26 16:57:46.227: I neuronx_distributed/parallel_layers/parallel_state.py:672] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_model_group -s: replica_groups.ep_model_groups=[[0]] [2025-07-26 16:57:46.227: I neuronx_distributed/parallel_layers/parallel_state.py:673] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_data_groups -: replica_groups.ep_data_groups=[[0]] WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! +or replica groups ret_logic=, 'Ascending Ring PG Group')> [2025-07-26 16:57:46.226: I neuronx_distributed/parallel_layers/parallel_state.py:668] [rank_0_pp-1_tp-1_dp-1_cp-1] tp_groups: rep +lica_groups.tp_groups=[[0]] [2025-07-26 16:57:46.226: I neuronx_distributed/parallel_layers/parallel_state.py:669] [rank_0_pp-1_tp-1_dp-1_cp-1] dp_groups: rep +lica_groups.dp_groups=[[0]] [2025-07-26 16:57:46.227: I neuronx_distributed/parallel_layers/parallel_state.py:670] [rank_0_pp-1_tp-1_dp-1_cp-1] pp_groups: rep +lica_groups.pp_groups=[[0]] [2025-07-26 16:57:46.227: I neuronx_distributed/parallel_layers/parallel_state.py:671] [rank_0_pp-1_tp-1_dp-1_cp-1] cp_groups: rep +lica_groups.cp_groups=[[0]] [2025-07-26 16:57:46.227: I neuronx_distributed/parallel_layers/parallel_state.py:672] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_model_group +s: replica_groups.ep_model_groups=[[0]] [2025-07-26 16:57:46.227: I neuronx_distributed/parallel_layers/parallel_state.py:673] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_data_groups +: replica_groups.ep_data_groups=[[0]] WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! @@ -2044,7 +2015,7 @@ WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding atte WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/trace/trace.py:640: UserWarning: Removing redundant keys from - checkpoint: ['layers.0.self_attn.q_proj.weight', 'layers.0.self_attn.k_proj.weight', 'layers.0.self_attn.v_proj.weight', 'layers.0.self_attn.o_proj.weight', 'layers.1.self_attn.q_proj.weight', 'layers.1.self_attn.k_proj.weight', 'layers.1.self_attn.v_proj.weight', 'layers.1.self_attn.o_proj.weight', 'layers.2.self_attn.q_proj.weight', 'layers.2.self_attn.k_proj.weight', 'layers.2.self_attn.v_proj.weight', 'layers.2.self_attn.o_proj.weight', 'layers.3.self_attn.q_proj.weight', 'layers.3.self_attn.k_proj.weight', 'layers.3.self_attn.v_proj.weight', 'layers.3.self_attn.o_proj.weight', 'layers.4.self_attn.q_proj.weight', 'layers.4.self_attn.k_proj.weight', 'layers.4.self_attn.v_proj.weight', 'layers.4.self_attn.o_proj.weight', 'layers.5.self_attn.q_proj.weight', 'layers.5.self_attn.k_proj.weight', 'layers.5.self_attn.v_proj.weight', 'layers.5.self_attn.o_proj.weight', 'layers.6.self_attn.q_proj.weight', 'layers.6.self_attn.k_proj.weight', 'layers.6.self_attn.v_proj.weight', 'layers.6.self_attn.o_proj.weight', 'layers.7.self_attn.q_proj.weight', 'layers.7.self_attn.k_proj.weight', 'layers.7.self_attn.v_proj.weight', 'layers.7.self_attn.o_proj.weight', 'layers.8.self_attn.q_proj.weight', 'layers.8.self_attn.k_proj.weight', 'layers.8.self_attn.v_proj.weight', 'layers.8.self_attn.o_proj.weight', 'layers.9.self_attn.q_proj.weight', 'layers.9.self_attn.k_proj.weight', 'layers.9.self_attn.v_proj.weight', 'layers.9.self_attn.o_proj.weight', 'layers.10.self_attn.q_proj.weight', 'layers.10.self_attn.k_proj.weight', 'layers.10.self_attn.v_proj.weight', 'layers.10.self_attn.o_proj.weight', 'layers.11.self_attn.q_proj.weight', 'layers.11.self_attn.k_proj.weight', 'layers.11.self_attn.v_proj.weight', 'layers.11.self_attn.o_proj.weight', 'layers.12.self_attn.q_proj.weight', 'layers.12.self_attn.k_proj.weight', 'layers.12.self_attn.v_proj.weight', 'layers.12.self_attn.o_proj.weight', 'layers.13.self_attn.q_proj.weight', 'layers.13.self_attn.k_proj.weight', 'layers.13.self_attn.v_proj.weight', 'layers.13.self_attn.o_proj.weight', 'layers.14.self_attn.q_proj.weight', 'layers.14.self_attn.k_proj.weight', 'layers.14.self_attn.v_proj.weight', 'layers.14.self_attn.o_proj.weight', 'layers.15.self_attn.q_proj.weight', 'layers.15.self_attn.k_proj.weight', 'layers.15.self_attn.v_proj.weight', 'layers.15.self_attn.o_proj.weight'] warnings.warn(f"Removing redundant keys from checkpoint: {keys_to_delete}") +checkpoint: ['layers.0.self_attn.q_proj.weight', 'layers.0.self_attn.k_proj.weight', 'layers.0.self_attn.v_proj.weight', 'layers.0.self_attn.o_proj.weight', 'layers.1.self_attn.q_proj.weight', 'layers.1.self_attn.k_proj.weight', 'layers.1.self_attn.v_proj.weight', 'layers.1.self_attn.o_proj.weight', 'layers.2.self_attn.q_proj.weight', 'layers.2.self_attn.k_proj.weight', 'layers.2.self_attn.v_proj.weight', 'layers.2.self_attn.o_proj.weight', 'layers.3.self_attn.q_proj.weight', 'layers.3.self_attn.k_proj.weight', 'layers.3.self_attn.v_proj.weight', 'layers.3.self_attn.o_proj.weight', 'layers.4.self_attn.q_proj.weight', 'layers.4.self_attn.k_proj.weight', 'layers.4.self_attn.v_proj.weight', 'layers.4.self_attn.o_proj.weight', 'layers.5.self_attn.q_proj.weight', 'layers.5.self_attn.k_proj.weight', 'layers.5.self_attn.v_proj.weight', 'layers.5.self_attn.o_proj.weight', 'layers.6.self_attn.q_proj.weight', 'layers.6.self_attn.k_proj.weight', 'layers.6.self_attn.v_proj.weight', 'layers.6.self_attn.o_proj.weight', 'layers.7.self_attn.q_proj.weight', 'layers.7.self_attn.k_proj.weight', 'layers.7.self_attn.v_proj.weight', 'layers.7.self_attn.o_proj.weight', 'layers.8.self_attn.q_proj.weight', 'layers.8.self_attn.k_proj.weight', 'layers.8.self_attn.v_proj.weight', 'layers.8.self_attn.o_proj.weight', 'layers.9.self_attn.q_proj.weight', 'layers.9.self_attn.k_proj.weight', 'layers.9.self_attn.v_proj.weight', 'layers.9.self_attn.o_proj.weight', 'layers.10.self_attn.q_proj.weight', 'layers.10.self_attn.k_proj.weight', 'layers.10.self_attn.v_proj.weight', 'layers.10.self_attn.o_proj.weight', 'layers.11.self_attn.q_proj.weight', 'layers.11.self_attn.k_proj.weight', 'layers.11.self_attn.v_proj.weight', 'layers.11.self_attn.o_proj.weight', 'layers.12.self_attn.q_proj.weight', 'layers.12.self_attn.k_proj.weight', 'layers.12.self_attn.v_proj.weight', 'layers.12.self_attn.o_proj.weight', 'layers.13.self_attn.q_proj.weight', 'layers.13.self_attn.k_proj.weight', 'layers.13.self_attn.v_proj.weight', 'layers.13.self_attn.o_proj.weight', 'layers.14.self_attn.q_proj.weight', 'layers.14.self_attn.k_proj.weight', 'layers.14.self_attn.v_proj.weight', 'layers.14.self_attn.o_proj.weight', 'layers.15.self_attn.q_proj.weight', 'layers.15.self_attn.k_proj.weight', 'layers.15.self_attn.v_proj.weight', 'layers.15.self_attn.o_proj.weight'] warnings.warn(f"Removing redundant keys from checkpoint: {keys_to_delete}") WARNING:Neuron:casting embed_tokens.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.0.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.0.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to torch.float32 @@ -2196,11 +2167,11 @@ INFO:Neuron:Done Sharding weights in 0.5069742330815643 INFO:Neuron:Finished weights loading in 12.252695644972846 seconds INFO:Neuron:Warming up the model. INFO:Neuron:Warmup completed in 0.0524439811706543 seconds. -WARNING:__main__:AutoTokenizer failed: The checkpoint you are trying to load has model type `llama3_neuron` but Transformers does -not recognize this architecture. This could be because of an issue with the checkpoint, or because your version of Transformers is out of date. +WARNING:**main**:AutoTokenizer failed: The checkpoint you are trying to load has model type `llama3_neuron` but Transformers does +not recognize this architecture. This could be because of an issue with the checkpoint, or because your version of Transformers is out of date. You can update Transformers with the command `pip install --upgrade transformers`. If this does not work, and the checkpoint is ve -ry new, then there may not be a release version that supports this model yet. In this case, you can get the most up-to-date code by installing Transformers from source with the command `pip install git+https://github.com/huggingface/transformers.git` You are using the default legacy behaviour of the . This is e -xpected, and simply means that the `legacy` (previous) behavior will be used so nothing changes for you. If you want to use the new behaviour, set `legacy=False`. This should only be set if you understand what it means, and thoroughly read the reason why this was added as explained in https://github.com/huggingface/transformers/pull/24565 - if you loaded a llama tokenizer from a GGUF file you can ignore this message +ry new, then there may not be a release version that supports this model yet. In this case, you can get the most up-to-date code by installing Transformers from source with the command `pip install git+https://github.com/huggingface/transformers.git` You are using the default legacy behaviour of the . This is e +xpected, and simply means that the `legacy` (previous) behavior will be used so nothing changes for you. If you want to use the new behaviour, set `legacy=False`. This should only be set if you understand what it means, and thoroughly read the reason why this was added as explained in https://github.com/huggingface/transformers/pull/24565 - if you loaded a llama tokenizer from a GGUF file you can ignore this message ============================================================ LLAMA3 INFERENCE RESULTS ============================================================ @@ -2208,9 +2179,9 @@ Model: ./llama3_compiled Generation settings: max_tokens=5, temperature=0.6, top_p=0.9 ============================================================ -[1/1] Prompt: Hello, how are you? ----------------------------------------- -ERROR:__main__:Inference failed: 'bool' object is not callable +## [1/1] Prompt: Hello, how are you? + +ERROR:**main**:Inference failed: 'bool' object is not callable Great progress! The tokenizer is now loading successfully (I can see the LlamaTokenizer warning which means it loaded), but there's a new error: 'bool' object is not callable. Let me check what's causing this error in the generation code. Searched workspace @@ -2228,11 +2199,11 @@ The issue might be in the model forward pass. Let me check if there's an issue w Canceled edits -this is the error: +this is the error: ----------------------------------------- +--- -ERROR:__main__:Inference failed: 'bool' object is not callable +ERROR:**main**:Inference failed: 'bool' object is not callable can you focus on that @@ -2282,59 +2253,58 @@ Accepted edits to run_inference.py - Now let's test the inference again with the proper HuggingFaceGenerationAdapter: Command python run_inference.py --model_path ./llama3_compiled --prompt "Hello, how are you?" --max_new_tokens 5 /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuro -nx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +nx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuro -nx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +nx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuro -nx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +nx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuron -x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuron -x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuron -x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuron -x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuron -x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuron -x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inference/modules/attention/utils.py:14: Deprecation -Warning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.custom_calls import neuron_cumsum +Warning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.custom_calls import neuron_cumsum /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: Dep -recationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttention_QKV +recationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: Dep -recationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttention_QKV +recationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: Dep -recationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttention_QKV +recationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is de -precated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base import NeuronAttentionBase +precated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is de -precated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base import NeuronAttentionBase +precated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is de -precated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base import NeuronAttentionBase +precated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base import NeuronAttentionBase WARNING:root:NeuronConfig init: Unexpected keyword arguments: {'apply_seq_ids_mask': False, 'enable_long_context_mode': False, 'en -able_output_completion_notifications': False, 'enable_token_tree': False, 'is_chunked_prefill': False, 'is_prefill_stage': None, 'kv_cache_tiling': False, 'scratchpad_page_size': None, 'skip_warmup': False, 'tile_cc': False, 'weights_to_skip_layout_optimization': []} INFO:Neuron:Sharding weights on load... +able_output_completion_notifications': False, 'enable_token_tree': False, 'is_chunked_prefill': False, 'is_prefill_stage': None, 'kv_cache_tiling': False, 'scratchpad_page_size': None, 'skip_warmup': False, 'tile_cc': False, 'weights_to_skip_layout_optimization': []} INFO:Neuron:Sharding weights on load... INFO:Neuron:Sharding Weights for ranks: 0...0 [2025-07-26 17:00:35.674: I neuronx_distributed/parallel_layers/parallel_state.py:628] > initializing tensor model parallel with s -ize 1 [2025-07-26 17:00:35.674: I neuronx_distributed/parallel_layers/parallel_state.py:629] > initializing pipeline model parallel with - size 1 [2025-07-26 17:00:35.674: I neuronx_distributed/parallel_layers/parallel_state.py:630] > initializing context model parallel with -size 1 [2025-07-26 17:00:35.674: I neuronx_distributed/parallel_layers/parallel_state.py:631] > initializing data parallel with size 1 +ize 1 [2025-07-26 17:00:35.674: I neuronx_distributed/parallel_layers/parallel_state.py:629] > initializing pipeline model parallel with +size 1 [2025-07-26 17:00:35.674: I neuronx_distributed/parallel_layers/parallel_state.py:630] > initializing context model parallel with +size 1 [2025-07-26 17:00:35.674: I neuronx_distributed/parallel_layers/parallel_state.py:631] > initializing data parallel with size 1 [2025-07-26 17:00:35.674: I neuronx_distributed/parallel_layers/parallel_state.py:632] > initializing world size to 1 [2025-07-26 17:00:35.675: I neuronx_distributed/parallel_layers/parallel_state.py:379] [rank_0_pp-1_tp-1_dp-1_cp-1] Chosen Logic f -or replica groups ret_logic=, 'Ascending Ring PG Group')> [2025-07-26 17:00:35.676: I neuronx_distributed/parallel_layers/parallel_state.py:668] [rank_0_pp-1_tp-1_dp-1_cp-1] tp_groups: rep -lica_groups.tp_groups=[[0]] [2025-07-26 17:00:35.676: I neuronx_distributed/parallel_layers/parallel_state.py:669] [rank_0_pp-1_tp-1_dp-1_cp-1] dp_groups: rep -lica_groups.dp_groups=[[0]] [2025-07-26 17:00:35.676: I neuronx_distributed/parallel_layers/parallel_state.py:670] [rank_0_pp-1_tp-1_dp-1_cp-1] pp_groups: rep -lica_groups.pp_groups=[[0]] [2025-07-26 17:00:35.676: I neuronx_distributed/parallel_layers/parallel_state.py:671] [rank_0_pp-1_tp-1_dp-1_cp-1] cp_groups: rep -lica_groups.cp_groups=[[0]] [2025-07-26 17:00:35.676: I neuronx_distributed/parallel_layers/parallel_state.py:672] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_model_group -s: replica_groups.ep_model_groups=[[0]] [2025-07-26 17:00:35.676: I neuronx_distributed/parallel_layers/parallel_state.py:673] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_data_groups -: replica_groups.ep_data_groups=[[0]] WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! +or replica groups ret_logic=, 'Ascending Ring PG Group')> [2025-07-26 17:00:35.676: I neuronx_distributed/parallel_layers/parallel_state.py:668] [rank_0_pp-1_tp-1_dp-1_cp-1] tp_groups: rep +lica_groups.tp_groups=[[0]] [2025-07-26 17:00:35.676: I neuronx_distributed/parallel_layers/parallel_state.py:669] [rank_0_pp-1_tp-1_dp-1_cp-1] dp_groups: rep +lica_groups.dp_groups=[[0]] [2025-07-26 17:00:35.676: I neuronx_distributed/parallel_layers/parallel_state.py:670] [rank_0_pp-1_tp-1_dp-1_cp-1] pp_groups: rep +lica_groups.pp_groups=[[0]] [2025-07-26 17:00:35.676: I neuronx_distributed/parallel_layers/parallel_state.py:671] [rank_0_pp-1_tp-1_dp-1_cp-1] cp_groups: rep +lica_groups.cp_groups=[[0]] [2025-07-26 17:00:35.676: I neuronx_distributed/parallel_layers/parallel_state.py:672] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_model_group +s: replica_groups.ep_model_groups=[[0]] [2025-07-26 17:00:35.676: I neuronx_distributed/parallel_layers/parallel_state.py:673] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_data_groups +: replica_groups.ep_data_groups=[[0]] WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! @@ -2368,7 +2338,7 @@ WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding atte WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/trace/trace.py:640: UserWarning: Removing redundant keys from - checkpoint: ['layers.0.self_attn.q_proj.weight', 'layers.0.self_attn.k_proj.weight', 'layers.0.self_attn.v_proj.weight', 'layers.0.self_attn.o_proj.weight', 'layers.1.self_attn.q_proj.weight', 'layers.1.self_attn.k_proj.weight', 'layers.1.self_attn.v_proj.weight', 'layers.1.self_attn.o_proj.weight', 'layers.2.self_attn.q_proj.weight', 'layers.2.self_attn.k_proj.weight', 'layers.2.self_attn.v_proj.weight', 'layers.2.self_attn.o_proj.weight', 'layers.3.self_attn.q_proj.weight', 'layers.3.self_attn.k_proj.weight', 'layers.3.self_attn.v_proj.weight', 'layers.3.self_attn.o_proj.weight', 'layers.4.self_attn.q_proj.weight', 'layers.4.self_attn.k_proj.weight', 'layers.4.self_attn.v_proj.weight', 'layers.4.self_attn.o_proj.weight', 'layers.5.self_attn.q_proj.weight', 'layers.5.self_attn.k_proj.weight', 'layers.5.self_attn.v_proj.weight', 'layers.5.self_attn.o_proj.weight', 'layers.6.self_attn.q_proj.weight', 'layers.6.self_attn.k_proj.weight', 'layers.6.self_attn.v_proj.weight', 'layers.6.self_attn.o_proj.weight', 'layers.7.self_attn.q_proj.weight', 'layers.7.self_attn.k_proj.weight', 'layers.7.self_attn.v_proj.weight', 'layers.7.self_attn.o_proj.weight', 'layers.8.self_attn.q_proj.weight', 'layers.8.self_attn.k_proj.weight', 'layers.8.self_attn.v_proj.weight', 'layers.8.self_attn.o_proj.weight', 'layers.9.self_attn.q_proj.weight', 'layers.9.self_attn.k_proj.weight', 'layers.9.self_attn.v_proj.weight', 'layers.9.self_attn.o_proj.weight', 'layers.10.self_attn.q_proj.weight', 'layers.10.self_attn.k_proj.weight', 'layers.10.self_attn.v_proj.weight', 'layers.10.self_attn.o_proj.weight', 'layers.11.self_attn.q_proj.weight', 'layers.11.self_attn.k_proj.weight', 'layers.11.self_attn.v_proj.weight', 'layers.11.self_attn.o_proj.weight', 'layers.12.self_attn.q_proj.weight', 'layers.12.self_attn.k_proj.weight', 'layers.12.self_attn.v_proj.weight', 'layers.12.self_attn.o_proj.weight', 'layers.13.self_attn.q_proj.weight', 'layers.13.self_attn.k_proj.weight', 'layers.13.self_attn.v_proj.weight', 'layers.13.self_attn.o_proj.weight', 'layers.14.self_attn.q_proj.weight', 'layers.14.self_attn.k_proj.weight', 'layers.14.self_attn.v_proj.weight', 'layers.14.self_attn.o_proj.weight', 'layers.15.self_attn.q_proj.weight', 'layers.15.self_attn.k_proj.weight', 'layers.15.self_attn.v_proj.weight', 'layers.15.self_attn.o_proj.weight'] warnings.warn(f"Removing redundant keys from checkpoint: {keys_to_delete}") +checkpoint: ['layers.0.self_attn.q_proj.weight', 'layers.0.self_attn.k_proj.weight', 'layers.0.self_attn.v_proj.weight', 'layers.0.self_attn.o_proj.weight', 'layers.1.self_attn.q_proj.weight', 'layers.1.self_attn.k_proj.weight', 'layers.1.self_attn.v_proj.weight', 'layers.1.self_attn.o_proj.weight', 'layers.2.self_attn.q_proj.weight', 'layers.2.self_attn.k_proj.weight', 'layers.2.self_attn.v_proj.weight', 'layers.2.self_attn.o_proj.weight', 'layers.3.self_attn.q_proj.weight', 'layers.3.self_attn.k_proj.weight', 'layers.3.self_attn.v_proj.weight', 'layers.3.self_attn.o_proj.weight', 'layers.4.self_attn.q_proj.weight', 'layers.4.self_attn.k_proj.weight', 'layers.4.self_attn.v_proj.weight', 'layers.4.self_attn.o_proj.weight', 'layers.5.self_attn.q_proj.weight', 'layers.5.self_attn.k_proj.weight', 'layers.5.self_attn.v_proj.weight', 'layers.5.self_attn.o_proj.weight', 'layers.6.self_attn.q_proj.weight', 'layers.6.self_attn.k_proj.weight', 'layers.6.self_attn.v_proj.weight', 'layers.6.self_attn.o_proj.weight', 'layers.7.self_attn.q_proj.weight', 'layers.7.self_attn.k_proj.weight', 'layers.7.self_attn.v_proj.weight', 'layers.7.self_attn.o_proj.weight', 'layers.8.self_attn.q_proj.weight', 'layers.8.self_attn.k_proj.weight', 'layers.8.self_attn.v_proj.weight', 'layers.8.self_attn.o_proj.weight', 'layers.9.self_attn.q_proj.weight', 'layers.9.self_attn.k_proj.weight', 'layers.9.self_attn.v_proj.weight', 'layers.9.self_attn.o_proj.weight', 'layers.10.self_attn.q_proj.weight', 'layers.10.self_attn.k_proj.weight', 'layers.10.self_attn.v_proj.weight', 'layers.10.self_attn.o_proj.weight', 'layers.11.self_attn.q_proj.weight', 'layers.11.self_attn.k_proj.weight', 'layers.11.self_attn.v_proj.weight', 'layers.11.self_attn.o_proj.weight', 'layers.12.self_attn.q_proj.weight', 'layers.12.self_attn.k_proj.weight', 'layers.12.self_attn.v_proj.weight', 'layers.12.self_attn.o_proj.weight', 'layers.13.self_attn.q_proj.weight', 'layers.13.self_attn.k_proj.weight', 'layers.13.self_attn.v_proj.weight', 'layers.13.self_attn.o_proj.weight', 'layers.14.self_attn.q_proj.weight', 'layers.14.self_attn.k_proj.weight', 'layers.14.self_attn.v_proj.weight', 'layers.14.self_attn.o_proj.weight', 'layers.15.self_attn.q_proj.weight', 'layers.15.self_attn.k_proj.weight', 'layers.15.self_attn.v_proj.weight', 'layers.15.self_attn.o_proj.weight'] warnings.warn(f"Removing redundant keys from checkpoint: {keys_to_delete}") WARNING:Neuron:casting embed_tokens.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.0.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.0.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to torch.float32 @@ -2520,11 +2490,11 @@ INFO:Neuron:Done Sharding weights in 0.5226453840732574 INFO:Neuron:Finished weights loading in 12.173774991882965 seconds INFO:Neuron:Warming up the model. INFO:Neuron:Warmup completed in 0.052486419677734375 seconds. -WARNING:__main__:AutoTokenizer failed: The checkpoint you are trying to load has model type `llama3_neuron` but Transformers does -not recognize this architecture. This could be because of an issue with the checkpoint, or because your version of Transformers is out of date. +WARNING:**main**:AutoTokenizer failed: The checkpoint you are trying to load has model type `llama3_neuron` but Transformers does +not recognize this architecture. This could be because of an issue with the checkpoint, or because your version of Transformers is out of date. You can update Transformers with the command `pip install --upgrade transformers`. If this does not work, and the checkpoint is ve -ry new, then there may not be a release version that supports this model yet. In this case, you can get the most up-to-date code by installing Transformers from source with the command `pip install git+https://github.com/huggingface/transformers.git` You are using the default legacy behaviour of the . This is e -xpected, and simply means that the `legacy` (previous) behavior will be used so nothing changes for you. If you want to use the new behaviour, set `legacy=False`. This should only be set if you understand what it means, and thoroughly read the reason why this was added as explained in https://github.com/huggingface/transformers/pull/24565 - if you loaded a llama tokenizer from a GGUF file you can ignore this message +ry new, then there may not be a release version that supports this model yet. In this case, you can get the most up-to-date code by installing Transformers from source with the command `pip install git+https://github.com/huggingface/transformers.git` You are using the default legacy behaviour of the . This is e +xpected, and simply means that the `legacy` (previous) behavior will be used so nothing changes for you. If you want to use the new behaviour, set `legacy=False`. This should only be set if you understand what it means, and thoroughly read the reason why this was added as explained in https://github.com/huggingface/transformers/pull/24565 - if you loaded a llama tokenizer from a GGUF file you can ignore this message ============================================================ LLAMA3 INFERENCE RESULTS ============================================================ @@ -2532,15 +2502,16 @@ Model: ./llama3_compiled Generation settings: max_tokens=5, temperature=0.6, top_p=0.9 ============================================================ -[1/1] Prompt: Hello, how are you? ----------------------------------------- -ERROR:__main__:Generation failed: 'bool' object has no attribute 'eos_token_id' +## [1/1] Prompt: Hello, how are you? + +ERROR:**main**:Generation failed: 'bool' object has no attribute 'eos_token_id' Traceback (most recent call last): - File "/home/ec2-user/NeuronxSDK/neuronx_llama3/run_inference.py", line 285, in generate_text - pad_token_id=tokenizer.eos_token_id, +File "/home/ec2-user/NeuronxSDK/neuronx_llama3/run_inference.py", line 285, in generate_text +pad_token_id=tokenizer.eos_token_id, AttributeError: 'bool' object has no attribute 'eos_token_id' Generated: Error during generation: 'bool' object has no attribute 'eos_token_id' ----------------------------------------- + +--- Inference completed! Great progress! The tokenizer is loading successfully, but there's an issue with the tokenizer object being a boolean instead of the actual tokenizer. Let me check what's happening in the tokenizer loading logic. @@ -2558,59 +2529,58 @@ Accepted edits to run_inference.py - Now let's test again to see what's happening with the tokenizer: Command python run_inference.py --model_path ./llama3_compiled --prompt "Hello, how are you?" --max_new_tokens 5 /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuro -nx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +nx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuro -nx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +nx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuro -nx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +nx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuron -x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuron -x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuron -x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuron -x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuron -x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuron -x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inference/modules/attention/utils.py:14: Deprecation -Warning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.custom_calls import neuron_cumsum +Warning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.custom_calls import neuron_cumsum /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: Dep -recationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttention_QKV +recationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: Dep -recationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttention_QKV +recationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: Dep -recationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttention_QKV +recationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is de -precated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base import NeuronAttentionBase +precated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is de -precated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base import NeuronAttentionBase +precated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is de -precated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base import NeuronAttentionBase +precated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base import NeuronAttentionBase WARNING:root:NeuronConfig init: Unexpected keyword arguments: {'apply_seq_ids_mask': False, 'enable_long_context_mode': False, 'en -able_output_completion_notifications': False, 'enable_token_tree': False, 'is_chunked_prefill': False, 'is_prefill_stage': None, 'kv_cache_tiling': False, 'scratchpad_page_size': None, 'skip_warmup': False, 'tile_cc': False, 'weights_to_skip_layout_optimization': []} INFO:Neuron:Sharding weights on load... +able_output_completion_notifications': False, 'enable_token_tree': False, 'is_chunked_prefill': False, 'is_prefill_stage': None, 'kv_cache_tiling': False, 'scratchpad_page_size': None, 'skip_warmup': False, 'tile_cc': False, 'weights_to_skip_layout_optimization': []} INFO:Neuron:Sharding weights on load... INFO:Neuron:Sharding Weights for ranks: 0...0 [2025-07-26 17:01:55.839: I neuronx_distributed/parallel_layers/parallel_state.py:628] > initializing tensor model parallel with s -ize 1 [2025-07-26 17:01:55.839: I neuronx_distributed/parallel_layers/parallel_state.py:629] > initializing pipeline model parallel with - size 1 [2025-07-26 17:01:55.839: I neuronx_distributed/parallel_layers/parallel_state.py:630] > initializing context model parallel with -size 1 [2025-07-26 17:01:55.839: I neuronx_distributed/parallel_layers/parallel_state.py:631] > initializing data parallel with size 1 +ize 1 [2025-07-26 17:01:55.839: I neuronx_distributed/parallel_layers/parallel_state.py:629] > initializing pipeline model parallel with +size 1 [2025-07-26 17:01:55.839: I neuronx_distributed/parallel_layers/parallel_state.py:630] > initializing context model parallel with +size 1 [2025-07-26 17:01:55.839: I neuronx_distributed/parallel_layers/parallel_state.py:631] > initializing data parallel with size 1 [2025-07-26 17:01:55.839: I neuronx_distributed/parallel_layers/parallel_state.py:632] > initializing world size to 1 [2025-07-26 17:01:55.840: I neuronx_distributed/parallel_layers/parallel_state.py:379] [rank_0_pp-1_tp-1_dp-1_cp-1] Chosen Logic f -or replica groups ret_logic=, 'Ascending Ring PG Group')> [2025-07-26 17:01:55.840: I neuronx_distributed/parallel_layers/parallel_state.py:668] [rank_0_pp-1_tp-1_dp-1_cp-1] tp_groups: rep -lica_groups.tp_groups=[[0]] [2025-07-26 17:01:55.840: I neuronx_distributed/parallel_layers/parallel_state.py:669] [rank_0_pp-1_tp-1_dp-1_cp-1] dp_groups: rep -lica_groups.dp_groups=[[0]] [2025-07-26 17:01:55.841: I neuronx_distributed/parallel_layers/parallel_state.py:670] [rank_0_pp-1_tp-1_dp-1_cp-1] pp_groups: rep -lica_groups.pp_groups=[[0]] [2025-07-26 17:01:55.841: I neuronx_distributed/parallel_layers/parallel_state.py:671] [rank_0_pp-1_tp-1_dp-1_cp-1] cp_groups: rep -lica_groups.cp_groups=[[0]] [2025-07-26 17:01:55.841: I neuronx_distributed/parallel_layers/parallel_state.py:672] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_model_group -s: replica_groups.ep_model_groups=[[0]] [2025-07-26 17:01:55.841: I neuronx_distributed/parallel_layers/parallel_state.py:673] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_data_groups -: replica_groups.ep_data_groups=[[0]] WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! +or replica groups ret_logic=, 'Ascending Ring PG Group')> [2025-07-26 17:01:55.840: I neuronx_distributed/parallel_layers/parallel_state.py:668] [rank_0_pp-1_tp-1_dp-1_cp-1] tp_groups: rep +lica_groups.tp_groups=[[0]] [2025-07-26 17:01:55.840: I neuronx_distributed/parallel_layers/parallel_state.py:669] [rank_0_pp-1_tp-1_dp-1_cp-1] dp_groups: rep +lica_groups.dp_groups=[[0]] [2025-07-26 17:01:55.841: I neuronx_distributed/parallel_layers/parallel_state.py:670] [rank_0_pp-1_tp-1_dp-1_cp-1] pp_groups: rep +lica_groups.pp_groups=[[0]] [2025-07-26 17:01:55.841: I neuronx_distributed/parallel_layers/parallel_state.py:671] [rank_0_pp-1_tp-1_dp-1_cp-1] cp_groups: rep +lica_groups.cp_groups=[[0]] [2025-07-26 17:01:55.841: I neuronx_distributed/parallel_layers/parallel_state.py:672] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_model_group +s: replica_groups.ep_model_groups=[[0]] [2025-07-26 17:01:55.841: I neuronx_distributed/parallel_layers/parallel_state.py:673] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_data_groups +: replica_groups.ep_data_groups=[[0]] WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! @@ -2644,7 +2614,7 @@ WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding atte WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/trace/trace.py:640: UserWarning: Removing redundant keys from - checkpoint: ['layers.0.self_attn.q_proj.weight', 'layers.0.self_attn.k_proj.weight', 'layers.0.self_attn.v_proj.weight', 'layers.0.self_attn.o_proj.weight', 'layers.1.self_attn.q_proj.weight', 'layers.1.self_attn.k_proj.weight', 'layers.1.self_attn.v_proj.weight', 'layers.1.self_attn.o_proj.weight', 'layers.2.self_attn.q_proj.weight', 'layers.2.self_attn.k_proj.weight', 'layers.2.self_attn.v_proj.weight', 'layers.2.self_attn.o_proj.weight', 'layers.3.self_attn.q_proj.weight', 'layers.3.self_attn.k_proj.weight', 'layers.3.self_attn.v_proj.weight', 'layers.3.self_attn.o_proj.weight', 'layers.4.self_attn.q_proj.weight', 'layers.4.self_attn.k_proj.weight', 'layers.4.self_attn.v_proj.weight', 'layers.4.self_attn.o_proj.weight', 'layers.5.self_attn.q_proj.weight', 'layers.5.self_attn.k_proj.weight', 'layers.5.self_attn.v_proj.weight', 'layers.5.self_attn.o_proj.weight', 'layers.6.self_attn.q_proj.weight', 'layers.6.self_attn.k_proj.weight', 'layers.6.self_attn.v_proj.weight', 'layers.6.self_attn.o_proj.weight', 'layers.7.self_attn.q_proj.weight', 'layers.7.self_attn.k_proj.weight', 'layers.7.self_attn.v_proj.weight', 'layers.7.self_attn.o_proj.weight', 'layers.8.self_attn.q_proj.weight', 'layers.8.self_attn.k_proj.weight', 'layers.8.self_attn.v_proj.weight', 'layers.8.self_attn.o_proj.weight', 'layers.9.self_attn.q_proj.weight', 'layers.9.self_attn.k_proj.weight', 'layers.9.self_attn.v_proj.weight', 'layers.9.self_attn.o_proj.weight', 'layers.10.self_attn.q_proj.weight', 'layers.10.self_attn.k_proj.weight', 'layers.10.self_attn.v_proj.weight', 'layers.10.self_attn.o_proj.weight', 'layers.11.self_attn.q_proj.weight', 'layers.11.self_attn.k_proj.weight', 'layers.11.self_attn.v_proj.weight', 'layers.11.self_attn.o_proj.weight', 'layers.12.self_attn.q_proj.weight', 'layers.12.self_attn.k_proj.weight', 'layers.12.self_attn.v_proj.weight', 'layers.12.self_attn.o_proj.weight', 'layers.13.self_attn.q_proj.weight', 'layers.13.self_attn.k_proj.weight', 'layers.13.self_attn.v_proj.weight', 'layers.13.self_attn.o_proj.weight', 'layers.14.self_attn.q_proj.weight', 'layers.14.self_attn.k_proj.weight', 'layers.14.self_attn.v_proj.weight', 'layers.14.self_attn.o_proj.weight', 'layers.15.self_attn.q_proj.weight', 'layers.15.self_attn.k_proj.weight', 'layers.15.self_attn.v_proj.weight', 'layers.15.self_attn.o_proj.weight'] warnings.warn(f"Removing redundant keys from checkpoint: {keys_to_delete}") +checkpoint: ['layers.0.self_attn.q_proj.weight', 'layers.0.self_attn.k_proj.weight', 'layers.0.self_attn.v_proj.weight', 'layers.0.self_attn.o_proj.weight', 'layers.1.self_attn.q_proj.weight', 'layers.1.self_attn.k_proj.weight', 'layers.1.self_attn.v_proj.weight', 'layers.1.self_attn.o_proj.weight', 'layers.2.self_attn.q_proj.weight', 'layers.2.self_attn.k_proj.weight', 'layers.2.self_attn.v_proj.weight', 'layers.2.self_attn.o_proj.weight', 'layers.3.self_attn.q_proj.weight', 'layers.3.self_attn.k_proj.weight', 'layers.3.self_attn.v_proj.weight', 'layers.3.self_attn.o_proj.weight', 'layers.4.self_attn.q_proj.weight', 'layers.4.self_attn.k_proj.weight', 'layers.4.self_attn.v_proj.weight', 'layers.4.self_attn.o_proj.weight', 'layers.5.self_attn.q_proj.weight', 'layers.5.self_attn.k_proj.weight', 'layers.5.self_attn.v_proj.weight', 'layers.5.self_attn.o_proj.weight', 'layers.6.self_attn.q_proj.weight', 'layers.6.self_attn.k_proj.weight', 'layers.6.self_attn.v_proj.weight', 'layers.6.self_attn.o_proj.weight', 'layers.7.self_attn.q_proj.weight', 'layers.7.self_attn.k_proj.weight', 'layers.7.self_attn.v_proj.weight', 'layers.7.self_attn.o_proj.weight', 'layers.8.self_attn.q_proj.weight', 'layers.8.self_attn.k_proj.weight', 'layers.8.self_attn.v_proj.weight', 'layers.8.self_attn.o_proj.weight', 'layers.9.self_attn.q_proj.weight', 'layers.9.self_attn.k_proj.weight', 'layers.9.self_attn.v_proj.weight', 'layers.9.self_attn.o_proj.weight', 'layers.10.self_attn.q_proj.weight', 'layers.10.self_attn.k_proj.weight', 'layers.10.self_attn.v_proj.weight', 'layers.10.self_attn.o_proj.weight', 'layers.11.self_attn.q_proj.weight', 'layers.11.self_attn.k_proj.weight', 'layers.11.self_attn.v_proj.weight', 'layers.11.self_attn.o_proj.weight', 'layers.12.self_attn.q_proj.weight', 'layers.12.self_attn.k_proj.weight', 'layers.12.self_attn.v_proj.weight', 'layers.12.self_attn.o_proj.weight', 'layers.13.self_attn.q_proj.weight', 'layers.13.self_attn.k_proj.weight', 'layers.13.self_attn.v_proj.weight', 'layers.13.self_attn.o_proj.weight', 'layers.14.self_attn.q_proj.weight', 'layers.14.self_attn.k_proj.weight', 'layers.14.self_attn.v_proj.weight', 'layers.14.self_attn.o_proj.weight', 'layers.15.self_attn.q_proj.weight', 'layers.15.self_attn.k_proj.weight', 'layers.15.self_attn.v_proj.weight', 'layers.15.self_attn.o_proj.weight'] warnings.warn(f"Removing redundant keys from checkpoint: {keys_to_delete}") WARNING:Neuron:casting embed_tokens.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.0.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.0.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to torch.float32 @@ -2796,11 +2766,11 @@ INFO:Neuron:Done Sharding weights in 0.5082141170278192 INFO:Neuron:Finished weights loading in 12.171694420976564 seconds INFO:Neuron:Warming up the model. INFO:Neuron:Warmup completed in 0.05229949951171875 seconds. -WARNING:__main__:AutoTokenizer failed: The checkpoint you are trying to load has model type `llama3_neuron` but Transformers does -not recognize this architecture. This could be because of an issue with the checkpoint, or because your version of Transformers is out of date. +WARNING:**main**:AutoTokenizer failed: The checkpoint you are trying to load has model type `llama3_neuron` but Transformers does +not recognize this architecture. This could be because of an issue with the checkpoint, or because your version of Transformers is out of date. You can update Transformers with the command `pip install --upgrade transformers`. If this does not work, and the checkpoint is ve -ry new, then there may not be a release version that supports this model yet. In this case, you can get the most up-to-date code by installing Transformers from source with the command `pip install git+https://github.com/huggingface/transformers.git` You are using the default legacy behaviour of the . This is e -xpected, and simply means that the `legacy` (previous) behavior will be used so nothing changes for you. If you want to use the new behaviour, set `legacy=False`. This should only be set if you understand what it means, and thoroughly read the reason why this was added as explained in https://github.com/huggingface/transformers/pull/24565 - if you loaded a llama tokenizer from a GGUF file you can ignore this message +ry new, then there may not be a release version that supports this model yet. In this case, you can get the most up-to-date code by installing Transformers from source with the command `pip install git+https://github.com/huggingface/transformers.git` You are using the default legacy behaviour of the . This is e +xpected, and simply means that the `legacy` (previous) behavior will be used so nothing changes for you. If you want to use the new behaviour, set `legacy=False`. This should only be set if you understand what it means, and thoroughly read the reason why this was added as explained in https://github.com/huggingface/transformers/pull/24565 - if you loaded a llama tokenizer from a GGUF file you can ignore this message ============================================================ LLAMA3 INFERENCE RESULTS ============================================================ @@ -2808,15 +2778,16 @@ Model: ./llama3_compiled Generation settings: max_tokens=5, temperature=0.6, top_p=0.9 ============================================================ -[1/1] Prompt: Hello, how are you? ----------------------------------------- -ERROR:__main__:Generation failed: 'bool' object has no attribute 'eos_token_id' +## [1/1] Prompt: Hello, how are you? + +ERROR:**main**:Generation failed: 'bool' object has no attribute 'eos_token_id' Traceback (most recent call last): - File "/home/ec2-user/NeuronxSDK/neuronx_llama3/run_inference.py", line 293, in generate_text - pad_token_id=tokenizer.eos_token_id, +File "/home/ec2-user/NeuronxSDK/neuronx_llama3/run_inference.py", line 293, in generate_text +pad_token_id=tokenizer.eos_token_id, AttributeError: 'bool' object has no attribute 'eos_token_id' Generated: Error during generation: 'bool' object has no attribute 'eos_token_id' ----------------------------------------- + +--- Inference completed! I can see the tokenizer is loading successfully (LlamaTokenizer message appears), but somehow the tokenizer variable is still being passed as a boolean. The issue is that the tokenizer loading debug messages aren't showing up, which means the function is returning before the tokenizer loading logic runs. Let me check if there's an issue with the function flow. @@ -2850,25 +2821,25 @@ from transformers import AutoTokenizer import os # Load tokenizer from original checkpoint and save to compiled directory + try: - tokenizer = AutoTokenizer.from_pretrained('./llama3_neuron_checkpoint') - tokenizer.save_pretrained('./llama3_compiled') - print('Tokenizer saved successfully to compiled directory') +tokenizer = AutoTokenizer.from_pretrained('./llama3_neuron_checkpoint') +tokenizer.save_pretrained('./llama3_compiled') +print('Tokenizer saved successfully to compiled directory') except Exception as e: - print(f'Failed to save tokenizer: {e}') - # Try from original Llama checkpoint - try: - tokenizer = AutoTokenizer.from_pretrained('/home/ec2-user/.llama/checkpoints/Llama3.2-1B', use_fast=False) - tokenizer.save_pretrained('./llama3_compiled') - print('Tokenizer saved from original Llama checkpoint') - except Exception as e2: - print(f'Also failed from original checkpoint: {e2}') +print(f'Failed to save tokenizer: {e}') # Try from original Llama checkpoint +try: +tokenizer = AutoTokenizer.from_pretrained('/home/ec2-user/.llama/checkpoints/Llama3.2-1B', use_fast=False) +tokenizer.save_pretrained('./llama3_compiled') +print('Tokenizer saved from original Llama checkpoint') +except Exception as e2: +print(f'Also failed from original checkpoint: {e2}') " Failed to save tokenizer: The checkpoint you are trying to load has model type `llama3_neuron` but Transformers does not recognize - this architecture. This could be because of an issue with the checkpoint, or because your version of Transformers is out of date. +this architecture. This could be because of an issue with the checkpoint, or because your version of Transformers is out of date. You can update Transformers with the command `pip install --upgrade transformers`. If this does not work, and the checkpoint is ve -ry new, then there may not be a release version that supports this model yet. In this case, you can get the most up-to-date code by installing Transformers from source with the command `pip install git+https://github.com/huggingface/transformers.git` You are using the default legacy behaviour of the . This is e -xpected, and simply means that the `legacy` (previous) behavior will be used so nothing changes for you. If you want to use the new behaviour, set `legacy=False`. This should only be set if you understand what it means, and thoroughly read the reason why this was added as explained in https://github.com/huggingface/transformers/pull/24565 - if you loaded a llama tokenizer from a GGUF file you can ignore this message Also failed from original checkpoint: 'bool' object has no attribute 'save_pretrained' +ry new, then there may not be a release version that supports this model yet. In this case, you can get the most up-to-date code by installing Transformers from source with the command `pip install git+https://github.com/huggingface/transformers.git` You are using the default legacy behaviour of the . This is e +xpected, and simply means that the `legacy` (previous) behavior will be used so nothing changes for you. If you want to use the new behaviour, set `legacy=False`. This should only be set if you understand what it means, and thoroughly read the reason why this was added as explained in https://github.com/huggingface/transformers/pull/24565 - if you loaded a llama tokenizer from a GGUF file you can ignore this message Also failed from original checkpoint: 'bool' object has no attribute 'save_pretrained' I see the issue. Let me try a different approach - let me load the tokenizer from the HuggingFace Hub and save it: Command @@ -2876,12 +2847,13 @@ python -c " from transformers import AutoTokenizer # Load tokenizer from HuggingFace Hub and save to compiled directory + try: - tokenizer = AutoTokenizer.from_pretrained('meta-llama/Llama-3.2-1B') - tokenizer.save_pretrained('./llama3_compiled') - print('Tokenizer saved successfully from HuggingFace Hub') +tokenizer = AutoTokenizer.from_pretrained('meta-llama/Llama-3.2-1B') +tokenizer.save_pretrained('./llama3_compiled') +print('Tokenizer saved successfully from HuggingFace Hub') except Exception as e: - print(f'Failed to save tokenizer from HuggingFace: {e}') +print(f'Failed to save tokenizer from HuggingFace: {e}') " Failed to save tokenizer from HuggingFace: You are trying to access a gated repo. Make sure to have access to it at https://huggingface.co/meta-llama/Llama-3.2-1B. @@ -2889,7 +2861,7 @@ Make sure to have access to it at https://huggingface.co/meta-llama/Llama-3.2-1B Cannot access gated repo for url https://huggingface.co/meta-llama/Llama-3.2-1B/resolve/main/config.json. Access to model meta-llama/Llama-3.2-1B is restricted. You must have access to it and be authenticated to access it. Please log in -. +. Let me try a simpler approach - let me just simplify our tokenizer loading logic to match the working example pattern, but handle the case where we don't have the proper tokenizer files: Error(s) while editing @@ -2925,45 +2897,45 @@ The imports are there, but they're inside the function. The error 'bool' object Command python -c " try: - from neuronx_distributed_inference.utils.hf_adapter import HuggingFaceGenerationAdapter - print('✅ HuggingFaceGenerationAdapter imported successfully') - print(f'Type: {type(HuggingFaceGenerationAdapter)}') +from neuronx_distributed_inference.utils.hf_adapter import HuggingFaceGenerationAdapter +print('✅ HuggingFaceGenerationAdapter imported successfully') +print(f'Type: {type(HuggingFaceGenerationAdapter)}') except Exception as e: - print(f'❌ Failed to import HuggingFaceGenerationAdapter: {e}') +print(f'❌ Failed to import HuggingFaceGenerationAdapter: {e}') try: - from neuronx_distributed_inference.modules.generation.sampling import prepare_sampling_params - print('✅ prepare_sampling_params imported successfully') - print(f'Type: {type(prepare_sampling_params)}') +from neuronx_distributed_inference.modules.generation.sampling import prepare_sampling_params +print('✅ prepare_sampling_params imported successfully') +print(f'Type: {type(prepare_sampling_params)}') except Exception as e: - print(f'❌ Failed to import prepare_sampling_params: {e}') +print(f'❌ Failed to import prepare_sampling_params: {e}') " /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuro -nx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +nx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuro -nx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +nx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuro -nx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +nx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuron -x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuron -x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuron -x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuron -x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuron -x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuron -x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inference/modules/attention/utils.py:14: Deprecation -Warning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.custom_calls import neuron_cumsum +Warning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.custom_calls import neuron_cumsum /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: Dep -recationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttention_QKV +recationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: Dep -recationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttention_QKV +recationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: Dep -recationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttention_QKV +recationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttention_QKV ✅ HuggingFaceGenerationAdapter imported successfully Type: ✅ prepare_sampling_params imported successfully @@ -2986,59 +2958,58 @@ Created simple_inference.py - Now let's test this simplified version: Command python simple_inference.py --model_path ./llama3_compiled --prompt "Hello, how are you?" --max_new_tokens 5 /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuro -nx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +nx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuro -nx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +nx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuro -nx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +nx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuron -x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuron -x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuron -x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuron -x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuron -x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuron -x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inference/modules/attention/utils.py:14: Deprecation -Warning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.custom_calls import neuron_cumsum +Warning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.custom_calls import neuron_cumsum /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: Dep -recationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttention_QKV +recationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: Dep -recationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttention_QKV +recationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: Dep -recationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttention_QKV +recationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is de -precated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base import NeuronAttentionBase +precated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is de -precated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base import NeuronAttentionBase +precated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is de -precated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base import NeuronAttentionBase +precated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base import NeuronAttentionBase WARNING:root:NeuronConfig init: Unexpected keyword arguments: {'apply_seq_ids_mask': False, 'enable_long_context_mode': False, 'en -able_output_completion_notifications': False, 'enable_token_tree': False, 'is_chunked_prefill': False, 'is_prefill_stage': None, 'kv_cache_tiling': False, 'scratchpad_page_size': None, 'skip_warmup': False, 'tile_cc': False, 'weights_to_skip_layout_optimization': []} INFO:Neuron:Sharding weights on load... +able_output_completion_notifications': False, 'enable_token_tree': False, 'is_chunked_prefill': False, 'is_prefill_stage': None, 'kv_cache_tiling': False, 'scratchpad_page_size': None, 'skip_warmup': False, 'tile_cc': False, 'weights_to_skip_layout_optimization': []} INFO:Neuron:Sharding weights on load... INFO:Neuron:Sharding Weights for ranks: 0...0 [2025-07-26 17:05:48.481: I neuronx_distributed/parallel_layers/parallel_state.py:628] > initializing tensor model parallel with s -ize 1 [2025-07-26 17:05:48.481: I neuronx_distributed/parallel_layers/parallel_state.py:629] > initializing pipeline model parallel with - size 1 [2025-07-26 17:05:48.482: I neuronx_distributed/parallel_layers/parallel_state.py:630] > initializing context model parallel with -size 1 [2025-07-26 17:05:48.482: I neuronx_distributed/parallel_layers/parallel_state.py:631] > initializing data parallel with size 1 +ize 1 [2025-07-26 17:05:48.481: I neuronx_distributed/parallel_layers/parallel_state.py:629] > initializing pipeline model parallel with +size 1 [2025-07-26 17:05:48.482: I neuronx_distributed/parallel_layers/parallel_state.py:630] > initializing context model parallel with +size 1 [2025-07-26 17:05:48.482: I neuronx_distributed/parallel_layers/parallel_state.py:631] > initializing data parallel with size 1 [2025-07-26 17:05:48.482: I neuronx_distributed/parallel_layers/parallel_state.py:632] > initializing world size to 1 [2025-07-26 17:05:48.483: I neuronx_distributed/parallel_layers/parallel_state.py:379] [rank_0_pp-1_tp-1_dp-1_cp-1] Chosen Logic f -or replica groups ret_logic=, 'Ascending Ring PG Group')> [2025-07-26 17:05:48.483: I neuronx_distributed/parallel_layers/parallel_state.py:668] [rank_0_pp-1_tp-1_dp-1_cp-1] tp_groups: rep -lica_groups.tp_groups=[[0]] [2025-07-26 17:05:48.483: I neuronx_distributed/parallel_layers/parallel_state.py:669] [rank_0_pp-1_tp-1_dp-1_cp-1] dp_groups: rep -lica_groups.dp_groups=[[0]] [2025-07-26 17:05:48.483: I neuronx_distributed/parallel_layers/parallel_state.py:670] [rank_0_pp-1_tp-1_dp-1_cp-1] pp_groups: rep -lica_groups.pp_groups=[[0]] [2025-07-26 17:05:48.483: I neuronx_distributed/parallel_layers/parallel_state.py:671] [rank_0_pp-1_tp-1_dp-1_cp-1] cp_groups: rep -lica_groups.cp_groups=[[0]] [2025-07-26 17:05:48.483: I neuronx_distributed/parallel_layers/parallel_state.py:672] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_model_group -s: replica_groups.ep_model_groups=[[0]] [2025-07-26 17:05:48.483: I neuronx_distributed/parallel_layers/parallel_state.py:673] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_data_groups -: replica_groups.ep_data_groups=[[0]] WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! +or replica groups ret_logic=, 'Ascending Ring PG Group')> [2025-07-26 17:05:48.483: I neuronx_distributed/parallel_layers/parallel_state.py:668] [rank_0_pp-1_tp-1_dp-1_cp-1] tp_groups: rep +lica_groups.tp_groups=[[0]] [2025-07-26 17:05:48.483: I neuronx_distributed/parallel_layers/parallel_state.py:669] [rank_0_pp-1_tp-1_dp-1_cp-1] dp_groups: rep +lica_groups.dp_groups=[[0]] [2025-07-26 17:05:48.483: I neuronx_distributed/parallel_layers/parallel_state.py:670] [rank_0_pp-1_tp-1_dp-1_cp-1] pp_groups: rep +lica_groups.pp_groups=[[0]] [2025-07-26 17:05:48.483: I neuronx_distributed/parallel_layers/parallel_state.py:671] [rank_0_pp-1_tp-1_dp-1_cp-1] cp_groups: rep +lica_groups.cp_groups=[[0]] [2025-07-26 17:05:48.483: I neuronx_distributed/parallel_layers/parallel_state.py:672] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_model_group +s: replica_groups.ep_model_groups=[[0]] [2025-07-26 17:05:48.483: I neuronx_distributed/parallel_layers/parallel_state.py:673] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_data_groups +: replica_groups.ep_data_groups=[[0]] WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! @@ -3072,7 +3043,7 @@ WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding atte WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/trace/trace.py:640: UserWarning: Removing redundant keys from - checkpoint: ['layers.0.self_attn.q_proj.weight', 'layers.0.self_attn.k_proj.weight', 'layers.0.self_attn.v_proj.weight', 'layers.0.self_attn.o_proj.weight', 'layers.1.self_attn.q_proj.weight', 'layers.1.self_attn.k_proj.weight', 'layers.1.self_attn.v_proj.weight', 'layers.1.self_attn.o_proj.weight', 'layers.2.self_attn.q_proj.weight', 'layers.2.self_attn.k_proj.weight', 'layers.2.self_attn.v_proj.weight', 'layers.2.self_attn.o_proj.weight', 'layers.3.self_attn.q_proj.weight', 'layers.3.self_attn.k_proj.weight', 'layers.3.self_attn.v_proj.weight', 'layers.3.self_attn.o_proj.weight', 'layers.4.self_attn.q_proj.weight', 'layers.4.self_attn.k_proj.weight', 'layers.4.self_attn.v_proj.weight', 'layers.4.self_attn.o_proj.weight', 'layers.5.self_attn.q_proj.weight', 'layers.5.self_attn.k_proj.weight', 'layers.5.self_attn.v_proj.weight', 'layers.5.self_attn.o_proj.weight', 'layers.6.self_attn.q_proj.weight', 'layers.6.self_attn.k_proj.weight', 'layers.6.self_attn.v_proj.weight', 'layers.6.self_attn.o_proj.weight', 'layers.7.self_attn.q_proj.weight', 'layers.7.self_attn.k_proj.weight', 'layers.7.self_attn.v_proj.weight', 'layers.7.self_attn.o_proj.weight', 'layers.8.self_attn.q_proj.weight', 'layers.8.self_attn.k_proj.weight', 'layers.8.self_attn.v_proj.weight', 'layers.8.self_attn.o_proj.weight', 'layers.9.self_attn.q_proj.weight', 'layers.9.self_attn.k_proj.weight', 'layers.9.self_attn.v_proj.weight', 'layers.9.self_attn.o_proj.weight', 'layers.10.self_attn.q_proj.weight', 'layers.10.self_attn.k_proj.weight', 'layers.10.self_attn.v_proj.weight', 'layers.10.self_attn.o_proj.weight', 'layers.11.self_attn.q_proj.weight', 'layers.11.self_attn.k_proj.weight', 'layers.11.self_attn.v_proj.weight', 'layers.11.self_attn.o_proj.weight', 'layers.12.self_attn.q_proj.weight', 'layers.12.self_attn.k_proj.weight', 'layers.12.self_attn.v_proj.weight', 'layers.12.self_attn.o_proj.weight', 'layers.13.self_attn.q_proj.weight', 'layers.13.self_attn.k_proj.weight', 'layers.13.self_attn.v_proj.weight', 'layers.13.self_attn.o_proj.weight', 'layers.14.self_attn.q_proj.weight', 'layers.14.self_attn.k_proj.weight', 'layers.14.self_attn.v_proj.weight', 'layers.14.self_attn.o_proj.weight', 'layers.15.self_attn.q_proj.weight', 'layers.15.self_attn.k_proj.weight', 'layers.15.self_attn.v_proj.weight', 'layers.15.self_attn.o_proj.weight'] warnings.warn(f"Removing redundant keys from checkpoint: {keys_to_delete}") +checkpoint: ['layers.0.self_attn.q_proj.weight', 'layers.0.self_attn.k_proj.weight', 'layers.0.self_attn.v_proj.weight', 'layers.0.self_attn.o_proj.weight', 'layers.1.self_attn.q_proj.weight', 'layers.1.self_attn.k_proj.weight', 'layers.1.self_attn.v_proj.weight', 'layers.1.self_attn.o_proj.weight', 'layers.2.self_attn.q_proj.weight', 'layers.2.self_attn.k_proj.weight', 'layers.2.self_attn.v_proj.weight', 'layers.2.self_attn.o_proj.weight', 'layers.3.self_attn.q_proj.weight', 'layers.3.self_attn.k_proj.weight', 'layers.3.self_attn.v_proj.weight', 'layers.3.self_attn.o_proj.weight', 'layers.4.self_attn.q_proj.weight', 'layers.4.self_attn.k_proj.weight', 'layers.4.self_attn.v_proj.weight', 'layers.4.self_attn.o_proj.weight', 'layers.5.self_attn.q_proj.weight', 'layers.5.self_attn.k_proj.weight', 'layers.5.self_attn.v_proj.weight', 'layers.5.self_attn.o_proj.weight', 'layers.6.self_attn.q_proj.weight', 'layers.6.self_attn.k_proj.weight', 'layers.6.self_attn.v_proj.weight', 'layers.6.self_attn.o_proj.weight', 'layers.7.self_attn.q_proj.weight', 'layers.7.self_attn.k_proj.weight', 'layers.7.self_attn.v_proj.weight', 'layers.7.self_attn.o_proj.weight', 'layers.8.self_attn.q_proj.weight', 'layers.8.self_attn.k_proj.weight', 'layers.8.self_attn.v_proj.weight', 'layers.8.self_attn.o_proj.weight', 'layers.9.self_attn.q_proj.weight', 'layers.9.self_attn.k_proj.weight', 'layers.9.self_attn.v_proj.weight', 'layers.9.self_attn.o_proj.weight', 'layers.10.self_attn.q_proj.weight', 'layers.10.self_attn.k_proj.weight', 'layers.10.self_attn.v_proj.weight', 'layers.10.self_attn.o_proj.weight', 'layers.11.self_attn.q_proj.weight', 'layers.11.self_attn.k_proj.weight', 'layers.11.self_attn.v_proj.weight', 'layers.11.self_attn.o_proj.weight', 'layers.12.self_attn.q_proj.weight', 'layers.12.self_attn.k_proj.weight', 'layers.12.self_attn.v_proj.weight', 'layers.12.self_attn.o_proj.weight', 'layers.13.self_attn.q_proj.weight', 'layers.13.self_attn.k_proj.weight', 'layers.13.self_attn.v_proj.weight', 'layers.13.self_attn.o_proj.weight', 'layers.14.self_attn.q_proj.weight', 'layers.14.self_attn.k_proj.weight', 'layers.14.self_attn.v_proj.weight', 'layers.14.self_attn.o_proj.weight', 'layers.15.self_attn.q_proj.weight', 'layers.15.self_attn.k_proj.weight', 'layers.15.self_attn.v_proj.weight', 'layers.15.self_attn.o_proj.weight'] warnings.warn(f"Removing redundant keys from checkpoint: {keys_to_delete}") WARNING:Neuron:casting embed_tokens.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.0.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.0.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to torch.float32 @@ -3225,27 +3196,27 @@ INFO:Neuron:Finished weights loading in 19.60225479491055 seconds INFO:Neuron:Warming up the model. INFO:Neuron:Warmup completed in 0.05246686935424805 seconds. Traceback (most recent call last): - File "/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages/transformers/models/auto/configuration_auto.p -y", line 1131, in from_pretrained config_class = CONFIG_MAPPING[config_dict["model_type"]] - File "/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages/transformers/models/auto/configuration_auto.p -y", line 833, in __getitem__ raise KeyError(key) +File "/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages/transformers/models/auto/configuration_auto.p +y", line 1131, in from_pretrained config_class = CONFIG_MAPPING[config_dict["model_type"]] +File "/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages/transformers/models/auto/configuration_auto.p +y", line 833, in **getitem** raise KeyError(key) KeyError: 'llama3_neuron' During handling of the above exception, another exception occurred: Traceback (most recent call last): - File "/home/ec2-user/NeuronxSDK/neuronx_llama3/simple_inference.py", line 87, in - main() - File "/home/ec2-user/NeuronxSDK/neuronx_llama3/simple_inference.py", line 38, in main - tokenizer = AutoTokenizer.from_pretrained(args.model_path) - File "/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages/transformers/models/auto/tokenization_auto.py -", line 966, in from_pretrained config = AutoConfig.from_pretrained( - File "/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages/transformers/models/auto/configuration_auto.p -y", line 1133, in from_pretrained raise ValueError( +File "/home/ec2-user/NeuronxSDK/neuronx_llama3/simple_inference.py", line 87, in +main() +File "/home/ec2-user/NeuronxSDK/neuronx_llama3/simple_inference.py", line 38, in main +tokenizer = AutoTokenizer.from_pretrained(args.model_path) +File "/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages/transformers/models/auto/tokenization_auto.py +", line 966, in from_pretrained config = AutoConfig.from_pretrained( +File "/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages/transformers/models/auto/configuration_auto.p +y", line 1133, in from_pretrained raise ValueError( ValueError: The checkpoint you are trying to load has model type `llama3_neuron` but Transformers does not recognize this architec -ture. This could be because of an issue with the checkpoint, or because your version of Transformers is out of date. +ture. This could be because of an issue with the checkpoint, or because your version of Transformers is out of date. You can update Transformers with the command `pip install --upgrade transformers`. If this does not work, and the checkpoint is ve -ry new, then there may not be a release version that supports this model yet. In this case, you can get the most up-to-date code by installing Transformers from source with the command `pip install git+https://github.com/huggingface/transformers.git` +ry new, then there may not be a release version that supports this model yet. In this case, you can get the most up-to-date code by installing Transformers from source with the command `pip install git+https://github.com/huggingface/transformers.git` Perfect! Now I can see the exact issue. The problem is that the compiled directory has a config.json with model_type: "llama3_neuron" which transformers doesn't recognize. This is exactly what you meant about overcomplicating the tokenizer loading. The solution is simple: we need to copy the proper tokenizer files from the original checkpoint to the compiled directory, just like the working example does during compilation. Let me fix this: @@ -3257,33 +3228,36 @@ import shutil import os # Copy tokenizer files from original checkpoint to compiled directory + original_checkpoint = './llama3_neuron_checkpoint' compiled_dir = './llama3_compiled' # Files to copy for tokenizer + tokenizer_files = [ - 'tokenizer.json', - 'tokenizer_config.json', - 'special_tokens_map.json', - 'tokenizer.model' +'tokenizer.json', +'tokenizer_config.json', +'special_tokens_map.json', +'tokenizer.model' ] print('Copying tokenizer files...') for file in tokenizer_files: - src = os.path.join(original_checkpoint, file) - dst = os.path.join(compiled_dir, file) - if os.path.exists(src): - shutil.copy2(src, dst) - print(f'✅ Copied {file}') - else: - print(f'❌ {file} not found in original checkpoint') +src = os.path.join(original_checkpoint, file) +dst = os.path.join(compiled_dir, file) +if os.path.exists(src): +shutil.copy2(src, dst) +print(f'✅ Copied {file}') +else: +print(f'❌ {file} not found in original checkpoint') # Also need to fix the config.json to have the right model_type + config_path = os.path.join(compiled_dir, 'config.json') if os.path.exists(config_path): - with open(config_path, 'r') as f: - config = json.load(f) - +with open(config_path, 'r') as f: +config = json.load(f) + # Change model_type back to llama for tokenizer compatibility if config.get('model_type') == 'llama3_neuron': config['model_type'] = 'llama' @@ -3292,8 +3266,9 @@ if os.path.exists(config_path): print('✅ Fixed config.json model_type') else: print(f'Config model_type is: {config.get(\"model_type\")}') + else: - print('❌ config.json not found') +print('❌ config.json not found') " Copying tokenizer files... ❌ tokenizer.json not found in original checkpoint @@ -3308,31 +3283,33 @@ python -c " import json # Create minimal tokenizer_config.json + tokenizer_config = { - 'tokenizer_class': 'LlamaTokenizer', - 'model_max_length': 2048, - 'padding_side': 'right', - 'special_tokens_map_file': None, - 'name_or_path': './llama3_compiled', - 'tokenizer_file': None, - 'legacy': True +'tokenizer_class': 'LlamaTokenizer', +'model_max_length': 2048, +'padding_side': 'right', +'special_tokens_map_file': None, +'name_or_path': './llama3_compiled', +'tokenizer_file': None, +'legacy': True } with open('./llama3_compiled/tokenizer_config.json', 'w') as f: - json.dump(tokenizer_config, f, indent=2) +json.dump(tokenizer_config, f, indent=2) print('✅ Created tokenizer_config.json') # Create minimal special_tokens_map.json + special_tokens = { - 'bos_token': '<|begin_of_text|>', - 'eos_token': '<|end_of_text|>', - 'pad_token': '<|end_of_text|>', - 'unk_token': '' +'bos_token': '<|begin_of_text|>', +'eos_token': '<|end_of_text|>', +'pad_token': '<|end_of_text|>', +'unk_token': '' } with open('./llama3_compiled/special_tokens_map.json', 'w') as f: - json.dump(special_tokens, f, indent=2) +json.dump(special_tokens, f, indent=2) print('✅ Created special_tokens_map.json') " @@ -3343,53 +3320,53 @@ Now let's test the simple inference script again: Command python simple_inference.py --model_path ./llama3_compiled --prompt "Hello, how are you?" --max_new_tokens 5 /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuro -nx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +nx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuro -nx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +nx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/parallel_layers/layers.py:14: DeprecationWarning: torch_neuro -nx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +nx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuron -x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuron -x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuron -x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuron -x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuron -x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: DeprecationWarning: torch_neuron -x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +x.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inference/modules/attention/utils.py:14: Deprecation -Warning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.custom_calls import neuron_cumsum +Warning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.custom_calls import neuron_cumsum /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: Dep -recationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttention_QKV +recationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: Dep -recationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttention_QKV +recationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inference/modules/lora_serving/lora_model.py:12: Dep -recationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttention_QKV +recationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is de -precated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base import NeuronAttentionBase +precated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is de -precated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base import NeuronAttentionBase +precated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_llama3/modeling_llama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is de -precated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base import NeuronAttentionBase +precated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base import NeuronAttentionBase WARNING:root:NeuronConfig init: Unexpected keyword arguments: {'apply_seq_ids_mask': False, 'enable_long_context_mode': False, 'en -able_output_completion_notifications': False, 'enable_token_tree': False, 'is_chunked_prefill': False, 'is_prefill_stage': None, 'kv_cache_tiling': False, 'scratchpad_page_size': None, 'skip_warmup': False, 'tile_cc': False, 'weights_to_skip_layout_optimization': []} INFO:Neuron:Sharding weights on load... +able_output_completion_notifications': False, 'enable_token_tree': False, 'is_chunked_prefill': False, 'is_prefill_stage': None, 'kv_cache_tiling': False, 'scratchpad_page_size': None, 'skip_warmup': False, 'tile_cc': False, 'weights_to_skip_layout_optimization': []} INFO:Neuron:Sharding weights on load... INFO:Neuron:Sharding Weights for ranks: 0...0 [2025-07-26 17:06:45.745: I neuronx_distributed/parallel_layers/parallel_state.py:628] > initializing tensor model parallel with s -ize 1 [2025-07-26 17:06:45.745: I neuronx_distributed/parallel_layers/parallel_state.py:629] > initializing pipeline model parallel with - size 1 [2025-07-26 17:06:45.745: I neuronx_distributed/parallel_layers/parallel_state.py:630] > initializing context model parallel with -size 1 [2025-07-26 17:06:45.745: I neuronx_distributed/parallel_layers/parallel_state.py:631] > initializing data parallel with size 1 +ize 1 [2025-07-26 17:06:45.745: I neuronx_distributed/parallel_layers/parallel_state.py:629] > initializing pipeline model parallel with +size 1 [2025-07-26 17:06:45.745: I neuronx_distributed/parallel_layers/parallel_state.py:630] > initializing context model parallel with +size 1 [2025-07-26 17:06:45.745: I neuronx_distributed/parallel_layers/parallel_state.py:631] > initializing data parallel with size 1 [2025-07-26 17:06:45.745: I neuronx_distributed/parallel_layers/parallel_state.py:632] > initializing world size to 1 [2025-07-26 17:06:45.746: I neuronx_distributed/parallel_layers/parallel_state.py:379] [rank_0_pp-1_tp-1_dp-1_cp-1] Chosen Logic f -or replica groups ret_logic=, 'Ascending Ring PG Group')> [2025-07-26 17:06:45.747: I neuronx_distributed/parallel_layers/parallel_state.py:668] [rank_0_pp-1_tp-1_dp-1_cp-1] tp_groups: rep -lica_groups.tp_groups=[[0]] [2025-07-26 17:06:45.747: I neuronx_distributed/parallel_layers/parallel_state.py:669] [rank_0_pp-1_tp-1_dp-1_cp-1] dp_groups: rep -lica_groups.dp_groups=[[0]] [2025-07-26 17:06:45.747: I neuronx_distributed/parallel_layers/parallel_state.py:670] [rank_0_pp-1_tp-1_dp-1_cp-1] pp_groups: rep -lica_groups.pp_groups=[[0]] [2025-07-26 17:06:45.747: I neuronx_distributed/parallel_layers/parallel_state.py:671] [rank_0_pp-1_tp-1_dp-1_cp-1] cp_groups: rep -lica_groups.cp_groups=[[0]] [2025-07-26 17:06:45.747: I neuronx_distributed/parallel_layers/parallel_state.py:672] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_model_group -s: replica_groups.ep_model_groups=[[0]] [2025-07-26 17:06:45.747: I neuronx_distributed/parallel_layers/parallel_state.py:673] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_data_groups -: replica_groups.ep_data_groups=[[0]] WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! +or replica groups ret_logic=, 'Ascending Ring PG Group')> [2025-07-26 17:06:45.747: I neuronx_distributed/parallel_layers/parallel_state.py:668] [rank_0_pp-1_tp-1_dp-1_cp-1] tp_groups: rep +lica_groups.tp_groups=[[0]] [2025-07-26 17:06:45.747: I neuronx_distributed/parallel_layers/parallel_state.py:669] [rank_0_pp-1_tp-1_dp-1_cp-1] dp_groups: rep +lica_groups.dp_groups=[[0]] [2025-07-26 17:06:45.747: I neuronx_distributed/parallel_layers/parallel_state.py:670] [rank_0_pp-1_tp-1_dp-1_cp-1] pp_groups: rep +lica_groups.pp_groups=[[0]] [2025-07-26 17:06:45.747: I neuronx_distributed/parallel_layers/parallel_state.py:671] [rank_0_pp-1_tp-1_dp-1_cp-1] cp_groups: rep +lica_groups.cp_groups=[[0]] [2025-07-26 17:06:45.747: I neuronx_distributed/parallel_layers/parallel_state.py:672] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_model_group +s: replica_groups.ep_model_groups=[[0]] [2025-07-26 17:06:45.747: I neuronx_distributed/parallel_layers/parallel_state.py:673] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_data_groups +: replica_groups.ep_data_groups=[[0]] WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! @@ -3423,7 +3400,7 @@ WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding atte WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA! /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/trace/trace.py:640: UserWarning: Removing redundant keys from - checkpoint: ['layers.0.self_attn.q_proj.weight', 'layers.0.self_attn.k_proj.weight', 'layers.0.self_attn.v_proj.weight', 'layers.0.self_attn.o_proj.weight', 'layers.1.self_attn.q_proj.weight', 'layers.1.self_attn.k_proj.weight', 'layers.1.self_attn.v_proj.weight', 'layers.1.self_attn.o_proj.weight', 'layers.2.self_attn.q_proj.weight', 'layers.2.self_attn.k_proj.weight', 'layers.2.self_attn.v_proj.weight', 'layers.2.self_attn.o_proj.weight', 'layers.3.self_attn.q_proj.weight', 'layers.3.self_attn.k_proj.weight', 'layers.3.self_attn.v_proj.weight', 'layers.3.self_attn.o_proj.weight', 'layers.4.self_attn.q_proj.weight', 'layers.4.self_attn.k_proj.weight', 'layers.4.self_attn.v_proj.weight', 'layers.4.self_attn.o_proj.weight', 'layers.5.self_attn.q_proj.weight', 'layers.5.self_attn.k_proj.weight', 'layers.5.self_attn.v_proj.weight', 'layers.5.self_attn.o_proj.weight', 'layers.6.self_attn.q_proj.weight', 'layers.6.self_attn.k_proj.weight', 'layers.6.self_attn.v_proj.weight', 'layers.6.self_attn.o_proj.weight', 'layers.7.self_attn.q_proj.weight', 'layers.7.self_attn.k_proj.weight', 'layers.7.self_attn.v_proj.weight', 'layers.7.self_attn.o_proj.weight', 'layers.8.self_attn.q_proj.weight', 'layers.8.self_attn.k_proj.weight', 'layers.8.self_attn.v_proj.weight', 'layers.8.self_attn.o_proj.weight', 'layers.9.self_attn.q_proj.weight', 'layers.9.self_attn.k_proj.weight', 'layers.9.self_attn.v_proj.weight', 'layers.9.self_attn.o_proj.weight', 'layers.10.self_attn.q_proj.weight', 'layers.10.self_attn.k_proj.weight', 'layers.10.self_attn.v_proj.weight', 'layers.10.self_attn.o_proj.weight', 'layers.11.self_attn.q_proj.weight', 'layers.11.self_attn.k_proj.weight', 'layers.11.self_attn.v_proj.weight', 'layers.11.self_attn.o_proj.weight', 'layers.12.self_attn.q_proj.weight', 'layers.12.self_attn.k_proj.weight', 'layers.12.self_attn.v_proj.weight', 'layers.12.self_attn.o_proj.weight', 'layers.13.self_attn.q_proj.weight', 'layers.13.self_attn.k_proj.weight', 'layers.13.self_attn.v_proj.weight', 'layers.13.self_attn.o_proj.weight', 'layers.14.self_attn.q_proj.weight', 'layers.14.self_attn.k_proj.weight', 'layers.14.self_attn.v_proj.weight', 'layers.14.self_attn.o_proj.weight', 'layers.15.self_attn.q_proj.weight', 'layers.15.self_attn.k_proj.weight', 'layers.15.self_attn.v_proj.weight', 'layers.15.self_attn.o_proj.weight'] warnings.warn(f"Removing redundant keys from checkpoint: {keys_to_delete}") +checkpoint: ['layers.0.self_attn.q_proj.weight', 'layers.0.self_attn.k_proj.weight', 'layers.0.self_attn.v_proj.weight', 'layers.0.self_attn.o_proj.weight', 'layers.1.self_attn.q_proj.weight', 'layers.1.self_attn.k_proj.weight', 'layers.1.self_attn.v_proj.weight', 'layers.1.self_attn.o_proj.weight', 'layers.2.self_attn.q_proj.weight', 'layers.2.self_attn.k_proj.weight', 'layers.2.self_attn.v_proj.weight', 'layers.2.self_attn.o_proj.weight', 'layers.3.self_attn.q_proj.weight', 'layers.3.self_attn.k_proj.weight', 'layers.3.self_attn.v_proj.weight', 'layers.3.self_attn.o_proj.weight', 'layers.4.self_attn.q_proj.weight', 'layers.4.self_attn.k_proj.weight', 'layers.4.self_attn.v_proj.weight', 'layers.4.self_attn.o_proj.weight', 'layers.5.self_attn.q_proj.weight', 'layers.5.self_attn.k_proj.weight', 'layers.5.self_attn.v_proj.weight', 'layers.5.self_attn.o_proj.weight', 'layers.6.self_attn.q_proj.weight', 'layers.6.self_attn.k_proj.weight', 'layers.6.self_attn.v_proj.weight', 'layers.6.self_attn.o_proj.weight', 'layers.7.self_attn.q_proj.weight', 'layers.7.self_attn.k_proj.weight', 'layers.7.self_attn.v_proj.weight', 'layers.7.self_attn.o_proj.weight', 'layers.8.self_attn.q_proj.weight', 'layers.8.self_attn.k_proj.weight', 'layers.8.self_attn.v_proj.weight', 'layers.8.self_attn.o_proj.weight', 'layers.9.self_attn.q_proj.weight', 'layers.9.self_attn.k_proj.weight', 'layers.9.self_attn.v_proj.weight', 'layers.9.self_attn.o_proj.weight', 'layers.10.self_attn.q_proj.weight', 'layers.10.self_attn.k_proj.weight', 'layers.10.self_attn.v_proj.weight', 'layers.10.self_attn.o_proj.weight', 'layers.11.self_attn.q_proj.weight', 'layers.11.self_attn.k_proj.weight', 'layers.11.self_attn.v_proj.weight', 'layers.11.self_attn.o_proj.weight', 'layers.12.self_attn.q_proj.weight', 'layers.12.self_attn.k_proj.weight', 'layers.12.self_attn.v_proj.weight', 'layers.12.self_attn.o_proj.weight', 'layers.13.self_attn.q_proj.weight', 'layers.13.self_attn.k_proj.weight', 'layers.13.self_attn.v_proj.weight', 'layers.13.self_attn.o_proj.weight', 'layers.14.self_attn.q_proj.weight', 'layers.14.self_attn.k_proj.weight', 'layers.14.self_attn.v_proj.weight', 'layers.14.self_attn.o_proj.weight', 'layers.15.self_attn.q_proj.weight', 'layers.15.self_attn.k_proj.weight', 'layers.15.self_attn.v_proj.weight', 'layers.15.self_attn.o_proj.weight'] warnings.warn(f"Removing redundant keys from checkpoint: {keys_to_delete}") WARNING:Neuron:casting embed_tokens.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.0.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.0.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to torch.float32 @@ -3576,39 +3553,39 @@ INFO:Neuron:Finished weights loading in 19.251377552980557 seconds INFO:Neuron:Warming up the model. INFO:Neuron:Warmup completed in 0.05246591567993164 seconds. Can't set use_return_dict with value True for PretrainedConfig { - "fused_spec_config": null, - "hidden_act": "silu", - "hidden_size": 2048, - "intermediate_size": 8192, - "max_position_embeddings": 2048, - "metadata": null, - "num_attention_heads": 32, - "num_cores_per_group": 1, - "num_hidden_layers": 16, - "num_key_value_heads": 8, - "rms_norm_eps": 1e-05, - "rope_theta": 500000.0, - "tie_word_embeddings": false, - "torch_dtype": "float32", - "transformers_version": "4.51.3", - "use_cache": true +"fused_spec_config": null, +"hidden_act": "silu", +"hidden_size": 2048, +"intermediate_size": 8192, +"max_position_embeddings": 2048, +"metadata": null, +"num_attention_heads": 32, +"num_cores_per_group": 1, +"num_hidden_layers": 16, +"num_key_value_heads": 8, +"rms_norm_eps": 1e-05, +"rope_theta": 500000.0, +"tie_word_embeddings": false, +"torch_dtype": "float32", +"transformers_version": "4.51.3", +"use_cache": true } Traceback (most recent call last): - File "/home/ec2-user/NeuronxSDK/neuronx_llama3/simple_inference.py", line 87, in - main() - File "/home/ec2-user/NeuronxSDK/neuronx_llama3/simple_inference.py", line 48, in main - generation_model = HuggingFaceGenerationAdapter(model) - File "/home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inference/utils/hf_adapter.py", line 101, in - __init__ hf_config = to_pretrained_config(model.config) - File "/home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inference/utils/hf_adapter.py", line 96, in -to_pretrained_config return PretrainedConfig(**config_dict) - File "/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages/transformers/configuration_utils.py", line 32 -3, in __init__ raise err - File "/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages/transformers/configuration_utils.py", line 32 -0, in __init__ setattr(self, key, value) - File "/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages/transformers/configuration_utils.py", line 20 -5, in __setattr__ super().__setattr__(key, value) +File "/home/ec2-user/NeuronxSDK/neuronx_llama3/simple_inference.py", line 87, in +main() +File "/home/ec2-user/NeuronxSDK/neuronx_llama3/simple_inference.py", line 48, in main +generation_model = HuggingFaceGenerationAdapter(model) +File "/home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inference/utils/hf_adapter.py", line 101, in +**init** hf_config = to_pretrained_config(model.config) +File "/home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inference/utils/hf_adapter.py", line 96, in +to_pretrained_config return PretrainedConfig(\*\*config_dict) +File "/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages/transformers/configuration_utils.py", line 32 +3, in **init** raise err +File "/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages/transformers/configuration_utils.py", line 32 +0, in **init** setattr(self, key, value) +File "/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages/transformers/configuration_utils.py", line 20 +5, in **setattr** super().**setattr**(key, value) AttributeError: can't set attribute 'use_return_dict' Great progress! The tokenizer is now loading successfully, but there's a configuration issue with the HuggingFaceGenerationAdapter. The error shows that use_return_dict can't be set on the config. This is a compatibility issue between the model config and the HuggingFace adapter. @@ -3621,758 +3598,754 @@ The config looks fine. The issue is that the HuggingFaceGenerationAdapter is try Let me try a different approach - let me use the model's built-in generate method directly instead of the HuggingFaceGenerationAdapter: - 2 edits to file simple_inference.py - Now let's test the simplified version: Command python simple_inference.py --model_path ./llama3_compiled --prompt "Hello, how are you?" --max_new_tokens 5 - are you?" --max_new_tokens 5 /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/pa -rallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +are you?" --max_new_tokens 5 /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/pa +rallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/pa -rallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +rallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/pa -rallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +rallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/mo -dules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +dules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/mo -dules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +dules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/mo -dules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +dules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/mo -dules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +dules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/mo -dules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +dules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/mo -dules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +dules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distr -ibuted_inference/modules/attention/utils.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.custom_calls import neuron -_cumsum /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distr -ibuted_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, -GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distr -ibuted_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, -GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distr -ibuted_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, -GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_llama3/modeling_ll -ama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base i -mport NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_llama3/modeling_ll -ama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base i -mport NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_llama3/modeling_ll -ama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base i -mport NeuronAttentionBase WARNING:root:NeuronConfig init: Unexpected keyword arguments: {'apply_s -eq_ids_mask': False, 'enable_long_context_mode': False, 'enable_output_completion_notifications': False, 'enable_token_tree': False, 'is_chunked_prefill': False, 'is_prefill_stage': None, 'kv_cache_tiling': False, 'scratchpad_page_size': None, 'skip_warmup': False, 'tile_cc': False, 'weights_to_skip_layout_optimization': []} INFO:Neuron:Sharding weights on load... +ibuted_inference/modules/attention/utils.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.custom_calls import neuron +\_cumsum /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distr +ibuted_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, +GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distr +ibuted_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, +GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distr +ibuted_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, +GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_llama3/modeling_ll +ama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base i +mport NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_llama3/modeling_ll +ama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base i +mport NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_llama3/modeling_ll +ama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base i +mport NeuronAttentionBase WARNING:root:NeuronConfig init: Unexpected keyword arguments: {'apply_s +eq_ids_mask': False, 'enable_long_context_mode': False, 'enable_output_completion_notifications': False, 'enable_token_tree': False, 'is_chunked_prefill': False, 'is_prefill_stage': None, 'kv_cache_tiling': False, 'scratchpad_page_size': None, 'skip_warmup': False, 'tile_cc': False, 'weights_to_skip_layout_optimization': []} INFO:Neuron:Sharding weights on load... INFO:Neuron:Sharding Weights for ranks: 0...0 [2025-07-26 17:07:52.416: I neuronx_distributed/parallel_layers/paralle -l_state.py:628] > initializing tensor model parallel with size 1 [2025-07-26 17:07:52.417: I neuronx_distributed/parallel_layers/paralle -l_state.py:629] > initializing pipeline model parallel with size 1 [2025-07-26 17:07:52.417: I neuronx_distributed/parallel_layers/paralle -l_state.py:630] > initializing context model parallel with size 1 [2025-07-26 17:07:52.417: I neuronx_distributed/parallel_layers/paralle -l_state.py:631] > initializing data parallel with size 1 [2025-07-26 17:07:52.417: I neuronx_distributed/parallel_layers/paralle -l_state.py:632] > initializing world size to 1 [2025-07-26 17:07:52.418: I neuronx_distributed/parallel_layers/paralle -l_state.py:379] [rank_0_pp-1_tp-1_dp-1_cp-1] Chosen Logic for replica groups ret_logic=, 'Ascending Ring PG Group')> [2025-07-26 17:07:52.418: I neuronx_distributed/parallel_layers/paralle -l_state.py:668] [rank_0_pp-1_tp-1_dp-1_cp-1] tp_groups: replica_groups.tp_groups=[[0]] [2025-07-26 17:07:52.418: I neuronx_distributed/parallel_layers/paralle -l_state.py:669] [rank_0_pp-1_tp-1_dp-1_cp-1] dp_groups: replica_groups.dp_groups=[[0]] [2025-07-26 17:07:52.418: I neuronx_distributed/parallel_layers/paralle -l_state.py:670] [rank_0_pp-1_tp-1_dp-1_cp-1] pp_groups: replica_groups.pp_groups=[[0]] [2025-07-26 17:07:52.418: I neuronx_distributed/parallel_layers/paralle -l_state.py:671] [rank_0_pp-1_tp-1_dp-1_cp-1] cp_groups: replica_groups.cp_groups=[[0]] [2025-07-26 17:07:52.418: I neuronx_distributed/parallel_layers/paralle -l_state.py:672] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_model_groups: replica_groups.ep_model_groups=[[0]] [2025-07-26 17:07:52.419: I neuronx_distributed/parallel_layers/paralle -l_state.py:673] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_data_groups: replica_groups.ep_data_groups=[[0]] WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/tr -ace/trace.py:640: UserWarning: Removing redundant keys from checkpoint: ['layers.0.self_attn.q_proj.weight', 'layers.0.self_attn.k_proj.weight', 'layers.0.self_attn.v_proj.weight', 'layers.0.self_attn.o_proj.weight', 'layers.1.self_attn.q_proj.weight', 'layers.1.self_attn.k_proj.weight', 'layers.1.self_attn.v_proj.weight', 'layers.1.self_attn.o_proj.weight', 'layers.2.self_attn.q_proj.weight', 'layers.2.self_attn.k_proj.weight', 'layers.2.self_attn.v_proj.weight', 'layers.2.self_attn.o_proj.weight', 'layers.3.self_attn.q_proj.weight', 'layers.3.self_attn.k_proj.weight', 'layers.3.self_attn.v_proj.weight', 'layers.3.self_attn.o_proj.weight', 'layers.4.self_attn.q_proj.weight', 'layers.4.self_attn.k_proj.weight', 'layers.4.self_attn.v_proj.weight', 'layers.4.self_attn.o_proj.weight', 'layers.5.self_attn.q_proj.weight', 'layers.5.self_attn.k_proj.weight', 'layers.5.self_attn.v_proj.weight', 'layers.5.self_attn.o_proj.weight', 'layers.6.self_attn.q_proj.weight', 'layers.6.self_attn.k_proj.weight', 'layers.6.self_attn.v_proj.weight', 'layers.6.self_attn.o_proj.weight', 'layers.7.self_attn.q_proj.weight', 'layers.7.self_attn.k_proj.weight', 'layers.7.self_attn.v_proj.weight', 'layers.7.self_attn.o_proj.weight', 'layers.8.self_attn.q_proj.weight', 'layers.8.self_attn.k_proj.weight', 'layers.8.self_attn.v_proj.weight', 'layers.8.self_attn.o_proj.weight', 'layers.9.self_attn.q_proj.weight', 'layers.9.self_attn.k_proj.weight', 'layers.9.self_attn.v_proj.weight', 'layers.9.self_attn.o_proj.weight', 'layers.10.self_attn.q_proj.weight', 'layers.10.self_attn.k_proj.weight', 'layers.10.self_attn.v_proj.weight', 'layers.10.self_attn.o_proj.weight', 'layers.11.self_attn.q_proj.weight', 'layers.11.self_attn.k_proj.weight', 'layers.11.self_attn.v_proj.weight', 'layers.11.self_attn.o_proj.weight', 'layers.12.self_attn.q_proj.weight', 'layers.12.self_attn.k_proj.weight', 'layers.12.self_attn.v_proj.weight', 'layers.12.self_attn.o_proj.weight', 'layers.13.self_attn.q_proj.weight', 'layers.13.self_attn.k_proj.weight', 'layers.13.self_attn.v_proj.weight', 'layers.13.self_attn.o_proj.weight', 'layers.14.self_attn.q_proj.weight', 'layers.14.self_attn.k_proj.weight', 'layers.14.self_attn.v_proj.weight', 'layers.14.self_attn.o_proj.weight', 'layers.15.self_attn.q_proj.weight', 'layers.15.self_attn.k_proj.weight', 'layers.15.self_attn.v_proj.weight', 'layers.15.self_attn.o_proj.weight'] warnings.warn(f"Removing redundant keys from checkpoint: {keys_to_del -ete}") WARNING:Neuron:casting embed_tokens.weight from torch.bfloat16 to torch -.float32 WARNING:Neuron:casting layers.0.self_attn.qkv_proj.q_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.0.self_attn.qkv_proj.k_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.0.self_attn.qkv_proj.v_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.0.self_attn.o_proj.o_proj.weight from tor -ch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.0.mlp.gate_proj.weight from torch.bfloat1 -6 to torch.float32 WARNING:Neuron:casting layers.0.mlp.up_proj.weight from torch.bfloat16 -to torch.float32 WARNING:Neuron:casting layers.0.mlp.down_proj.weight from torch.bfloat1 -6 to torch.float32 WARNING:Neuron:casting layers.0.input_layernorm.weight from torch.bfloa -t16 to torch.float32 WARNING:Neuron:casting layers.0.post_attention_layernorm.weight from to -rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.1.self_attn.qkv_proj.q_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.1.self_attn.qkv_proj.k_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.1.self_attn.qkv_proj.v_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.1.self_attn.o_proj.o_proj.weight from tor -ch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.1.mlp.gate_proj.weight from torch.bfloat1 -6 to torch.float32 WARNING:Neuron:casting layers.1.mlp.up_proj.weight from torch.bfloat16 -to torch.float32 WARNING:Neuron:casting layers.1.mlp.down_proj.weight from torch.bfloat1 -6 to torch.float32 WARNING:Neuron:casting layers.1.input_layernorm.weight from torch.bfloa -t16 to torch.float32 WARNING:Neuron:casting layers.1.post_attention_layernorm.weight from to -rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.2.self_attn.qkv_proj.q_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.2.self_attn.qkv_proj.k_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.2.self_attn.qkv_proj.v_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.2.self_attn.o_proj.o_proj.weight from tor -ch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.2.mlp.gate_proj.weight from torch.bfloat1 -6 to torch.float32 WARNING:Neuron:casting layers.2.mlp.up_proj.weight from torch.bfloat16 -to torch.float32 WARNING:Neuron:casting layers.2.mlp.down_proj.weight from torch.bfloat1 -6 to torch.float32 WARNING:Neuron:casting layers.2.input_layernorm.weight from torch.bfloa -t16 to torch.float32 WARNING:Neuron:casting layers.2.post_attention_layernorm.weight from to -rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.3.self_attn.qkv_proj.q_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.3.self_attn.qkv_proj.k_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.3.self_attn.qkv_proj.v_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.3.self_attn.o_proj.o_proj.weight from tor -ch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.3.mlp.gate_proj.weight from torch.bfloat1 -6 to torch.float32 WARNING:Neuron:casting layers.3.mlp.up_proj.weight from torch.bfloat16 -to torch.float32 WARNING:Neuron:casting layers.3.mlp.down_proj.weight from torch.bfloat1 -6 to torch.float32 WARNING:Neuron:casting layers.3.input_layernorm.weight from torch.bfloa -t16 to torch.float32 WARNING:Neuron:casting layers.3.post_attention_layernorm.weight from to -rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.4.self_attn.qkv_proj.q_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.4.self_attn.qkv_proj.k_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.4.self_attn.qkv_proj.v_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.4.self_attn.o_proj.o_proj.weight from tor -ch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.4.mlp.gate_proj.weight from torch.bfloat1 -6 to torch.float32 WARNING:Neuron:casting layers.4.mlp.up_proj.weight from torch.bfloat16 -to torch.float32 WARNING:Neuron:casting layers.4.mlp.down_proj.weight from torch.bfloat1 -6 to torch.float32 WARNING:Neuron:casting layers.4.input_layernorm.weight from torch.bfloa -t16 to torch.float32 WARNING:Neuron:casting layers.4.post_attention_layernorm.weight from to -rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.5.self_attn.qkv_proj.q_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.5.self_attn.qkv_proj.k_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.5.self_attn.qkv_proj.v_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.5.self_attn.o_proj.o_proj.weight from tor -ch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.5.mlp.gate_proj.weight from torch.bfloat1 -6 to torch.float32 WARNING:Neuron:casting layers.5.mlp.up_proj.weight from torch.bfloat16 -to torch.float32 WARNING:Neuron:casting layers.5.mlp.down_proj.weight from torch.bfloat1 -6 to torch.float32 WARNING:Neuron:casting layers.5.input_layernorm.weight from torch.bfloa -t16 to torch.float32 WARNING:Neuron:casting layers.5.post_attention_layernorm.weight from to -rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.6.self_attn.qkv_proj.q_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.6.self_attn.qkv_proj.k_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.6.self_attn.qkv_proj.v_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.6.self_attn.o_proj.o_proj.weight from tor -ch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.6.mlp.gate_proj.weight from torch.bfloat1 -6 to torch.float32 WARNING:Neuron:casting layers.6.mlp.up_proj.weight from torch.bfloat16 -to torch.float32 WARNING:Neuron:casting layers.6.mlp.down_proj.weight from torch.bfloat1 -6 to torch.float32 WARNING:Neuron:casting layers.6.input_layernorm.weight from torch.bfloa -t16 to torch.float32 WARNING:Neuron:casting layers.6.post_attention_layernorm.weight from to -rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.7.self_attn.qkv_proj.q_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.7.self_attn.qkv_proj.k_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.7.self_attn.qkv_proj.v_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.7.self_attn.o_proj.o_proj.weight from tor -ch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.7.mlp.gate_proj.weight from torch.bfloat1 -6 to torch.float32 WARNING:Neuron:casting layers.7.mlp.up_proj.weight from torch.bfloat16 -to torch.float32 WARNING:Neuron:casting layers.7.mlp.down_proj.weight from torch.bfloat1 -6 to torch.float32 WARNING:Neuron:casting layers.7.input_layernorm.weight from torch.bfloa -t16 to torch.float32 WARNING:Neuron:casting layers.7.post_attention_layernorm.weight from to -rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.8.self_attn.qkv_proj.q_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.8.self_attn.qkv_proj.k_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.8.self_attn.qkv_proj.v_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.8.self_attn.o_proj.o_proj.weight from tor -ch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.8.mlp.gate_proj.weight from torch.bfloat1 -6 to torch.float32 WARNING:Neuron:casting layers.8.mlp.up_proj.weight from torch.bfloat16 -to torch.float32 WARNING:Neuron:casting layers.8.mlp.down_proj.weight from torch.bfloat1 -6 to torch.float32 WARNING:Neuron:casting layers.8.input_layernorm.weight from torch.bfloa -t16 to torch.float32 WARNING:Neuron:casting layers.8.post_attention_layernorm.weight from to -rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.9.self_attn.qkv_proj.q_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.9.self_attn.qkv_proj.k_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.9.self_attn.qkv_proj.v_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.9.self_attn.o_proj.o_proj.weight from tor -ch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.9.mlp.gate_proj.weight from torch.bfloat1 -6 to torch.float32 WARNING:Neuron:casting layers.9.mlp.up_proj.weight from torch.bfloat16 -to torch.float32 WARNING:Neuron:casting layers.9.mlp.down_proj.weight from torch.bfloat1 -6 to torch.float32 WARNING:Neuron:casting layers.9.input_layernorm.weight from torch.bfloa -t16 to torch.float32 WARNING:Neuron:casting layers.9.post_attention_layernorm.weight from to -rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.10.self_attn.qkv_proj.q_proj.weight from -torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.10.self_attn.qkv_proj.k_proj.weight from -torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.10.self_attn.qkv_proj.v_proj.weight from -torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.10.self_attn.o_proj.o_proj.weight from to -rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.10.mlp.gate_proj.weight from torch.bfloat -16 to torch.float32 WARNING:Neuron:casting layers.10.mlp.up_proj.weight from torch.bfloat16 - to torch.float32 WARNING:Neuron:casting layers.10.mlp.down_proj.weight from torch.bfloat -16 to torch.float32 WARNING:Neuron:casting layers.10.input_layernorm.weight from torch.bflo -at16 to torch.float32 WARNING:Neuron:casting layers.10.post_attention_layernorm.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.11.self_attn.qkv_proj.q_proj.weight from -torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.11.self_attn.qkv_proj.k_proj.weight from -torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.11.self_attn.qkv_proj.v_proj.weight from -torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.11.self_attn.o_proj.o_proj.weight from to -rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.11.mlp.gate_proj.weight from torch.bfloat -16 to torch.float32 WARNING:Neuron:casting layers.11.mlp.up_proj.weight from torch.bfloat16 - to torch.float32 WARNING:Neuron:casting layers.11.mlp.down_proj.weight from torch.bfloat -16 to torch.float32 WARNING:Neuron:casting layers.11.input_layernorm.weight from torch.bflo -at16 to torch.float32 WARNING:Neuron:casting layers.11.post_attention_layernorm.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.12.self_attn.qkv_proj.q_proj.weight from -torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.12.self_attn.qkv_proj.k_proj.weight from -torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.12.self_attn.qkv_proj.v_proj.weight from -torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.12.self_attn.o_proj.o_proj.weight from to -rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.12.mlp.gate_proj.weight from torch.bfloat -16 to torch.float32 WARNING:Neuron:casting layers.12.mlp.up_proj.weight from torch.bfloat16 - to torch.float32 WARNING:Neuron:casting layers.12.mlp.down_proj.weight from torch.bfloat -16 to torch.float32 WARNING:Neuron:casting layers.12.input_layernorm.weight from torch.bflo -at16 to torch.float32 WARNING:Neuron:casting layers.12.post_attention_layernorm.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.13.self_attn.qkv_proj.q_proj.weight from -torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.13.self_attn.qkv_proj.k_proj.weight from -torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.13.self_attn.qkv_proj.v_proj.weight from -torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.13.self_attn.o_proj.o_proj.weight from to -rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.13.mlp.gate_proj.weight from torch.bfloat -16 to torch.float32 WARNING:Neuron:casting layers.13.mlp.up_proj.weight from torch.bfloat16 - to torch.float32 WARNING:Neuron:casting layers.13.mlp.down_proj.weight from torch.bfloat -16 to torch.float32 WARNING:Neuron:casting layers.13.input_layernorm.weight from torch.bflo -at16 to torch.float32 WARNING:Neuron:casting layers.13.post_attention_layernorm.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.14.self_attn.qkv_proj.q_proj.weight from -torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.14.self_attn.qkv_proj.k_proj.weight from -torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.14.self_attn.qkv_proj.v_proj.weight from -torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.14.self_attn.o_proj.o_proj.weight from to -rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.14.mlp.gate_proj.weight from torch.bfloat -16 to torch.float32 WARNING:Neuron:casting layers.14.mlp.up_proj.weight from torch.bfloat16 - to torch.float32 WARNING:Neuron:casting layers.14.mlp.down_proj.weight from torch.bfloat -16 to torch.float32 WARNING:Neuron:casting layers.14.input_layernorm.weight from torch.bflo -at16 to torch.float32 WARNING:Neuron:casting layers.14.post_attention_layernorm.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.15.self_attn.qkv_proj.q_proj.weight from -torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.15.self_attn.qkv_proj.k_proj.weight from -torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.15.self_attn.qkv_proj.v_proj.weight from -torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.15.self_attn.o_proj.o_proj.weight from to -rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.15.mlp.gate_proj.weight from torch.bfloat -16 to torch.float32 WARNING:Neuron:casting layers.15.mlp.up_proj.weight from torch.bfloat16 - to torch.float32 WARNING:Neuron:casting layers.15.mlp.down_proj.weight from torch.bfloat -16 to torch.float32 WARNING:Neuron:casting layers.15.input_layernorm.weight from torch.bflo -at16 to torch.float32 WARNING:Neuron:casting layers.15.post_attention_layernorm.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting norm.weight from torch.bfloat16 to torch.float32 +l_state.py:628] > initializing tensor model parallel with size 1 [2025-07-26 17:07:52.417: I neuronx_distributed/parallel_layers/paralle +l_state.py:629] > initializing pipeline model parallel with size 1 [2025-07-26 17:07:52.417: I neuronx_distributed/parallel_layers/paralle +l_state.py:630] > initializing context model parallel with size 1 [2025-07-26 17:07:52.417: I neuronx_distributed/parallel_layers/paralle +l_state.py:631] > initializing data parallel with size 1 [2025-07-26 17:07:52.417: I neuronx_distributed/parallel_layers/paralle +l_state.py:632] > initializing world size to 1 [2025-07-26 17:07:52.418: I neuronx_distributed/parallel_layers/paralle +l_state.py:379] [rank_0_pp-1_tp-1_dp-1_cp-1] Chosen Logic for replica groups ret_logic=, 'Ascending Ring PG Group')> [2025-07-26 17:07:52.418: I neuronx_distributed/parallel_layers/paralle +l_state.py:668] [rank_0_pp-1_tp-1_dp-1_cp-1] tp_groups: replica_groups.tp_groups=[[0]] [2025-07-26 17:07:52.418: I neuronx_distributed/parallel_layers/paralle +l_state.py:669] [rank_0_pp-1_tp-1_dp-1_cp-1] dp_groups: replica_groups.dp_groups=[[0]] [2025-07-26 17:07:52.418: I neuronx_distributed/parallel_layers/paralle +l_state.py:670] [rank_0_pp-1_tp-1_dp-1_cp-1] pp_groups: replica_groups.pp_groups=[[0]] [2025-07-26 17:07:52.418: I neuronx_distributed/parallel_layers/paralle +l_state.py:671] [rank_0_pp-1_tp-1_dp-1_cp-1] cp_groups: replica_groups.cp_groups=[[0]] [2025-07-26 17:07:52.418: I neuronx_distributed/parallel_layers/paralle +l_state.py:672] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_model_groups: replica_groups.ep_model_groups=[[0]] [2025-07-26 17:07:52.419: I neuronx_distributed/parallel_layers/paralle +l_state.py:673] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_data_groups: replica_groups.ep_data_groups=[[0]] WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/tr +ace/trace.py:640: UserWarning: Removing redundant keys from checkpoint: ['layers.0.self_attn.q_proj.weight', 'layers.0.self_attn.k_proj.weight', 'layers.0.self_attn.v_proj.weight', 'layers.0.self_attn.o_proj.weight', 'layers.1.self_attn.q_proj.weight', 'layers.1.self_attn.k_proj.weight', 'layers.1.self_attn.v_proj.weight', 'layers.1.self_attn.o_proj.weight', 'layers.2.self_attn.q_proj.weight', 'layers.2.self_attn.k_proj.weight', 'layers.2.self_attn.v_proj.weight', 'layers.2.self_attn.o_proj.weight', 'layers.3.self_attn.q_proj.weight', 'layers.3.self_attn.k_proj.weight', 'layers.3.self_attn.v_proj.weight', 'layers.3.self_attn.o_proj.weight', 'layers.4.self_attn.q_proj.weight', 'layers.4.self_attn.k_proj.weight', 'layers.4.self_attn.v_proj.weight', 'layers.4.self_attn.o_proj.weight', 'layers.5.self_attn.q_proj.weight', 'layers.5.self_attn.k_proj.weight', 'layers.5.self_attn.v_proj.weight', 'layers.5.self_attn.o_proj.weight', 'layers.6.self_attn.q_proj.weight', 'layers.6.self_attn.k_proj.weight', 'layers.6.self_attn.v_proj.weight', 'layers.6.self_attn.o_proj.weight', 'layers.7.self_attn.q_proj.weight', 'layers.7.self_attn.k_proj.weight', 'layers.7.self_attn.v_proj.weight', 'layers.7.self_attn.o_proj.weight', 'layers.8.self_attn.q_proj.weight', 'layers.8.self_attn.k_proj.weight', 'layers.8.self_attn.v_proj.weight', 'layers.8.self_attn.o_proj.weight', 'layers.9.self_attn.q_proj.weight', 'layers.9.self_attn.k_proj.weight', 'layers.9.self_attn.v_proj.weight', 'layers.9.self_attn.o_proj.weight', 'layers.10.self_attn.q_proj.weight', 'layers.10.self_attn.k_proj.weight', 'layers.10.self_attn.v_proj.weight', 'layers.10.self_attn.o_proj.weight', 'layers.11.self_attn.q_proj.weight', 'layers.11.self_attn.k_proj.weight', 'layers.11.self_attn.v_proj.weight', 'layers.11.self_attn.o_proj.weight', 'layers.12.self_attn.q_proj.weight', 'layers.12.self_attn.k_proj.weight', 'layers.12.self_attn.v_proj.weight', 'layers.12.self_attn.o_proj.weight', 'layers.13.self_attn.q_proj.weight', 'layers.13.self_attn.k_proj.weight', 'layers.13.self_attn.v_proj.weight', 'layers.13.self_attn.o_proj.weight', 'layers.14.self_attn.q_proj.weight', 'layers.14.self_attn.k_proj.weight', 'layers.14.self_attn.v_proj.weight', 'layers.14.self_attn.o_proj.weight', 'layers.15.self_attn.q_proj.weight', 'layers.15.self_attn.k_proj.weight', 'layers.15.self_attn.v_proj.weight', 'layers.15.self_attn.o_proj.weight'] warnings.warn(f"Removing redundant keys from checkpoint: {keys_to_del +ete}") WARNING:Neuron:casting embed_tokens.weight from torch.bfloat16 to torch +.float32 WARNING:Neuron:casting layers.0.self_attn.qkv_proj.q_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.0.self_attn.qkv_proj.k_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.0.self_attn.qkv_proj.v_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.0.self_attn.o_proj.o_proj.weight from tor +ch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.0.mlp.gate_proj.weight from torch.bfloat1 +6 to torch.float32 WARNING:Neuron:casting layers.0.mlp.up_proj.weight from torch.bfloat16 +to torch.float32 WARNING:Neuron:casting layers.0.mlp.down_proj.weight from torch.bfloat1 +6 to torch.float32 WARNING:Neuron:casting layers.0.input_layernorm.weight from torch.bfloa +t16 to torch.float32 WARNING:Neuron:casting layers.0.post_attention_layernorm.weight from to +rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.1.self_attn.qkv_proj.q_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.1.self_attn.qkv_proj.k_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.1.self_attn.qkv_proj.v_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.1.self_attn.o_proj.o_proj.weight from tor +ch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.1.mlp.gate_proj.weight from torch.bfloat1 +6 to torch.float32 WARNING:Neuron:casting layers.1.mlp.up_proj.weight from torch.bfloat16 +to torch.float32 WARNING:Neuron:casting layers.1.mlp.down_proj.weight from torch.bfloat1 +6 to torch.float32 WARNING:Neuron:casting layers.1.input_layernorm.weight from torch.bfloa +t16 to torch.float32 WARNING:Neuron:casting layers.1.post_attention_layernorm.weight from to +rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.2.self_attn.qkv_proj.q_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.2.self_attn.qkv_proj.k_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.2.self_attn.qkv_proj.v_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.2.self_attn.o_proj.o_proj.weight from tor +ch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.2.mlp.gate_proj.weight from torch.bfloat1 +6 to torch.float32 WARNING:Neuron:casting layers.2.mlp.up_proj.weight from torch.bfloat16 +to torch.float32 WARNING:Neuron:casting layers.2.mlp.down_proj.weight from torch.bfloat1 +6 to torch.float32 WARNING:Neuron:casting layers.2.input_layernorm.weight from torch.bfloa +t16 to torch.float32 WARNING:Neuron:casting layers.2.post_attention_layernorm.weight from to +rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.3.self_attn.qkv_proj.q_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.3.self_attn.qkv_proj.k_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.3.self_attn.qkv_proj.v_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.3.self_attn.o_proj.o_proj.weight from tor +ch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.3.mlp.gate_proj.weight from torch.bfloat1 +6 to torch.float32 WARNING:Neuron:casting layers.3.mlp.up_proj.weight from torch.bfloat16 +to torch.float32 WARNING:Neuron:casting layers.3.mlp.down_proj.weight from torch.bfloat1 +6 to torch.float32 WARNING:Neuron:casting layers.3.input_layernorm.weight from torch.bfloa +t16 to torch.float32 WARNING:Neuron:casting layers.3.post_attention_layernorm.weight from to +rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.4.self_attn.qkv_proj.q_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.4.self_attn.qkv_proj.k_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.4.self_attn.qkv_proj.v_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.4.self_attn.o_proj.o_proj.weight from tor +ch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.4.mlp.gate_proj.weight from torch.bfloat1 +6 to torch.float32 WARNING:Neuron:casting layers.4.mlp.up_proj.weight from torch.bfloat16 +to torch.float32 WARNING:Neuron:casting layers.4.mlp.down_proj.weight from torch.bfloat1 +6 to torch.float32 WARNING:Neuron:casting layers.4.input_layernorm.weight from torch.bfloa +t16 to torch.float32 WARNING:Neuron:casting layers.4.post_attention_layernorm.weight from to +rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.5.self_attn.qkv_proj.q_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.5.self_attn.qkv_proj.k_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.5.self_attn.qkv_proj.v_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.5.self_attn.o_proj.o_proj.weight from tor +ch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.5.mlp.gate_proj.weight from torch.bfloat1 +6 to torch.float32 WARNING:Neuron:casting layers.5.mlp.up_proj.weight from torch.bfloat16 +to torch.float32 WARNING:Neuron:casting layers.5.mlp.down_proj.weight from torch.bfloat1 +6 to torch.float32 WARNING:Neuron:casting layers.5.input_layernorm.weight from torch.bfloa +t16 to torch.float32 WARNING:Neuron:casting layers.5.post_attention_layernorm.weight from to +rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.6.self_attn.qkv_proj.q_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.6.self_attn.qkv_proj.k_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.6.self_attn.qkv_proj.v_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.6.self_attn.o_proj.o_proj.weight from tor +ch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.6.mlp.gate_proj.weight from torch.bfloat1 +6 to torch.float32 WARNING:Neuron:casting layers.6.mlp.up_proj.weight from torch.bfloat16 +to torch.float32 WARNING:Neuron:casting layers.6.mlp.down_proj.weight from torch.bfloat1 +6 to torch.float32 WARNING:Neuron:casting layers.6.input_layernorm.weight from torch.bfloa +t16 to torch.float32 WARNING:Neuron:casting layers.6.post_attention_layernorm.weight from to +rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.7.self_attn.qkv_proj.q_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.7.self_attn.qkv_proj.k_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.7.self_attn.qkv_proj.v_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.7.self_attn.o_proj.o_proj.weight from tor +ch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.7.mlp.gate_proj.weight from torch.bfloat1 +6 to torch.float32 WARNING:Neuron:casting layers.7.mlp.up_proj.weight from torch.bfloat16 +to torch.float32 WARNING:Neuron:casting layers.7.mlp.down_proj.weight from torch.bfloat1 +6 to torch.float32 WARNING:Neuron:casting layers.7.input_layernorm.weight from torch.bfloa +t16 to torch.float32 WARNING:Neuron:casting layers.7.post_attention_layernorm.weight from to +rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.8.self_attn.qkv_proj.q_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.8.self_attn.qkv_proj.k_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.8.self_attn.qkv_proj.v_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.8.self_attn.o_proj.o_proj.weight from tor +ch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.8.mlp.gate_proj.weight from torch.bfloat1 +6 to torch.float32 WARNING:Neuron:casting layers.8.mlp.up_proj.weight from torch.bfloat16 +to torch.float32 WARNING:Neuron:casting layers.8.mlp.down_proj.weight from torch.bfloat1 +6 to torch.float32 WARNING:Neuron:casting layers.8.input_layernorm.weight from torch.bfloa +t16 to torch.float32 WARNING:Neuron:casting layers.8.post_attention_layernorm.weight from to +rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.9.self_attn.qkv_proj.q_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.9.self_attn.qkv_proj.k_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.9.self_attn.qkv_proj.v_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.9.self_attn.o_proj.o_proj.weight from tor +ch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.9.mlp.gate_proj.weight from torch.bfloat1 +6 to torch.float32 WARNING:Neuron:casting layers.9.mlp.up_proj.weight from torch.bfloat16 +to torch.float32 WARNING:Neuron:casting layers.9.mlp.down_proj.weight from torch.bfloat1 +6 to torch.float32 WARNING:Neuron:casting layers.9.input_layernorm.weight from torch.bfloa +t16 to torch.float32 WARNING:Neuron:casting layers.9.post_attention_layernorm.weight from to +rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.10.self_attn.qkv_proj.q_proj.weight from +torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.10.self_attn.qkv_proj.k_proj.weight from +torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.10.self_attn.qkv_proj.v_proj.weight from +torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.10.self_attn.o_proj.o_proj.weight from to +rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.10.mlp.gate_proj.weight from torch.bfloat +16 to torch.float32 WARNING:Neuron:casting layers.10.mlp.up_proj.weight from torch.bfloat16 +to torch.float32 WARNING:Neuron:casting layers.10.mlp.down_proj.weight from torch.bfloat +16 to torch.float32 WARNING:Neuron:casting layers.10.input_layernorm.weight from torch.bflo +at16 to torch.float32 WARNING:Neuron:casting layers.10.post_attention_layernorm.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.11.self_attn.qkv_proj.q_proj.weight from +torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.11.self_attn.qkv_proj.k_proj.weight from +torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.11.self_attn.qkv_proj.v_proj.weight from +torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.11.self_attn.o_proj.o_proj.weight from to +rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.11.mlp.gate_proj.weight from torch.bfloat +16 to torch.float32 WARNING:Neuron:casting layers.11.mlp.up_proj.weight from torch.bfloat16 +to torch.float32 WARNING:Neuron:casting layers.11.mlp.down_proj.weight from torch.bfloat +16 to torch.float32 WARNING:Neuron:casting layers.11.input_layernorm.weight from torch.bflo +at16 to torch.float32 WARNING:Neuron:casting layers.11.post_attention_layernorm.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.12.self_attn.qkv_proj.q_proj.weight from +torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.12.self_attn.qkv_proj.k_proj.weight from +torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.12.self_attn.qkv_proj.v_proj.weight from +torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.12.self_attn.o_proj.o_proj.weight from to +rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.12.mlp.gate_proj.weight from torch.bfloat +16 to torch.float32 WARNING:Neuron:casting layers.12.mlp.up_proj.weight from torch.bfloat16 +to torch.float32 WARNING:Neuron:casting layers.12.mlp.down_proj.weight from torch.bfloat +16 to torch.float32 WARNING:Neuron:casting layers.12.input_layernorm.weight from torch.bflo +at16 to torch.float32 WARNING:Neuron:casting layers.12.post_attention_layernorm.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.13.self_attn.qkv_proj.q_proj.weight from +torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.13.self_attn.qkv_proj.k_proj.weight from +torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.13.self_attn.qkv_proj.v_proj.weight from +torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.13.self_attn.o_proj.o_proj.weight from to +rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.13.mlp.gate_proj.weight from torch.bfloat +16 to torch.float32 WARNING:Neuron:casting layers.13.mlp.up_proj.weight from torch.bfloat16 +to torch.float32 WARNING:Neuron:casting layers.13.mlp.down_proj.weight from torch.bfloat +16 to torch.float32 WARNING:Neuron:casting layers.13.input_layernorm.weight from torch.bflo +at16 to torch.float32 WARNING:Neuron:casting layers.13.post_attention_layernorm.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.14.self_attn.qkv_proj.q_proj.weight from +torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.14.self_attn.qkv_proj.k_proj.weight from +torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.14.self_attn.qkv_proj.v_proj.weight from +torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.14.self_attn.o_proj.o_proj.weight from to +rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.14.mlp.gate_proj.weight from torch.bfloat +16 to torch.float32 WARNING:Neuron:casting layers.14.mlp.up_proj.weight from torch.bfloat16 +to torch.float32 WARNING:Neuron:casting layers.14.mlp.down_proj.weight from torch.bfloat +16 to torch.float32 WARNING:Neuron:casting layers.14.input_layernorm.weight from torch.bflo +at16 to torch.float32 WARNING:Neuron:casting layers.14.post_attention_layernorm.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.15.self_attn.qkv_proj.q_proj.weight from +torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.15.self_attn.qkv_proj.k_proj.weight from +torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.15.self_attn.qkv_proj.v_proj.weight from +torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.15.self_attn.o_proj.o_proj.weight from to +rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.15.mlp.gate_proj.weight from torch.bfloat +16 to torch.float32 WARNING:Neuron:casting layers.15.mlp.up_proj.weight from torch.bfloat16 +to torch.float32 WARNING:Neuron:casting layers.15.mlp.down_proj.weight from torch.bfloat +16 to torch.float32 WARNING:Neuron:casting layers.15.input_layernorm.weight from torch.bflo +at16 to torch.float32 WARNING:Neuron:casting layers.15.post_attention_layernorm.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting norm.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting lm_head.weight from torch.bfloat16 to torch.floa -t32 INFO:Neuron:Done Sharding weights in 0.506343635963276 +t32 INFO:Neuron:Done Sharding weights in 0.506343635963276 INFO:Neuron:Finished weights loading in 12.236704375129193 seconds INFO:Neuron:Warming up the model. INFO:Neuron:Warmup completed in 0.05245780944824219 seconds. Traceback (most recent call last): - File "/home/ec2-user/NeuronxSDK/neuronx_llama3/simple_inference.py", -line 67, in main() - File "/home/ec2-user/NeuronxSDK/neuronx_llama3/simple_inference.py", -line 49, in main outputs = model.generate( - File "/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/ -site-packages/torch/nn/modules/module.py", line 1940, in __getattr__ raise AttributeError( +File "/home/ec2-user/NeuronxSDK/neuronx_llama3/simple_inference.py", +line 67, in main() +File "/home/ec2-user/NeuronxSDK/neuronx_llama3/simple_inference.py", +line 49, in main outputs = model.generate( +File "/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/ +site-packages/torch/nn/modules/module.py", line 1940, in **getattr** raise AttributeError( AttributeError: 'NeuronLlama3ForCausalLM' object has no attribute 'gene -rate' +rate' Perfect! Now we're getting to the core issue. The model doesn't have a generate method. This is exactly what the HuggingFaceGenerationAdapter was supposed to provide. Let me check what methods the model actually has and implement a simple generation loop: Accepted edits to simple_inference.py - Now let's test this simplified generation approach: Command python simple_inference.py --model_path ./llama3_compiled --prompt "Hello, how are you?" --max_new_tokens 5 /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/pa -rallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +rallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/pa -rallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +rallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/pa -rallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +rallel_layers/layers.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/mo -dules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +dules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/mo -dules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +dules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/mo -dules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +dules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/mo -dules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +dules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/mo -dules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +dules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/mo -dules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +dules/moe/blockwise.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distr -ibuted_inference/modules/attention/utils.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.custom_calls import neuron -_cumsum /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distr -ibuted_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, -GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distr -ibuted_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, -GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distr -ibuted_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, -GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_llama3/modeling_ll -ama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base i -mport NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_llama3/modeling_ll -ama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base i -mport NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_llama3/modeling_ll -ama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base i -mport NeuronAttentionBase WARNING:root:NeuronConfig init: Unexpected keyword arguments: {'apply_s -eq_ids_mask': False, 'enable_long_context_mode': False, 'enable_output_completion_notifications': False, 'enable_token_tree': False, 'is_chunked_prefill': False, 'is_prefill_stage': None, 'kv_cache_tiling': False, 'scratchpad_page_size': None, 'skip_warmup': False, 'tile_cc': False, 'weights_to_skip_layout_optimization': []} INFO:Neuron:Sharding weights on load... +ibuted_inference/modules/attention/utils.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.custom_calls import neuron +\_cumsum /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distr +ibuted_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, +GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distr +ibuted_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, +GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distr +ibuted_inference/modules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, +GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_llama3/modeling_ll +ama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base i +mport NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_llama3/modeling_ll +ama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base i +mport NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_llama3/modeling_ll +ama3.py:61: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base i +mport NeuronAttentionBase WARNING:root:NeuronConfig init: Unexpected keyword arguments: {'apply_s +eq_ids_mask': False, 'enable_long_context_mode': False, 'enable_output_completion_notifications': False, 'enable_token_tree': False, 'is_chunked_prefill': False, 'is_prefill_stage': None, 'kv_cache_tiling': False, 'scratchpad_page_size': None, 'skip_warmup': False, 'tile_cc': False, 'weights_to_skip_layout_optimization': []} INFO:Neuron:Sharding weights on load... INFO:Neuron:Sharding Weights for ranks: 0...0 [2025-07-26 17:08:45.719: I neuronx_distributed/parallel_layers/paralle -l_state.py:628] > initializing tensor model parallel with size 1 [2025-07-26 17:08:45.720: I neuronx_distributed/parallel_layers/paralle -l_state.py:629] > initializing pipeline model parallel with size 1 [2025-07-26 17:08:45.720: I neuronx_distributed/parallel_layers/paralle -l_state.py:630] > initializing context model parallel with size 1 [2025-07-26 17:08:45.720: I neuronx_distributed/parallel_layers/paralle -l_state.py:631] > initializing data parallel with size 1 [2025-07-26 17:08:45.720: I neuronx_distributed/parallel_layers/paralle -l_state.py:632] > initializing world size to 1 [2025-07-26 17:08:45.721: I neuronx_distributed/parallel_layers/paralle -l_state.py:379] [rank_0_pp-1_tp-1_dp-1_cp-1] Chosen Logic for replica groups ret_logic=, 'Ascending Ring PG Group')> [2025-07-26 17:08:45.721: I neuronx_distributed/parallel_layers/paralle -l_state.py:668] [rank_0_pp-1_tp-1_dp-1_cp-1] tp_groups: replica_groups.tp_groups=[[0]] [2025-07-26 17:08:45.721: I neuronx_distributed/parallel_layers/paralle -l_state.py:669] [rank_0_pp-1_tp-1_dp-1_cp-1] dp_groups: replica_groups.dp_groups=[[0]] [2025-07-26 17:08:45.721: I neuronx_distributed/parallel_layers/paralle -l_state.py:670] [rank_0_pp-1_tp-1_dp-1_cp-1] pp_groups: replica_groups.pp_groups=[[0]] [2025-07-26 17:08:45.721: I neuronx_distributed/parallel_layers/paralle -l_state.py:671] [rank_0_pp-1_tp-1_dp-1_cp-1] cp_groups: replica_groups.cp_groups=[[0]] [2025-07-26 17:08:45.721: I neuronx_distributed/parallel_layers/paralle -l_state.py:672] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_model_groups: replica_groups.ep_model_groups=[[0]] [2025-07-26 17:08:45.722: I neuronx_distributed/parallel_layers/paralle -l_state.py:673] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_data_groups: replica_groups.ep_data_groups=[[0]] WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri -ding attention sharding strategy to GQA.CONVERT_TO_MHA! /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/tr -ace/trace.py:640: UserWarning: Removing redundant keys from checkpoint: ['layers.0.self_attn.q_proj.weight', 'layers.0.self_attn.k_proj.weight', 'layers.0.self_attn.v_proj.weight', 'layers.0.self_attn.o_proj.weight', 'layers.1.self_attn.q_proj.weight', 'layers.1.self_attn.k_proj.weight', 'layers.1.self_attn.v_proj.weight', 'layers.1.self_attn.o_proj.weight', 'layers.2.self_attn.q_proj.weight', 'layers.2.self_attn.k_proj.weight', 'layers.2.self_attn.v_proj.weight', 'layers.2.self_attn.o_proj.weight', 'layers.3.self_attn.q_proj.weight', 'layers.3.self_attn.k_proj.weight', 'layers.3.self_attn.v_proj.weight', 'layers.3.self_attn.o_proj.weight', 'layers.4.self_attn.q_proj.weight', 'layers.4.self_attn.k_proj.weight', 'layers.4.self_attn.v_proj.weight', 'layers.4.self_attn.o_proj.weight', 'layers.5.self_attn.q_proj.weight', 'layers.5.self_attn.k_proj.weight', 'layers.5.self_attn.v_proj.weight', 'layers.5.self_attn.o_proj.weight', 'layers.6.self_attn.q_proj.weight', 'layers.6.self_attn.k_proj.weight', 'layers.6.self_attn.v_proj.weight', 'layers.6.self_attn.o_proj.weight', 'layers.7.self_attn.q_proj.weight', 'layers.7.self_attn.k_proj.weight', 'layers.7.self_attn.v_proj.weight', 'layers.7.self_attn.o_proj.weight', 'layers.8.self_attn.q_proj.weight', 'layers.8.self_attn.k_proj.weight', 'layers.8.self_attn.v_proj.weight', 'layers.8.self_attn.o_proj.weight', 'layers.9.self_attn.q_proj.weight', 'layers.9.self_attn.k_proj.weight', 'layers.9.self_attn.v_proj.weight', 'layers.9.self_attn.o_proj.weight', 'layers.10.self_attn.q_proj.weight', 'layers.10.self_attn.k_proj.weight', 'layers.10.self_attn.v_proj.weight', 'layers.10.self_attn.o_proj.weight', 'layers.11.self_attn.q_proj.weight', 'layers.11.self_attn.k_proj.weight', 'layers.11.self_attn.v_proj.weight', 'layers.11.self_attn.o_proj.weight', 'layers.12.self_attn.q_proj.weight', 'layers.12.self_attn.k_proj.weight', 'layers.12.self_attn.v_proj.weight', 'layers.12.self_attn.o_proj.weight', 'layers.13.self_attn.q_proj.weight', 'layers.13.self_attn.k_proj.weight', 'layers.13.self_attn.v_proj.weight', 'layers.13.self_attn.o_proj.weight', 'layers.14.self_attn.q_proj.weight', 'layers.14.self_attn.k_proj.weight', 'layers.14.self_attn.v_proj.weight', 'layers.14.self_attn.o_proj.weight', 'layers.15.self_attn.q_proj.weight', 'layers.15.self_attn.k_proj.weight', 'layers.15.self_attn.v_proj.weight', 'layers.15.self_attn.o_proj.weight'] warnings.warn(f"Removing redundant keys from checkpoint: {keys_to_del -ete}") WARNING:Neuron:casting embed_tokens.weight from torch.bfloat16 to torch -.float32 WARNING:Neuron:casting layers.0.self_attn.qkv_proj.q_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.0.self_attn.qkv_proj.k_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.0.self_attn.qkv_proj.v_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.0.self_attn.o_proj.o_proj.weight from tor -ch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.0.mlp.gate_proj.weight from torch.bfloat1 -6 to torch.float32 WARNING:Neuron:casting layers.0.mlp.up_proj.weight from torch.bfloat16 -to torch.float32 WARNING:Neuron:casting layers.0.mlp.down_proj.weight from torch.bfloat1 -6 to torch.float32 WARNING:Neuron:casting layers.0.input_layernorm.weight from torch.bfloa -t16 to torch.float32 WARNING:Neuron:casting layers.0.post_attention_layernorm.weight from to -rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.1.self_attn.qkv_proj.q_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.1.self_attn.qkv_proj.k_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.1.self_attn.qkv_proj.v_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.1.self_attn.o_proj.o_proj.weight from tor -ch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.1.mlp.gate_proj.weight from torch.bfloat1 -6 to torch.float32 WARNING:Neuron:casting layers.1.mlp.up_proj.weight from torch.bfloat16 -to torch.float32 WARNING:Neuron:casting layers.1.mlp.down_proj.weight from torch.bfloat1 -6 to torch.float32 WARNING:Neuron:casting layers.1.input_layernorm.weight from torch.bfloa -t16 to torch.float32 WARNING:Neuron:casting layers.1.post_attention_layernorm.weight from to -rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.2.self_attn.qkv_proj.q_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.2.self_attn.qkv_proj.k_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.2.self_attn.qkv_proj.v_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.2.self_attn.o_proj.o_proj.weight from tor -ch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.2.mlp.gate_proj.weight from torch.bfloat1 -6 to torch.float32 WARNING:Neuron:casting layers.2.mlp.up_proj.weight from torch.bfloat16 -to torch.float32 WARNING:Neuron:casting layers.2.mlp.down_proj.weight from torch.bfloat1 -6 to torch.float32 WARNING:Neuron:casting layers.2.input_layernorm.weight from torch.bfloa -t16 to torch.float32 WARNING:Neuron:casting layers.2.post_attention_layernorm.weight from to -rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.3.self_attn.qkv_proj.q_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.3.self_attn.qkv_proj.k_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.3.self_attn.qkv_proj.v_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.3.self_attn.o_proj.o_proj.weight from tor -ch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.3.mlp.gate_proj.weight from torch.bfloat1 -6 to torch.float32 WARNING:Neuron:casting layers.3.mlp.up_proj.weight from torch.bfloat16 -to torch.float32 WARNING:Neuron:casting layers.3.mlp.down_proj.weight from torch.bfloat1 -6 to torch.float32 WARNING:Neuron:casting layers.3.input_layernorm.weight from torch.bfloa -t16 to torch.float32 WARNING:Neuron:casting layers.3.post_attention_layernorm.weight from to -rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.4.self_attn.qkv_proj.q_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.4.self_attn.qkv_proj.k_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.4.self_attn.qkv_proj.v_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.4.self_attn.o_proj.o_proj.weight from tor -ch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.4.mlp.gate_proj.weight from torch.bfloat1 -6 to torch.float32 WARNING:Neuron:casting layers.4.mlp.up_proj.weight from torch.bfloat16 -to torch.float32 WARNING:Neuron:casting layers.4.mlp.down_proj.weight from torch.bfloat1 -6 to torch.float32 WARNING:Neuron:casting layers.4.input_layernorm.weight from torch.bfloa -t16 to torch.float32 WARNING:Neuron:casting layers.4.post_attention_layernorm.weight from to -rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.5.self_attn.qkv_proj.q_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.5.self_attn.qkv_proj.k_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.5.self_attn.qkv_proj.v_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.5.self_attn.o_proj.o_proj.weight from tor -ch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.5.mlp.gate_proj.weight from torch.bfloat1 -6 to torch.float32 WARNING:Neuron:casting layers.5.mlp.up_proj.weight from torch.bfloat16 -to torch.float32 WARNING:Neuron:casting layers.5.mlp.down_proj.weight from torch.bfloat1 -6 to torch.float32 WARNING:Neuron:casting layers.5.input_layernorm.weight from torch.bfloa -t16 to torch.float32 WARNING:Neuron:casting layers.5.post_attention_layernorm.weight from to -rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.6.self_attn.qkv_proj.q_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.6.self_attn.qkv_proj.k_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.6.self_attn.qkv_proj.v_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.6.self_attn.o_proj.o_proj.weight from tor -ch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.6.mlp.gate_proj.weight from torch.bfloat1 -6 to torch.float32 WARNING:Neuron:casting layers.6.mlp.up_proj.weight from torch.bfloat16 -to torch.float32 WARNING:Neuron:casting layers.6.mlp.down_proj.weight from torch.bfloat1 -6 to torch.float32 WARNING:Neuron:casting layers.6.input_layernorm.weight from torch.bfloa -t16 to torch.float32 WARNING:Neuron:casting layers.6.post_attention_layernorm.weight from to -rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.7.self_attn.qkv_proj.q_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.7.self_attn.qkv_proj.k_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.7.self_attn.qkv_proj.v_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.7.self_attn.o_proj.o_proj.weight from tor -ch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.7.mlp.gate_proj.weight from torch.bfloat1 -6 to torch.float32 WARNING:Neuron:casting layers.7.mlp.up_proj.weight from torch.bfloat16 -to torch.float32 WARNING:Neuron:casting layers.7.mlp.down_proj.weight from torch.bfloat1 -6 to torch.float32 WARNING:Neuron:casting layers.7.input_layernorm.weight from torch.bfloa -t16 to torch.float32 WARNING:Neuron:casting layers.7.post_attention_layernorm.weight from to -rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.8.self_attn.qkv_proj.q_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.8.self_attn.qkv_proj.k_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.8.self_attn.qkv_proj.v_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.8.self_attn.o_proj.o_proj.weight from tor -ch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.8.mlp.gate_proj.weight from torch.bfloat1 -6 to torch.float32 WARNING:Neuron:casting layers.8.mlp.up_proj.weight from torch.bfloat16 -to torch.float32 WARNING:Neuron:casting layers.8.mlp.down_proj.weight from torch.bfloat1 -6 to torch.float32 WARNING:Neuron:casting layers.8.input_layernorm.weight from torch.bfloa -t16 to torch.float32 WARNING:Neuron:casting layers.8.post_attention_layernorm.weight from to -rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.9.self_attn.qkv_proj.q_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.9.self_attn.qkv_proj.k_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.9.self_attn.qkv_proj.v_proj.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.9.self_attn.o_proj.o_proj.weight from tor -ch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.9.mlp.gate_proj.weight from torch.bfloat1 -6 to torch.float32 WARNING:Neuron:casting layers.9.mlp.up_proj.weight from torch.bfloat16 -to torch.float32 WARNING:Neuron:casting layers.9.mlp.down_proj.weight from torch.bfloat1 -6 to torch.float32 WARNING:Neuron:casting layers.9.input_layernorm.weight from torch.bfloa -t16 to torch.float32 WARNING:Neuron:casting layers.9.post_attention_layernorm.weight from to -rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.10.self_attn.qkv_proj.q_proj.weight from -torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.10.self_attn.qkv_proj.k_proj.weight from -torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.10.self_attn.qkv_proj.v_proj.weight from -torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.10.self_attn.o_proj.o_proj.weight from to -rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.10.mlp.gate_proj.weight from torch.bfloat -16 to torch.float32 WARNING:Neuron:casting layers.10.mlp.up_proj.weight from torch.bfloat16 - to torch.float32 WARNING:Neuron:casting layers.10.mlp.down_proj.weight from torch.bfloat -16 to torch.float32 WARNING:Neuron:casting layers.10.input_layernorm.weight from torch.bflo -at16 to torch.float32 WARNING:Neuron:casting layers.10.post_attention_layernorm.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.11.self_attn.qkv_proj.q_proj.weight from -torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.11.self_attn.qkv_proj.k_proj.weight from -torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.11.self_attn.qkv_proj.v_proj.weight from -torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.11.self_attn.o_proj.o_proj.weight from to -rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.11.mlp.gate_proj.weight from torch.bfloat -16 to torch.float32 WARNING:Neuron:casting layers.11.mlp.up_proj.weight from torch.bfloat16 - to torch.float32 WARNING:Neuron:casting layers.11.mlp.down_proj.weight from torch.bfloat -16 to torch.float32 WARNING:Neuron:casting layers.11.input_layernorm.weight from torch.bflo -at16 to torch.float32 WARNING:Neuron:casting layers.11.post_attention_layernorm.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.12.self_attn.qkv_proj.q_proj.weight from -torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.12.self_attn.qkv_proj.k_proj.weight from -torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.12.self_attn.qkv_proj.v_proj.weight from -torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.12.self_attn.o_proj.o_proj.weight from to -rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.12.mlp.gate_proj.weight from torch.bfloat -16 to torch.float32 WARNING:Neuron:casting layers.12.mlp.up_proj.weight from torch.bfloat16 - to torch.float32 WARNING:Neuron:casting layers.12.mlp.down_proj.weight from torch.bfloat -16 to torch.float32 WARNING:Neuron:casting layers.12.input_layernorm.weight from torch.bflo -at16 to torch.float32 WARNING:Neuron:casting layers.12.post_attention_layernorm.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.13.self_attn.qkv_proj.q_proj.weight from -torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.13.self_attn.qkv_proj.k_proj.weight from -torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.13.self_attn.qkv_proj.v_proj.weight from -torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.13.self_attn.o_proj.o_proj.weight from to -rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.13.mlp.gate_proj.weight from torch.bfloat -16 to torch.float32 WARNING:Neuron:casting layers.13.mlp.up_proj.weight from torch.bfloat16 - to torch.float32 WARNING:Neuron:casting layers.13.mlp.down_proj.weight from torch.bfloat -16 to torch.float32 WARNING:Neuron:casting layers.13.input_layernorm.weight from torch.bflo -at16 to torch.float32 WARNING:Neuron:casting layers.13.post_attention_layernorm.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.14.self_attn.qkv_proj.q_proj.weight from -torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.14.self_attn.qkv_proj.k_proj.weight from -torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.14.self_attn.qkv_proj.v_proj.weight from -torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.14.self_attn.o_proj.o_proj.weight from to -rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.14.mlp.gate_proj.weight from torch.bfloat -16 to torch.float32 WARNING:Neuron:casting layers.14.mlp.up_proj.weight from torch.bfloat16 - to torch.float32 WARNING:Neuron:casting layers.14.mlp.down_proj.weight from torch.bfloat -16 to torch.float32 WARNING:Neuron:casting layers.14.input_layernorm.weight from torch.bflo -at16 to torch.float32 WARNING:Neuron:casting layers.14.post_attention_layernorm.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.15.self_attn.qkv_proj.q_proj.weight from -torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.15.self_attn.qkv_proj.k_proj.weight from -torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.15.self_attn.qkv_proj.v_proj.weight from -torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.15.self_attn.o_proj.o_proj.weight from to -rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.15.mlp.gate_proj.weight from torch.bfloat -16 to torch.float32 WARNING:Neuron:casting layers.15.mlp.up_proj.weight from torch.bfloat16 - to torch.float32 WARNING:Neuron:casting layers.15.mlp.down_proj.weight from torch.bfloat -16 to torch.float32 WARNING:Neuron:casting layers.15.input_layernorm.weight from torch.bflo -at16 to torch.float32 WARNING:Neuron:casting layers.15.post_attention_layernorm.weight from t -orch.bfloat16 to torch.float32 WARNING:Neuron:casting norm.weight from torch.bfloat16 to torch.float32 +l_state.py:628] > initializing tensor model parallel with size 1 [2025-07-26 17:08:45.720: I neuronx_distributed/parallel_layers/paralle +l_state.py:629] > initializing pipeline model parallel with size 1 [2025-07-26 17:08:45.720: I neuronx_distributed/parallel_layers/paralle +l_state.py:630] > initializing context model parallel with size 1 [2025-07-26 17:08:45.720: I neuronx_distributed/parallel_layers/paralle +l_state.py:631] > initializing data parallel with size 1 [2025-07-26 17:08:45.720: I neuronx_distributed/parallel_layers/paralle +l_state.py:632] > initializing world size to 1 [2025-07-26 17:08:45.721: I neuronx_distributed/parallel_layers/paralle +l_state.py:379] [rank_0_pp-1_tp-1_dp-1_cp-1] Chosen Logic for replica groups ret_logic=, 'Ascending Ring PG Group')> [2025-07-26 17:08:45.721: I neuronx_distributed/parallel_layers/paralle +l_state.py:668] [rank_0_pp-1_tp-1_dp-1_cp-1] tp_groups: replica_groups.tp_groups=[[0]] [2025-07-26 17:08:45.721: I neuronx_distributed/parallel_layers/paralle +l_state.py:669] [rank_0_pp-1_tp-1_dp-1_cp-1] dp_groups: replica_groups.dp_groups=[[0]] [2025-07-26 17:08:45.721: I neuronx_distributed/parallel_layers/paralle +l_state.py:670] [rank_0_pp-1_tp-1_dp-1_cp-1] pp_groups: replica_groups.pp_groups=[[0]] [2025-07-26 17:08:45.721: I neuronx_distributed/parallel_layers/paralle +l_state.py:671] [rank_0_pp-1_tp-1_dp-1_cp-1] cp_groups: replica_groups.cp_groups=[[0]] [2025-07-26 17:08:45.721: I neuronx_distributed/parallel_layers/paralle +l_state.py:672] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_model_groups: replica_groups.ep_model_groups=[[0]] [2025-07-26 17:08:45.722: I neuronx_distributed/parallel_layers/paralle +l_state.py:673] [rank_0_pp-1_tp-1_dp-1_cp-1] ep_data_groups: replica_groups.ep_data_groups=[[0]] WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overri +ding attention sharding strategy to GQA.CONVERT_TO_MHA! /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/tr +ace/trace.py:640: UserWarning: Removing redundant keys from checkpoint: ['layers.0.self_attn.q_proj.weight', 'layers.0.self_attn.k_proj.weight', 'layers.0.self_attn.v_proj.weight', 'layers.0.self_attn.o_proj.weight', 'layers.1.self_attn.q_proj.weight', 'layers.1.self_attn.k_proj.weight', 'layers.1.self_attn.v_proj.weight', 'layers.1.self_attn.o_proj.weight', 'layers.2.self_attn.q_proj.weight', 'layers.2.self_attn.k_proj.weight', 'layers.2.self_attn.v_proj.weight', 'layers.2.self_attn.o_proj.weight', 'layers.3.self_attn.q_proj.weight', 'layers.3.self_attn.k_proj.weight', 'layers.3.self_attn.v_proj.weight', 'layers.3.self_attn.o_proj.weight', 'layers.4.self_attn.q_proj.weight', 'layers.4.self_attn.k_proj.weight', 'layers.4.self_attn.v_proj.weight', 'layers.4.self_attn.o_proj.weight', 'layers.5.self_attn.q_proj.weight', 'layers.5.self_attn.k_proj.weight', 'layers.5.self_attn.v_proj.weight', 'layers.5.self_attn.o_proj.weight', 'layers.6.self_attn.q_proj.weight', 'layers.6.self_attn.k_proj.weight', 'layers.6.self_attn.v_proj.weight', 'layers.6.self_attn.o_proj.weight', 'layers.7.self_attn.q_proj.weight', 'layers.7.self_attn.k_proj.weight', 'layers.7.self_attn.v_proj.weight', 'layers.7.self_attn.o_proj.weight', 'layers.8.self_attn.q_proj.weight', 'layers.8.self_attn.k_proj.weight', 'layers.8.self_attn.v_proj.weight', 'layers.8.self_attn.o_proj.weight', 'layers.9.self_attn.q_proj.weight', 'layers.9.self_attn.k_proj.weight', 'layers.9.self_attn.v_proj.weight', 'layers.9.self_attn.o_proj.weight', 'layers.10.self_attn.q_proj.weight', 'layers.10.self_attn.k_proj.weight', 'layers.10.self_attn.v_proj.weight', 'layers.10.self_attn.o_proj.weight', 'layers.11.self_attn.q_proj.weight', 'layers.11.self_attn.k_proj.weight', 'layers.11.self_attn.v_proj.weight', 'layers.11.self_attn.o_proj.weight', 'layers.12.self_attn.q_proj.weight', 'layers.12.self_attn.k_proj.weight', 'layers.12.self_attn.v_proj.weight', 'layers.12.self_attn.o_proj.weight', 'layers.13.self_attn.q_proj.weight', 'layers.13.self_attn.k_proj.weight', 'layers.13.self_attn.v_proj.weight', 'layers.13.self_attn.o_proj.weight', 'layers.14.self_attn.q_proj.weight', 'layers.14.self_attn.k_proj.weight', 'layers.14.self_attn.v_proj.weight', 'layers.14.self_attn.o_proj.weight', 'layers.15.self_attn.q_proj.weight', 'layers.15.self_attn.k_proj.weight', 'layers.15.self_attn.v_proj.weight', 'layers.15.self_attn.o_proj.weight'] warnings.warn(f"Removing redundant keys from checkpoint: {keys_to_del +ete}") WARNING:Neuron:casting embed_tokens.weight from torch.bfloat16 to torch +.float32 WARNING:Neuron:casting layers.0.self_attn.qkv_proj.q_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.0.self_attn.qkv_proj.k_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.0.self_attn.qkv_proj.v_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.0.self_attn.o_proj.o_proj.weight from tor +ch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.0.mlp.gate_proj.weight from torch.bfloat1 +6 to torch.float32 WARNING:Neuron:casting layers.0.mlp.up_proj.weight from torch.bfloat16 +to torch.float32 WARNING:Neuron:casting layers.0.mlp.down_proj.weight from torch.bfloat1 +6 to torch.float32 WARNING:Neuron:casting layers.0.input_layernorm.weight from torch.bfloa +t16 to torch.float32 WARNING:Neuron:casting layers.0.post_attention_layernorm.weight from to +rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.1.self_attn.qkv_proj.q_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.1.self_attn.qkv_proj.k_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.1.self_attn.qkv_proj.v_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.1.self_attn.o_proj.o_proj.weight from tor +ch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.1.mlp.gate_proj.weight from torch.bfloat1 +6 to torch.float32 WARNING:Neuron:casting layers.1.mlp.up_proj.weight from torch.bfloat16 +to torch.float32 WARNING:Neuron:casting layers.1.mlp.down_proj.weight from torch.bfloat1 +6 to torch.float32 WARNING:Neuron:casting layers.1.input_layernorm.weight from torch.bfloa +t16 to torch.float32 WARNING:Neuron:casting layers.1.post_attention_layernorm.weight from to +rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.2.self_attn.qkv_proj.q_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.2.self_attn.qkv_proj.k_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.2.self_attn.qkv_proj.v_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.2.self_attn.o_proj.o_proj.weight from tor +ch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.2.mlp.gate_proj.weight from torch.bfloat1 +6 to torch.float32 WARNING:Neuron:casting layers.2.mlp.up_proj.weight from torch.bfloat16 +to torch.float32 WARNING:Neuron:casting layers.2.mlp.down_proj.weight from torch.bfloat1 +6 to torch.float32 WARNING:Neuron:casting layers.2.input_layernorm.weight from torch.bfloa +t16 to torch.float32 WARNING:Neuron:casting layers.2.post_attention_layernorm.weight from to +rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.3.self_attn.qkv_proj.q_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.3.self_attn.qkv_proj.k_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.3.self_attn.qkv_proj.v_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.3.self_attn.o_proj.o_proj.weight from tor +ch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.3.mlp.gate_proj.weight from torch.bfloat1 +6 to torch.float32 WARNING:Neuron:casting layers.3.mlp.up_proj.weight from torch.bfloat16 +to torch.float32 WARNING:Neuron:casting layers.3.mlp.down_proj.weight from torch.bfloat1 +6 to torch.float32 WARNING:Neuron:casting layers.3.input_layernorm.weight from torch.bfloa +t16 to torch.float32 WARNING:Neuron:casting layers.3.post_attention_layernorm.weight from to +rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.4.self_attn.qkv_proj.q_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.4.self_attn.qkv_proj.k_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.4.self_attn.qkv_proj.v_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.4.self_attn.o_proj.o_proj.weight from tor +ch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.4.mlp.gate_proj.weight from torch.bfloat1 +6 to torch.float32 WARNING:Neuron:casting layers.4.mlp.up_proj.weight from torch.bfloat16 +to torch.float32 WARNING:Neuron:casting layers.4.mlp.down_proj.weight from torch.bfloat1 +6 to torch.float32 WARNING:Neuron:casting layers.4.input_layernorm.weight from torch.bfloa +t16 to torch.float32 WARNING:Neuron:casting layers.4.post_attention_layernorm.weight from to +rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.5.self_attn.qkv_proj.q_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.5.self_attn.qkv_proj.k_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.5.self_attn.qkv_proj.v_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.5.self_attn.o_proj.o_proj.weight from tor +ch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.5.mlp.gate_proj.weight from torch.bfloat1 +6 to torch.float32 WARNING:Neuron:casting layers.5.mlp.up_proj.weight from torch.bfloat16 +to torch.float32 WARNING:Neuron:casting layers.5.mlp.down_proj.weight from torch.bfloat1 +6 to torch.float32 WARNING:Neuron:casting layers.5.input_layernorm.weight from torch.bfloa +t16 to torch.float32 WARNING:Neuron:casting layers.5.post_attention_layernorm.weight from to +rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.6.self_attn.qkv_proj.q_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.6.self_attn.qkv_proj.k_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.6.self_attn.qkv_proj.v_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.6.self_attn.o_proj.o_proj.weight from tor +ch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.6.mlp.gate_proj.weight from torch.bfloat1 +6 to torch.float32 WARNING:Neuron:casting layers.6.mlp.up_proj.weight from torch.bfloat16 +to torch.float32 WARNING:Neuron:casting layers.6.mlp.down_proj.weight from torch.bfloat1 +6 to torch.float32 WARNING:Neuron:casting layers.6.input_layernorm.weight from torch.bfloa +t16 to torch.float32 WARNING:Neuron:casting layers.6.post_attention_layernorm.weight from to +rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.7.self_attn.qkv_proj.q_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.7.self_attn.qkv_proj.k_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.7.self_attn.qkv_proj.v_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.7.self_attn.o_proj.o_proj.weight from tor +ch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.7.mlp.gate_proj.weight from torch.bfloat1 +6 to torch.float32 WARNING:Neuron:casting layers.7.mlp.up_proj.weight from torch.bfloat16 +to torch.float32 WARNING:Neuron:casting layers.7.mlp.down_proj.weight from torch.bfloat1 +6 to torch.float32 WARNING:Neuron:casting layers.7.input_layernorm.weight from torch.bfloa +t16 to torch.float32 WARNING:Neuron:casting layers.7.post_attention_layernorm.weight from to +rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.8.self_attn.qkv_proj.q_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.8.self_attn.qkv_proj.k_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.8.self_attn.qkv_proj.v_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.8.self_attn.o_proj.o_proj.weight from tor +ch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.8.mlp.gate_proj.weight from torch.bfloat1 +6 to torch.float32 WARNING:Neuron:casting layers.8.mlp.up_proj.weight from torch.bfloat16 +to torch.float32 WARNING:Neuron:casting layers.8.mlp.down_proj.weight from torch.bfloat1 +6 to torch.float32 WARNING:Neuron:casting layers.8.input_layernorm.weight from torch.bfloa +t16 to torch.float32 WARNING:Neuron:casting layers.8.post_attention_layernorm.weight from to +rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.9.self_attn.qkv_proj.q_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.9.self_attn.qkv_proj.k_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.9.self_attn.qkv_proj.v_proj.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.9.self_attn.o_proj.o_proj.weight from tor +ch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.9.mlp.gate_proj.weight from torch.bfloat1 +6 to torch.float32 WARNING:Neuron:casting layers.9.mlp.up_proj.weight from torch.bfloat16 +to torch.float32 WARNING:Neuron:casting layers.9.mlp.down_proj.weight from torch.bfloat1 +6 to torch.float32 WARNING:Neuron:casting layers.9.input_layernorm.weight from torch.bfloa +t16 to torch.float32 WARNING:Neuron:casting layers.9.post_attention_layernorm.weight from to +rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.10.self_attn.qkv_proj.q_proj.weight from +torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.10.self_attn.qkv_proj.k_proj.weight from +torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.10.self_attn.qkv_proj.v_proj.weight from +torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.10.self_attn.o_proj.o_proj.weight from to +rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.10.mlp.gate_proj.weight from torch.bfloat +16 to torch.float32 WARNING:Neuron:casting layers.10.mlp.up_proj.weight from torch.bfloat16 +to torch.float32 WARNING:Neuron:casting layers.10.mlp.down_proj.weight from torch.bfloat +16 to torch.float32 WARNING:Neuron:casting layers.10.input_layernorm.weight from torch.bflo +at16 to torch.float32 WARNING:Neuron:casting layers.10.post_attention_layernorm.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.11.self_attn.qkv_proj.q_proj.weight from +torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.11.self_attn.qkv_proj.k_proj.weight from +torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.11.self_attn.qkv_proj.v_proj.weight from +torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.11.self_attn.o_proj.o_proj.weight from to +rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.11.mlp.gate_proj.weight from torch.bfloat +16 to torch.float32 WARNING:Neuron:casting layers.11.mlp.up_proj.weight from torch.bfloat16 +to torch.float32 WARNING:Neuron:casting layers.11.mlp.down_proj.weight from torch.bfloat +16 to torch.float32 WARNING:Neuron:casting layers.11.input_layernorm.weight from torch.bflo +at16 to torch.float32 WARNING:Neuron:casting layers.11.post_attention_layernorm.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.12.self_attn.qkv_proj.q_proj.weight from +torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.12.self_attn.qkv_proj.k_proj.weight from +torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.12.self_attn.qkv_proj.v_proj.weight from +torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.12.self_attn.o_proj.o_proj.weight from to +rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.12.mlp.gate_proj.weight from torch.bfloat +16 to torch.float32 WARNING:Neuron:casting layers.12.mlp.up_proj.weight from torch.bfloat16 +to torch.float32 WARNING:Neuron:casting layers.12.mlp.down_proj.weight from torch.bfloat +16 to torch.float32 WARNING:Neuron:casting layers.12.input_layernorm.weight from torch.bflo +at16 to torch.float32 WARNING:Neuron:casting layers.12.post_attention_layernorm.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.13.self_attn.qkv_proj.q_proj.weight from +torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.13.self_attn.qkv_proj.k_proj.weight from +torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.13.self_attn.qkv_proj.v_proj.weight from +torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.13.self_attn.o_proj.o_proj.weight from to +rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.13.mlp.gate_proj.weight from torch.bfloat +16 to torch.float32 WARNING:Neuron:casting layers.13.mlp.up_proj.weight from torch.bfloat16 +to torch.float32 WARNING:Neuron:casting layers.13.mlp.down_proj.weight from torch.bfloat +16 to torch.float32 WARNING:Neuron:casting layers.13.input_layernorm.weight from torch.bflo +at16 to torch.float32 WARNING:Neuron:casting layers.13.post_attention_layernorm.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.14.self_attn.qkv_proj.q_proj.weight from +torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.14.self_attn.qkv_proj.k_proj.weight from +torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.14.self_attn.qkv_proj.v_proj.weight from +torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.14.self_attn.o_proj.o_proj.weight from to +rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.14.mlp.gate_proj.weight from torch.bfloat +16 to torch.float32 WARNING:Neuron:casting layers.14.mlp.up_proj.weight from torch.bfloat16 +to torch.float32 WARNING:Neuron:casting layers.14.mlp.down_proj.weight from torch.bfloat +16 to torch.float32 WARNING:Neuron:casting layers.14.input_layernorm.weight from torch.bflo +at16 to torch.float32 WARNING:Neuron:casting layers.14.post_attention_layernorm.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.15.self_attn.qkv_proj.q_proj.weight from +torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.15.self_attn.qkv_proj.k_proj.weight from +torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.15.self_attn.qkv_proj.v_proj.weight from +torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.15.self_attn.o_proj.o_proj.weight from to +rch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.15.mlp.gate_proj.weight from torch.bfloat +16 to torch.float32 WARNING:Neuron:casting layers.15.mlp.up_proj.weight from torch.bfloat16 +to torch.float32 WARNING:Neuron:casting layers.15.mlp.down_proj.weight from torch.bfloat +16 to torch.float32 WARNING:Neuron:casting layers.15.input_layernorm.weight from torch.bflo +at16 to torch.float32 WARNING:Neuron:casting layers.15.post_attention_layernorm.weight from t +orch.bfloat16 to torch.float32 WARNING:Neuron:casting norm.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting lm_head.weight from torch.bfloat16 to torch.floa -t32 INFO:Neuron:Done Sharding weights in 0.5702150280121714 +t32 INFO:Neuron:Done Sharding weights in 0.5702150280121714 INFO:Neuron:Finished weights loading in 12.143054781015962 seconds INFO:Neuron:Warming up the model. INFO:Neuron:Warmup completed in 0.053960323333740234 seconds. Traceback (most recent call last): - File "/home/ec2-user/NeuronxSDK/neuronx_llama3/simple_inference.py", -line 84, in main() - File "/home/ec2-user/NeuronxSDK/neuronx_llama3/simple_inference.py", -line 54, in main outputs = model(generated_ids) - File "/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/ -site-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl return self._call_impl(*args, **kwargs) - File "/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/ -site-packages/torch/nn/modules/module.py", line 1762, in _call_impl return forward_call(*args, **kwargs) - File "/home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuro -nx_distributed_inference/models/model_base.py", line 3178, in forward input_ids, attention_mask, position_ids, seq_ids, sampling_params = - self.preprocess_inputs( File "/home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuro -nx_distributed_inference/models/model_base.py", line 3110, in preprocess_inputs attention_mask = self._infer_attention_mask( - File "/home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuro -nx_distributed_inference/models/model_base.py", line 3311, in _infer_attention_mask position_ids is not None +File "/home/ec2-user/NeuronxSDK/neuronx_llama3/simple_inference.py", +line 84, in main() +File "/home/ec2-user/NeuronxSDK/neuronx_llama3/simple_inference.py", +line 54, in main outputs = model(generated_ids) +File "/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/ +site-packages/torch/nn/modules/module.py", line 1751, in \_wrapped_call_impl return self.\_call_impl(*args, \*\*kwargs) +File "/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/ +site-packages/torch/nn/modules/module.py", line 1762, in \_call_impl return forward_call(*args, \*\*kwargs) +File "/home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuro +nx_distributed_inference/models/model_base.py", line 3178, in forward input_ids, attention_mask, position_ids, seq_ids, sampling_params = +self.preprocess_inputs( File "/home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuro +nx_distributed_inference/models/model_base.py", line 3110, in preprocess_inputs attention_mask = self.\_infer_attention_mask( +File "/home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuro +nx_distributed_inference/models/model_base.py", line 3311, in \_infer_attention_mask position_ids is not None AssertionError: need to call forward with position_ids if attention_mas -k is not provided +k is not provided Great! We're making progress. The model is loading and the forward pass is being called, but it needs position_ids. This is exactly what the working example handles. Let me fix this: Accepted edits to simple_inference.py - Now let's test again: Command python simple_inference.py --model_path ./llama3_compiled --prompt "Hello, how are you?" --max_new_tokens 5 /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/parallel_layers/layers.py:14: - DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/parallel_layers/layers.py:14: - DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/parallel_layers/layers.py:14: - DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( -/home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: -DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) -/home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: -DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) -/home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: -DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) -/home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: -DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) -/home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: -DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) -/home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: -DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +/home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: +DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +/home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: +DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +/home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: +DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +/home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: +DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +/home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: +DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +/home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwise.py:42: +DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inference/modules/at -tention/utils.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.custom_calls import neuron_cumsum +tention/utils.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.custom_calls import neuron_cumsum /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inference/modules/lo -ra_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttention_QKV +ra_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inference/modules/lo -ra_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttention_QKV +ra_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inference/modules/lo -ra_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttention_QKV +ra_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttention_QKV /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_llama3/modeling_llama3.py:61: DeprecationWarn -ing: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base import NeuronAttentionBase +ing: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_llama3/modeling_llama3.py:61: DeprecationWarn -ing: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base import NeuronAttentionBase +ing: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base import NeuronAttentionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_llama3/modeling_llama3.py:61: DeprecationWarn -ing: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base import NeuronAttentionBase +ing: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base import NeuronAttentionBase WARNING:root:NeuronConfig init: Unexpected keyword arguments: {'apply_seq_ids_mask': False, 'enabl -e_long_context_mode': False, 'enable_output_completion_notifications': False, 'enable_token_tree': False, 'is_chunked_prefill': False, 'is_prefill_stage': None, 'kv_cache_tiling': False, 'scratchpad_page_size': None, 'skip_warmup': False, 'tile_cc': False, 'weights_to_skip_layout_optimization': []} INFO:Neuron:Sharding weights on load... +e_long_context_mode': False, 'enable_output_completion_notifications': False, 'enable_token_tree': False, 'is_chunked_prefill': False, 'is_prefill_stage': None, 'kv_cache_tiling': False, 'scratchpad_page_size': None, 'skip_warmup': False, 'tile_cc': False, 'weights_to_skip_layout_optimization': []} INFO:Neuron:Sharding weights on load... INFO:Neuron:Sharding Weights for ranks: 0...0 [2025-07-26 17:09:41.573: I neuronx_distributed/parallel_layers/parallel_state.py:628] > initializ -ing tensor model parallel with size 1 [2025-07-26 17:09:41.573: I neuronx_distributed/parallel_layers/parallel_state.py:629] > initializ -ing pipeline model parallel with size 1 [2025-07-26 17:09:41.573: I neuronx_distributed/parallel_layers/parallel_state.py:630] > initializ -ing context model parallel with size 1 [2025-07-26 17:09:41.573: I neuronx_distributed/parallel_layers/parallel_state.py:631] > initializ -ing data parallel with size 1 [2025-07-26 17:09:41.573: I neuronx_distributed/parallel_layers/parallel_state.py:632] > initializ -ing world size to 1 [2025-07-26 17:09:41.574: I neuronx_distributed/parallel_layers/parallel_state.py:379] [rank_0_pp- -1_tp-1_dp-1_cp-1] Chosen Logic for replica groups ret_logic=, 'Ascending Ring PG Group')> [2025-07-26 17:09:41.575: I neuronx_distributed/parallel_layers/parallel_state.py:668] [rank_0_pp- -1_tp-1_dp-1_cp-1] tp_groups: replica_groups.tp_groups=[[0]] [2025-07-26 17:09:41.575: I neuronx_distributed/parallel_layers/parallel_state.py:669] [rank_0_pp- -1_tp-1_dp-1_cp-1] dp_groups: replica_groups.dp_groups=[[0]] [2025-07-26 17:09:41.575: I neuronx_distributed/parallel_layers/parallel_state.py:670] [rank_0_pp- -1_tp-1_dp-1_cp-1] pp_groups: replica_groups.pp_groups=[[0]] [2025-07-26 17:09:41.575: I neuronx_distributed/parallel_layers/parallel_state.py:671] [rank_0_pp- -1_tp-1_dp-1_cp-1] cp_groups: replica_groups.cp_groups=[[0]] [2025-07-26 17:09:41.575: I neuronx_distributed/parallel_layers/parallel_state.py:672] [rank_0_pp- -1_tp-1_dp-1_cp-1] ep_model_groups: replica_groups.ep_model_groups=[[0]] [2025-07-26 17:09:41.575: I neuronx_distributed/parallel_layers/parallel_state.py:673] [rank_0_pp- -1_tp-1_dp-1_cp-1] ep_data_groups: replica_groups.ep_data_groups=[[0]] WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str -ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str -ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str -ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str -ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str -ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str -ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str -ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str -ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str -ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str -ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str -ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str -ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str -ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str -ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str -ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str -ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str -ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str -ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str -ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str -ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str -ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str -ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str -ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str -ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str -ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str -ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str -ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str -ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str -ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str -ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str -ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str -ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str -ategy to GQA.CONVERT_TO_MHA! /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/trace/trace.py:640: UserWarni -ng: Removing redundant keys from checkpoint: ['layers.0.self_attn.q_proj.weight', 'layers.0.self_attn.k_proj.weight', 'layers.0.self_attn.v_proj.weight', 'layers.0.self_attn.o_proj.weight', 'layers.1.self_attn.q_proj.weight', 'layers.1.self_attn.k_proj.weight', 'layers.1.self_attn.v_proj.weight', 'layers.1.self_attn.o_proj.weight', 'layers.2.self_attn.q_proj.weight', 'layers.2.self_attn.k_proj.weight', 'layers.2.self_attn.v_proj.weight', 'layers.2.self_attn.o_proj.weight', 'layers.3.self_attn.q_proj.weight', 'layers.3.self_attn.k_proj.weight', 'layers.3.self_attn.v_proj.weight', 'layers.3.self_attn.o_proj.weight', 'layers.4.self_attn.q_proj.weight', 'layers.4.self_attn.k_proj.weight', 'layers.4.self_attn.v_proj.weight', 'layers.4.self_attn.o_proj.weight', 'layers.5.self_attn.q_proj.weight', 'layers.5.self_attn.k_proj.weight', 'layers.5.self_attn.v_proj.weight', 'layers.5.self_attn.o_proj.weight', 'layers.6.self_attn.q_proj.weight', 'layers.6.self_attn.k_proj.weight', 'layers.6.self_attn.v_proj.weight', 'layers.6.self_attn.o_proj.weight', 'layers.7.self_attn.q_proj.weight', 'layers.7.self_attn.k_proj.weight', 'layers.7.self_attn.v_proj.weight', 'layers.7.self_attn.o_proj.weight', 'layers.8.self_attn.q_proj.weight', 'layers.8.self_attn.k_proj.weight', 'layers.8.self_attn.v_proj.weight', 'layers.8.self_attn.o_proj.weight', 'layers.9.self_attn.q_proj.weight', 'layers.9.self_attn.k_proj.weight', 'layers.9.self_attn.v_proj.weight', 'layers.9.self_attn.o_proj.weight', 'layers.10.self_attn.q_proj.weight', 'layers.10.self_attn.k_proj.weight', 'layers.10.self_attn.v_proj.weight', 'layers.10.self_attn.o_proj.weight', 'layers.11.self_attn.q_proj.weight', 'layers.11.self_attn.k_proj.weight', 'layers.11.self_attn.v_proj.weight', 'layers.11.self_attn.o_proj.weight', 'layers.12.self_attn.q_proj.weight', 'layers.12.self_attn.k_proj.weight', 'layers.12.self_attn.v_proj.weight', 'layers.12.self_attn.o_proj.weight', 'layers.13.self_attn.q_proj.weight', 'layers.13.self_attn.k_proj.weight', 'layers.13.self_attn.v_proj.weight', 'layers.13.self_attn.o_proj.weight', 'layers.14.self_attn.q_proj.weight', 'layers.14.self_attn.k_proj.weight', 'layers.14.self_attn.v_proj.weight', 'layers.14.self_attn.o_proj.weight', 'layers.15.self_attn.q_proj.weight', 'layers.15.self_attn.k_proj.weight', 'layers.15.self_attn.v_proj.weight', 'layers.15.self_attn.o_proj.weight'] warnings.warn(f"Removing redundant keys from checkpoint: {keys_to_delete}") +ing tensor model parallel with size 1 [2025-07-26 17:09:41.573: I neuronx_distributed/parallel_layers/parallel_state.py:629] > initializ +ing pipeline model parallel with size 1 [2025-07-26 17:09:41.573: I neuronx_distributed/parallel_layers/parallel_state.py:630] > initializ +ing context model parallel with size 1 [2025-07-26 17:09:41.573: I neuronx_distributed/parallel_layers/parallel_state.py:631] > initializ +ing data parallel with size 1 [2025-07-26 17:09:41.573: I neuronx_distributed/parallel_layers/parallel_state.py:632] > initializ +ing world size to 1 [2025-07-26 17:09:41.574: I neuronx_distributed/parallel_layers/parallel_state.py:379] [rank_0_pp- +1_tp-1_dp-1_cp-1] Chosen Logic for replica groups ret_logic=, 'Ascending Ring PG Group')> [2025-07-26 17:09:41.575: I neuronx_distributed/parallel_layers/parallel_state.py:668] [rank_0_pp- +1_tp-1_dp-1_cp-1] tp_groups: replica_groups.tp_groups=[[0]] [2025-07-26 17:09:41.575: I neuronx_distributed/parallel_layers/parallel_state.py:669] [rank_0_pp- +1_tp-1_dp-1_cp-1] dp_groups: replica_groups.dp_groups=[[0]] [2025-07-26 17:09:41.575: I neuronx_distributed/parallel_layers/parallel_state.py:670] [rank_0_pp- +1_tp-1_dp-1_cp-1] pp_groups: replica_groups.pp_groups=[[0]] [2025-07-26 17:09:41.575: I neuronx_distributed/parallel_layers/parallel_state.py:671] [rank_0_pp- +1_tp-1_dp-1_cp-1] cp_groups: replica_groups.cp_groups=[[0]] [2025-07-26 17:09:41.575: I neuronx_distributed/parallel_layers/parallel_state.py:672] [rank_0_pp- +1_tp-1_dp-1_cp-1] ep_model_groups: replica_groups.ep_model_groups=[[0]] [2025-07-26 17:09:41.575: I neuronx_distributed/parallel_layers/parallel_state.py:673] [rank_0_pp- +1_tp-1_dp-1_cp-1] ep_data_groups: replica_groups.ep_data_groups=[[0]] WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str +ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str +ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str +ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str +ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str +ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str +ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str +ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str +ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str +ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str +ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str +ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str +ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str +ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str +ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str +ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str +ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str +ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str +ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str +ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str +ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str +ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str +ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str +ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str +ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str +ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str +ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str +ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str +ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str +ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str +ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str +ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str +ategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sharding str +ategy to GQA.CONVERT_TO_MHA! /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/trace/trace.py:640: UserWarni +ng: Removing redundant keys from checkpoint: ['layers.0.self_attn.q_proj.weight', 'layers.0.self_attn.k_proj.weight', 'layers.0.self_attn.v_proj.weight', 'layers.0.self_attn.o_proj.weight', 'layers.1.self_attn.q_proj.weight', 'layers.1.self_attn.k_proj.weight', 'layers.1.self_attn.v_proj.weight', 'layers.1.self_attn.o_proj.weight', 'layers.2.self_attn.q_proj.weight', 'layers.2.self_attn.k_proj.weight', 'layers.2.self_attn.v_proj.weight', 'layers.2.self_attn.o_proj.weight', 'layers.3.self_attn.q_proj.weight', 'layers.3.self_attn.k_proj.weight', 'layers.3.self_attn.v_proj.weight', 'layers.3.self_attn.o_proj.weight', 'layers.4.self_attn.q_proj.weight', 'layers.4.self_attn.k_proj.weight', 'layers.4.self_attn.v_proj.weight', 'layers.4.self_attn.o_proj.weight', 'layers.5.self_attn.q_proj.weight', 'layers.5.self_attn.k_proj.weight', 'layers.5.self_attn.v_proj.weight', 'layers.5.self_attn.o_proj.weight', 'layers.6.self_attn.q_proj.weight', 'layers.6.self_attn.k_proj.weight', 'layers.6.self_attn.v_proj.weight', 'layers.6.self_attn.o_proj.weight', 'layers.7.self_attn.q_proj.weight', 'layers.7.self_attn.k_proj.weight', 'layers.7.self_attn.v_proj.weight', 'layers.7.self_attn.o_proj.weight', 'layers.8.self_attn.q_proj.weight', 'layers.8.self_attn.k_proj.weight', 'layers.8.self_attn.v_proj.weight', 'layers.8.self_attn.o_proj.weight', 'layers.9.self_attn.q_proj.weight', 'layers.9.self_attn.k_proj.weight', 'layers.9.self_attn.v_proj.weight', 'layers.9.self_attn.o_proj.weight', 'layers.10.self_attn.q_proj.weight', 'layers.10.self_attn.k_proj.weight', 'layers.10.self_attn.v_proj.weight', 'layers.10.self_attn.o_proj.weight', 'layers.11.self_attn.q_proj.weight', 'layers.11.self_attn.k_proj.weight', 'layers.11.self_attn.v_proj.weight', 'layers.11.self_attn.o_proj.weight', 'layers.12.self_attn.q_proj.weight', 'layers.12.self_attn.k_proj.weight', 'layers.12.self_attn.v_proj.weight', 'layers.12.self_attn.o_proj.weight', 'layers.13.self_attn.q_proj.weight', 'layers.13.self_attn.k_proj.weight', 'layers.13.self_attn.v_proj.weight', 'layers.13.self_attn.o_proj.weight', 'layers.14.self_attn.q_proj.weight', 'layers.14.self_attn.k_proj.weight', 'layers.14.self_attn.v_proj.weight', 'layers.14.self_attn.o_proj.weight', 'layers.15.self_attn.q_proj.weight', 'layers.15.self_attn.k_proj.weight', 'layers.15.self_attn.v_proj.weight', 'layers.15.self_attn.o_proj.weight'] warnings.warn(f"Removing redundant keys from checkpoint: {keys_to_delete}") WARNING:Neuron:casting embed_tokens.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.0.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to torch.floa -t32 WARNING:Neuron:casting layers.0.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to torch.floa -t32 WARNING:Neuron:casting layers.0.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to torch.floa -t32 WARNING:Neuron:casting layers.0.self_attn.o_proj.o_proj.weight from torch.bfloat16 to torch.float3 -2 WARNING:Neuron:casting layers.0.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 +t32 WARNING:Neuron:casting layers.0.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to torch.floa +t32 WARNING:Neuron:casting layers.0.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to torch.floa +t32 WARNING:Neuron:casting layers.0.self_attn.o_proj.o_proj.weight from torch.bfloat16 to torch.float3 +2 WARNING:Neuron:casting layers.0.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.0.mlp.up_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.0.mlp.down_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.0.input_layernorm.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.0.post_attention_layernorm.weight from torch.bfloat16 to torch.float -32 WARNING:Neuron:casting layers.1.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to torch.floa -t32 WARNING:Neuron:casting layers.1.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to torch.floa -t32 WARNING:Neuron:casting layers.1.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to torch.floa -t32 WARNING:Neuron:casting layers.1.self_attn.o_proj.o_proj.weight from torch.bfloat16 to torch.float3 -2 WARNING:Neuron:casting layers.1.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 +32 WARNING:Neuron:casting layers.1.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to torch.floa +t32 WARNING:Neuron:casting layers.1.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to torch.floa +t32 WARNING:Neuron:casting layers.1.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to torch.floa +t32 WARNING:Neuron:casting layers.1.self_attn.o_proj.o_proj.weight from torch.bfloat16 to torch.float3 +2 WARNING:Neuron:casting layers.1.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.1.mlp.up_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.1.mlp.down_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.1.input_layernorm.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.1.post_attention_layernorm.weight from torch.bfloat16 to torch.float -32 WARNING:Neuron:casting layers.2.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to torch.floa -t32 WARNING:Neuron:casting layers.2.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to torch.floa -t32 WARNING:Neuron:casting layers.2.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to torch.floa -t32 WARNING:Neuron:casting layers.2.self_attn.o_proj.o_proj.weight from torch.bfloat16 to torch.float3 -2 WARNING:Neuron:casting layers.2.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 +32 WARNING:Neuron:casting layers.2.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to torch.floa +t32 WARNING:Neuron:casting layers.2.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to torch.floa +t32 WARNING:Neuron:casting layers.2.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to torch.floa +t32 WARNING:Neuron:casting layers.2.self_attn.o_proj.o_proj.weight from torch.bfloat16 to torch.float3 +2 WARNING:Neuron:casting layers.2.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.2.mlp.up_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.2.mlp.down_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.2.input_layernorm.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.2.post_attention_layernorm.weight from torch.bfloat16 to torch.float -32 WARNING:Neuron:casting layers.3.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to torch.floa -t32 WARNING:Neuron:casting layers.3.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to torch.floa -t32 WARNING:Neuron:casting layers.3.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to torch.floa -t32 WARNING:Neuron:casting layers.3.self_attn.o_proj.o_proj.weight from torch.bfloat16 to torch.float3 -2 WARNING:Neuron:casting layers.3.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 +32 WARNING:Neuron:casting layers.3.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to torch.floa +t32 WARNING:Neuron:casting layers.3.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to torch.floa +t32 WARNING:Neuron:casting layers.3.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to torch.floa +t32 WARNING:Neuron:casting layers.3.self_attn.o_proj.o_proj.weight from torch.bfloat16 to torch.float3 +2 WARNING:Neuron:casting layers.3.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.3.mlp.up_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.3.mlp.down_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.3.input_layernorm.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.3.post_attention_layernorm.weight from torch.bfloat16 to torch.float -32 WARNING:Neuron:casting layers.4.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to torch.floa -t32 WARNING:Neuron:casting layers.4.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to torch.floa -t32 WARNING:Neuron:casting layers.4.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to torch.floa -t32 WARNING:Neuron:casting layers.4.self_attn.o_proj.o_proj.weight from torch.bfloat16 to torch.float3 -2 WARNING:Neuron:casting layers.4.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 +32 WARNING:Neuron:casting layers.4.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to torch.floa +t32 WARNING:Neuron:casting layers.4.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to torch.floa +t32 WARNING:Neuron:casting layers.4.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to torch.floa +t32 WARNING:Neuron:casting layers.4.self_attn.o_proj.o_proj.weight from torch.bfloat16 to torch.float3 +2 WARNING:Neuron:casting layers.4.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.4.mlp.up_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.4.mlp.down_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.4.input_layernorm.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.4.post_attention_layernorm.weight from torch.bfloat16 to torch.float -32 WARNING:Neuron:casting layers.5.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to torch.floa -t32 WARNING:Neuron:casting layers.5.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to torch.floa -t32 WARNING:Neuron:casting layers.5.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to torch.floa -t32 WARNING:Neuron:casting layers.5.self_attn.o_proj.o_proj.weight from torch.bfloat16 to torch.float3 -2 WARNING:Neuron:casting layers.5.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 +32 WARNING:Neuron:casting layers.5.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to torch.floa +t32 WARNING:Neuron:casting layers.5.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to torch.floa +t32 WARNING:Neuron:casting layers.5.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to torch.floa +t32 WARNING:Neuron:casting layers.5.self_attn.o_proj.o_proj.weight from torch.bfloat16 to torch.float3 +2 WARNING:Neuron:casting layers.5.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.5.mlp.up_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.5.mlp.down_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.5.input_layernorm.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.5.post_attention_layernorm.weight from torch.bfloat16 to torch.float -32 WARNING:Neuron:casting layers.6.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to torch.floa -t32 WARNING:Neuron:casting layers.6.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to torch.floa -t32 WARNING:Neuron:casting layers.6.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to torch.floa -t32 WARNING:Neuron:casting layers.6.self_attn.o_proj.o_proj.weight from torch.bfloat16 to torch.float3 -2 WARNING:Neuron:casting layers.6.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 +32 WARNING:Neuron:casting layers.6.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to torch.floa +t32 WARNING:Neuron:casting layers.6.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to torch.floa +t32 WARNING:Neuron:casting layers.6.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to torch.floa +t32 WARNING:Neuron:casting layers.6.self_attn.o_proj.o_proj.weight from torch.bfloat16 to torch.float3 +2 WARNING:Neuron:casting layers.6.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.6.mlp.up_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.6.mlp.down_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.6.input_layernorm.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.6.post_attention_layernorm.weight from torch.bfloat16 to torch.float -32 WARNING:Neuron:casting layers.7.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to torch.floa -t32 WARNING:Neuron:casting layers.7.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to torch.floa -t32 WARNING:Neuron:casting layers.7.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to torch.floa -t32 WARNING:Neuron:casting layers.7.self_attn.o_proj.o_proj.weight from torch.bfloat16 to torch.float3 -2 WARNING:Neuron:casting layers.7.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 +32 WARNING:Neuron:casting layers.7.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to torch.floa +t32 WARNING:Neuron:casting layers.7.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to torch.floa +t32 WARNING:Neuron:casting layers.7.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to torch.floa +t32 WARNING:Neuron:casting layers.7.self_attn.o_proj.o_proj.weight from torch.bfloat16 to torch.float3 +2 WARNING:Neuron:casting layers.7.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.7.mlp.up_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.7.mlp.down_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.7.input_layernorm.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.7.post_attention_layernorm.weight from torch.bfloat16 to torch.float -32 WARNING:Neuron:casting layers.8.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to torch.floa -t32 WARNING:Neuron:casting layers.8.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to torch.floa -t32 WARNING:Neuron:casting layers.8.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to torch.floa -t32 WARNING:Neuron:casting layers.8.self_attn.o_proj.o_proj.weight from torch.bfloat16 to torch.float3 -2 WARNING:Neuron:casting layers.8.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 +32 WARNING:Neuron:casting layers.8.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to torch.floa +t32 WARNING:Neuron:casting layers.8.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to torch.floa +t32 WARNING:Neuron:casting layers.8.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to torch.floa +t32 WARNING:Neuron:casting layers.8.self_attn.o_proj.o_proj.weight from torch.bfloat16 to torch.float3 +2 WARNING:Neuron:casting layers.8.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.8.mlp.up_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.8.mlp.down_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.8.input_layernorm.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.8.post_attention_layernorm.weight from torch.bfloat16 to torch.float -32 WARNING:Neuron:casting layers.9.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to torch.floa -t32 WARNING:Neuron:casting layers.9.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to torch.floa -t32 WARNING:Neuron:casting layers.9.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to torch.floa -t32 WARNING:Neuron:casting layers.9.self_attn.o_proj.o_proj.weight from torch.bfloat16 to torch.float3 -2 WARNING:Neuron:casting layers.9.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 +32 WARNING:Neuron:casting layers.9.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to torch.floa +t32 WARNING:Neuron:casting layers.9.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to torch.floa +t32 WARNING:Neuron:casting layers.9.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to torch.floa +t32 WARNING:Neuron:casting layers.9.self_attn.o_proj.o_proj.weight from torch.bfloat16 to torch.float3 +2 WARNING:Neuron:casting layers.9.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.9.mlp.up_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.9.mlp.down_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.9.input_layernorm.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.9.post_attention_layernorm.weight from torch.bfloat16 to torch.float -32 WARNING:Neuron:casting layers.10.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to torch.flo -at32 WARNING:Neuron:casting layers.10.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to torch.flo -at32 WARNING:Neuron:casting layers.10.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to torch.flo -at32 WARNING:Neuron:casting layers.10.self_attn.o_proj.o_proj.weight from torch.bfloat16 to torch.float -32 WARNING:Neuron:casting layers.10.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 +32 WARNING:Neuron:casting layers.10.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to torch.flo +at32 WARNING:Neuron:casting layers.10.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to torch.flo +at32 WARNING:Neuron:casting layers.10.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to torch.flo +at32 WARNING:Neuron:casting layers.10.self_attn.o_proj.o_proj.weight from torch.bfloat16 to torch.float +32 WARNING:Neuron:casting layers.10.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.10.mlp.up_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.10.mlp.down_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.10.input_layernorm.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.10.post_attention_layernorm.weight from torch.bfloat16 to torch.floa -t32 WARNING:Neuron:casting layers.11.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to torch.flo -at32 WARNING:Neuron:casting layers.11.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to torch.flo -at32 WARNING:Neuron:casting layers.11.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to torch.flo -at32 WARNING:Neuron:casting layers.11.self_attn.o_proj.o_proj.weight from torch.bfloat16 to torch.float -32 WARNING:Neuron:casting layers.11.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 +t32 WARNING:Neuron:casting layers.11.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to torch.flo +at32 WARNING:Neuron:casting layers.11.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to torch.flo +at32 WARNING:Neuron:casting layers.11.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to torch.flo +at32 WARNING:Neuron:casting layers.11.self_attn.o_proj.o_proj.weight from torch.bfloat16 to torch.float +32 WARNING:Neuron:casting layers.11.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.11.mlp.up_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.11.mlp.down_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.11.input_layernorm.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.11.post_attention_layernorm.weight from torch.bfloat16 to torch.floa -t32 WARNING:Neuron:casting layers.12.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to torch.flo -at32 WARNING:Neuron:casting layers.12.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to torch.flo -at32 WARNING:Neuron:casting layers.12.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to torch.flo -at32 WARNING:Neuron:casting layers.12.self_attn.o_proj.o_proj.weight from torch.bfloat16 to torch.float -32 WARNING:Neuron:casting layers.12.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 +t32 WARNING:Neuron:casting layers.12.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to torch.flo +at32 WARNING:Neuron:casting layers.12.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to torch.flo +at32 WARNING:Neuron:casting layers.12.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to torch.flo +at32 WARNING:Neuron:casting layers.12.self_attn.o_proj.o_proj.weight from torch.bfloat16 to torch.float +32 WARNING:Neuron:casting layers.12.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.12.mlp.up_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.12.mlp.down_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.12.input_layernorm.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.12.post_attention_layernorm.weight from torch.bfloat16 to torch.floa -t32 WARNING:Neuron:casting layers.13.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to torch.flo -at32 WARNING:Neuron:casting layers.13.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to torch.flo -at32 WARNING:Neuron:casting layers.13.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to torch.flo -at32 WARNING:Neuron:casting layers.13.self_attn.o_proj.o_proj.weight from torch.bfloat16 to torch.float -32 WARNING:Neuron:casting layers.13.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 +t32 WARNING:Neuron:casting layers.13.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to torch.flo +at32 WARNING:Neuron:casting layers.13.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to torch.flo +at32 WARNING:Neuron:casting layers.13.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to torch.flo +at32 WARNING:Neuron:casting layers.13.self_attn.o_proj.o_proj.weight from torch.bfloat16 to torch.float +32 WARNING:Neuron:casting layers.13.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.13.mlp.up_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.13.mlp.down_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.13.input_layernorm.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.13.post_attention_layernorm.weight from torch.bfloat16 to torch.floa -t32 WARNING:Neuron:casting layers.14.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to torch.flo -at32 WARNING:Neuron:casting layers.14.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to torch.flo -at32 WARNING:Neuron:casting layers.14.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to torch.flo -at32 WARNING:Neuron:casting layers.14.self_attn.o_proj.o_proj.weight from torch.bfloat16 to torch.float -32 WARNING:Neuron:casting layers.14.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 +t32 WARNING:Neuron:casting layers.14.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to torch.flo +at32 WARNING:Neuron:casting layers.14.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to torch.flo +at32 WARNING:Neuron:casting layers.14.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to torch.flo +at32 WARNING:Neuron:casting layers.14.self_attn.o_proj.o_proj.weight from torch.bfloat16 to torch.float +32 WARNING:Neuron:casting layers.14.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.14.mlp.up_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.14.mlp.down_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.14.input_layernorm.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.14.post_attention_layernorm.weight from torch.bfloat16 to torch.floa -t32 WARNING:Neuron:casting layers.15.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to torch.flo -at32 WARNING:Neuron:casting layers.15.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to torch.flo -at32 WARNING:Neuron:casting layers.15.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to torch.flo -at32 WARNING:Neuron:casting layers.15.self_attn.o_proj.o_proj.weight from torch.bfloat16 to torch.float -32 WARNING:Neuron:casting layers.15.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 +t32 WARNING:Neuron:casting layers.15.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to torch.flo +at32 WARNING:Neuron:casting layers.15.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to torch.flo +at32 WARNING:Neuron:casting layers.15.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to torch.flo +at32 WARNING:Neuron:casting layers.15.self_attn.o_proj.o_proj.weight from torch.bfloat16 to torch.float +32 WARNING:Neuron:casting layers.15.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.15.mlp.up_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.15.mlp.down_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.15.input_layernorm.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.15.post_attention_layernorm.weight from torch.bfloat16 to torch.floa -t32 WARNING:Neuron:casting norm.weight from torch.bfloat16 to torch.float32 +t32 WARNING:Neuron:casting norm.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting lm_head.weight from torch.bfloat16 to torch.float32 INFO:Neuron:Done Sharding weights in 0.5333265441004187 INFO:Neuron:Finished weights loading in 12.090648926096037 seconds @@ -4409,7 +4382,6 @@ Created INFERENCE_TROUBLESHOOTING_GUIDE.md - I've created a comprehensive troubleshooting guide that documents all the errors we encountered and the fixes that were applied. The document covers: Model Initialization Error - The traced model not being properly initialized @@ -4457,276 +4429,275 @@ Created proper_inference.py - Command python proper_inference.py --model_path ./llama3_compiled --prompt "Hello, how are you?" --max_new_tokens 5 /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/parallel_layers/laye -rs.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +rs.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/parallel_layers/laye -rs.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +rs.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/parallel_layers/laye -rs.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +rs.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwis -e.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +e.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwis -e.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +e.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwis -e.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +e.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwis -e.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +e.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwis -e.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +e.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwis -e.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +e.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inference/m -odules/attention/utils.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.custom_calls import neuron_cumsum +odules/attention/utils.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.custom_calls import neuron_cumsum /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inference/m -odules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttentio -n_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inference/m -odules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttentio -n_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inference/m -odules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttentio -n_QKV /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_llama3/modeling_llama3.py:61: Deprec -ationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base import NeuronAttent -ionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_llama3/modeling_llama3.py:61: Deprec -ationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base import NeuronAttent -ionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_llama3/modeling_llama3.py:61: Deprec -ationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base import NeuronAttent -ionBase WARNING:root:NeuronConfig init: Unexpected keyword arguments: {'apply_seq_ids_mask': Fals -e, 'enable_long_context_mode': False, 'enable_output_completion_notifications': False, 'enable_token_tree': False, 'is_chunked_prefill': False, 'is_prefill_stage': None, 'kv_cache_tiling': False, 'scratchpad_page_size': None, 'skip_warmup': False, 'tile_cc': False, 'weights_to_skip_layout_optimization': []} INFO:Neuron:Sharding weights on load... +odules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttentio +n_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inference/m +odules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttentio +n_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inference/m +odules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttentio +n_QKV /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_llama3/modeling_llama3.py:61: Deprec +ationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base import NeuronAttent +ionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_llama3/modeling_llama3.py:61: Deprec +ationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base import NeuronAttent +ionBase /home/ec2-user/NeuronxSDK/neuronx_llama3/src/neuronx_llama3/modeling_llama3.py:61: Deprec +ationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.attention_base import NeuronAttent +ionBase WARNING:root:NeuronConfig init: Unexpected keyword arguments: {'apply_seq_ids_mask': Fals +e, 'enable_long_context_mode': False, 'enable_output_completion_notifications': False, 'enable_token_tree': False, 'is_chunked_prefill': False, 'is_prefill_stage': None, 'kv_cache_tiling': False, 'scratchpad_page_size': None, 'skip_warmup': False, 'tile_cc': False, 'weights_to_skip_layout_optimization': []} INFO:Neuron:Sharding weights on load... INFO:Neuron:Sharding Weights for ranks: 0...0 -[2025-07-26 17:20:07.427: I neuronx_distributed/parallel_layers/parallel_state.py:628] > -initializing tensor model parallel with size 1 [2025-07-26 17:20:07.427: I neuronx_distributed/parallel_layers/parallel_state.py:629] > -initializing pipeline model parallel with size 1 [2025-07-26 17:20:07.427: I neuronx_distributed/parallel_layers/parallel_state.py:630] > -initializing context model parallel with size 1 [2025-07-26 17:20:07.427: I neuronx_distributed/parallel_layers/parallel_state.py:631] > -initializing data parallel with size 1 [2025-07-26 17:20:07.427: I neuronx_distributed/parallel_layers/parallel_state.py:632] > -initializing world size to 1 [2025-07-26 17:20:07.428: I neuronx_distributed/parallel_layers/parallel_state.py:379] [r -ank_0_pp-1_tp-1_dp-1_cp-1] Chosen Logic for replica groups ret_logic=, 'Ascending Ring PG Group')> [2025-07-26 17:20:07.428: I neuronx_distributed/parallel_layers/parallel_state.py:668] [r -ank_0_pp-1_tp-1_dp-1_cp-1] tp_groups: replica_groups.tp_groups=[[0]] [2025-07-26 17:20:07.429: I neuronx_distributed/parallel_layers/parallel_state.py:669] [r -ank_0_pp-1_tp-1_dp-1_cp-1] dp_groups: replica_groups.dp_groups=[[0]] [2025-07-26 17:20:07.429: I neuronx_distributed/parallel_layers/parallel_state.py:670] [r -ank_0_pp-1_tp-1_dp-1_cp-1] pp_groups: replica_groups.pp_groups=[[0]] [2025-07-26 17:20:07.429: I neuronx_distributed/parallel_layers/parallel_state.py:671] [r -ank_0_pp-1_tp-1_dp-1_cp-1] cp_groups: replica_groups.cp_groups=[[0]] [2025-07-26 17:20:07.429: I neuronx_distributed/parallel_layers/parallel_state.py:672] [r -ank_0_pp-1_tp-1_dp-1_cp-1] ep_model_groups: replica_groups.ep_model_groups=[[0]] [2025-07-26 17:20:07.429: I neuronx_distributed/parallel_layers/parallel_state.py:673] [r -ank_0_pp-1_tp-1_dp-1_cp-1] ep_data_groups: replica_groups.ep_data_groups=[[0]] WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha -rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha -rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha -rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha -rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha -rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha -rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha -rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha -rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha -rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha -rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha -rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha -rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha -rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha -rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha -rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha -rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha -rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha -rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha -rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha -rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha -rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha -rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha -rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha -rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha -rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha -rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha -rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha -rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha -rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha -rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha -rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha -rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha -rding strategy to GQA.CONVERT_TO_MHA! /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/trace/trace.py:640: -UserWarning: Removing redundant keys from checkpoint: ['layers.0.self_attn.q_proj.weight', 'layers.0.self_attn.k_proj.weight', 'layers.0.self_attn.v_proj.weight', 'layers.0.self_attn.o_proj.weight', 'layers.1.self_attn.q_proj.weight', 'layers.1.self_attn.k_proj.weight', 'layers.1.self_attn.v_proj.weight', 'layers.1.self_attn.o_proj.weight', 'layers.2.self_attn.q_proj.weight', 'layers.2.self_attn.k_proj.weight', 'layers.2.self_attn.v_proj.weight', 'layers.2.self_attn.o_proj.weight', 'layers.3.self_attn.q_proj.weight', 'layers.3.self_attn.k_proj.weight', 'layers.3.self_attn.v_proj.weight', 'layers.3.self_attn.o_proj.weight', 'layers.4.self_attn.q_proj.weight', 'layers.4.self_attn.k_proj.weight', 'layers.4.self_attn.v_proj.weight', 'layers.4.self_attn.o_proj.weight', 'layers.5.self_attn.q_proj.weight', 'layers.5.self_attn.k_proj.weight', 'layers.5.self_attn.v_proj.weight', 'layers.5.self_attn.o_proj.weight', 'layers.6.self_attn.q_proj.weight', 'layers.6.self_attn.k_proj.weight', 'layers.6.self_attn.v_proj.weight', 'layers.6.self_attn.o_proj.weight', 'layers.7.self_attn.q_proj.weight', 'layers.7.self_attn.k_proj.weight', 'layers.7.self_attn.v_proj.weight', 'layers.7.self_attn.o_proj.weight', 'layers.8.self_attn.q_proj.weight', 'layers.8.self_attn.k_proj.weight', 'layers.8.self_attn.v_proj.weight', 'layers.8.self_attn.o_proj.weight', 'layers.9.self_attn.q_proj.weight', 'layers.9.self_attn.k_proj.weight', 'layers.9.self_attn.v_proj.weight', 'layers.9.self_attn.o_proj.weight', 'layers.10.self_attn.q_proj.weight', 'layers.10.self_attn.k_proj.weight', 'layers.10.self_attn.v_proj.weight', 'layers.10.self_attn.o_proj.weight', 'layers.11.self_attn.q_proj.weight', 'layers.11.self_attn.k_proj.weight', 'layers.11.self_attn.v_proj.weight', 'layers.11.self_attn.o_proj.weight', 'layers.12.self_attn.q_proj.weight', 'layers.12.self_attn.k_proj.weight', 'layers.12.self_attn.v_proj.weight', 'layers.12.self_attn.o_proj.weight', 'layers.13.self_attn.q_proj.weight', 'layers.13.self_attn.k_proj.weight', 'layers.13.self_attn.v_proj.weight', 'layers.13.self_attn.o_proj.weight', 'layers.14.self_attn.q_proj.weight', 'layers.14.self_attn.k_proj.weight', 'layers.14.self_attn.v_proj.weight', 'layers.14.self_attn.o_proj.weight', 'layers.15.self_attn.q_proj.weight', 'layers.15.self_attn.k_proj.weight', 'layers.15.self_attn.v_proj.weight', 'layers.15.self_attn.o_proj.weight'] warnings.warn(f"Removing redundant keys from checkpoint: {keys_to_delete}") +[2025-07-26 17:20:07.427: I neuronx_distributed/parallel_layers/parallel_state.py:628] > +initializing tensor model parallel with size 1 [2025-07-26 17:20:07.427: I neuronx_distributed/parallel_layers/parallel_state.py:629] > +initializing pipeline model parallel with size 1 [2025-07-26 17:20:07.427: I neuronx_distributed/parallel_layers/parallel_state.py:630] > +initializing context model parallel with size 1 [2025-07-26 17:20:07.427: I neuronx_distributed/parallel_layers/parallel_state.py:631] > +initializing data parallel with size 1 [2025-07-26 17:20:07.427: I neuronx_distributed/parallel_layers/parallel_state.py:632] > +initializing world size to 1 [2025-07-26 17:20:07.428: I neuronx_distributed/parallel_layers/parallel_state.py:379] [r +ank_0_pp-1_tp-1_dp-1_cp-1] Chosen Logic for replica groups ret_logic=, 'Ascending Ring PG Group')> [2025-07-26 17:20:07.428: I neuronx_distributed/parallel_layers/parallel_state.py:668] [r +ank_0_pp-1_tp-1_dp-1_cp-1] tp_groups: replica_groups.tp_groups=[[0]] [2025-07-26 17:20:07.429: I neuronx_distributed/parallel_layers/parallel_state.py:669] [r +ank_0_pp-1_tp-1_dp-1_cp-1] dp_groups: replica_groups.dp_groups=[[0]] [2025-07-26 17:20:07.429: I neuronx_distributed/parallel_layers/parallel_state.py:670] [r +ank_0_pp-1_tp-1_dp-1_cp-1] pp_groups: replica_groups.pp_groups=[[0]] [2025-07-26 17:20:07.429: I neuronx_distributed/parallel_layers/parallel_state.py:671] [r +ank_0_pp-1_tp-1_dp-1_cp-1] cp_groups: replica_groups.cp_groups=[[0]] [2025-07-26 17:20:07.429: I neuronx_distributed/parallel_layers/parallel_state.py:672] [r +ank_0_pp-1_tp-1_dp-1_cp-1] ep_model_groups: replica_groups.ep_model_groups=[[0]] [2025-07-26 17:20:07.429: I neuronx_distributed/parallel_layers/parallel_state.py:673] [r +ank_0_pp-1_tp-1_dp-1_cp-1] ep_data_groups: replica_groups.ep_data_groups=[[0]] WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha +rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha +rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha +rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha +rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha +rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha +rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha +rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha +rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha +rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha +rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha +rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha +rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha +rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha +rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha +rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha +rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha +rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha +rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha +rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha +rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha +rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha +rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha +rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha +rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha +rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha +rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha +rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha +rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha +rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha +rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha +rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha +rding strategy to GQA.CONVERT_TO_MHA! WARNING:Neuron:TP degree (1) and KV heads (8) are not divisible. Overriding attention sha +rding strategy to GQA.CONVERT_TO_MHA! /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/trace/trace.py:640: +UserWarning: Removing redundant keys from checkpoint: ['layers.0.self_attn.q_proj.weight', 'layers.0.self_attn.k_proj.weight', 'layers.0.self_attn.v_proj.weight', 'layers.0.self_attn.o_proj.weight', 'layers.1.self_attn.q_proj.weight', 'layers.1.self_attn.k_proj.weight', 'layers.1.self_attn.v_proj.weight', 'layers.1.self_attn.o_proj.weight', 'layers.2.self_attn.q_proj.weight', 'layers.2.self_attn.k_proj.weight', 'layers.2.self_attn.v_proj.weight', 'layers.2.self_attn.o_proj.weight', 'layers.3.self_attn.q_proj.weight', 'layers.3.self_attn.k_proj.weight', 'layers.3.self_attn.v_proj.weight', 'layers.3.self_attn.o_proj.weight', 'layers.4.self_attn.q_proj.weight', 'layers.4.self_attn.k_proj.weight', 'layers.4.self_attn.v_proj.weight', 'layers.4.self_attn.o_proj.weight', 'layers.5.self_attn.q_proj.weight', 'layers.5.self_attn.k_proj.weight', 'layers.5.self_attn.v_proj.weight', 'layers.5.self_attn.o_proj.weight', 'layers.6.self_attn.q_proj.weight', 'layers.6.self_attn.k_proj.weight', 'layers.6.self_attn.v_proj.weight', 'layers.6.self_attn.o_proj.weight', 'layers.7.self_attn.q_proj.weight', 'layers.7.self_attn.k_proj.weight', 'layers.7.self_attn.v_proj.weight', 'layers.7.self_attn.o_proj.weight', 'layers.8.self_attn.q_proj.weight', 'layers.8.self_attn.k_proj.weight', 'layers.8.self_attn.v_proj.weight', 'layers.8.self_attn.o_proj.weight', 'layers.9.self_attn.q_proj.weight', 'layers.9.self_attn.k_proj.weight', 'layers.9.self_attn.v_proj.weight', 'layers.9.self_attn.o_proj.weight', 'layers.10.self_attn.q_proj.weight', 'layers.10.self_attn.k_proj.weight', 'layers.10.self_attn.v_proj.weight', 'layers.10.self_attn.o_proj.weight', 'layers.11.self_attn.q_proj.weight', 'layers.11.self_attn.k_proj.weight', 'layers.11.self_attn.v_proj.weight', 'layers.11.self_attn.o_proj.weight', 'layers.12.self_attn.q_proj.weight', 'layers.12.self_attn.k_proj.weight', 'layers.12.self_attn.v_proj.weight', 'layers.12.self_attn.o_proj.weight', 'layers.13.self_attn.q_proj.weight', 'layers.13.self_attn.k_proj.weight', 'layers.13.self_attn.v_proj.weight', 'layers.13.self_attn.o_proj.weight', 'layers.14.self_attn.q_proj.weight', 'layers.14.self_attn.k_proj.weight', 'layers.14.self_attn.v_proj.weight', 'layers.14.self_attn.o_proj.weight', 'layers.15.self_attn.q_proj.weight', 'layers.15.self_attn.k_proj.weight', 'layers.15.self_attn.v_proj.weight', 'layers.15.self_attn.o_proj.weight'] warnings.warn(f"Removing redundant keys from checkpoint: {keys_to_delete}") WARNING:Neuron:casting embed_tokens.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.0.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to t -orch.float32 WARNING:Neuron:casting layers.0.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to t -orch.float32 WARNING:Neuron:casting layers.0.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to t -orch.float32 WARNING:Neuron:casting layers.0.self_attn.o_proj.o_proj.weight from torch.bfloat16 to tor -ch.float32 WARNING:Neuron:casting layers.0.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 +orch.float32 WARNING:Neuron:casting layers.0.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to t +orch.float32 WARNING:Neuron:casting layers.0.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to t +orch.float32 WARNING:Neuron:casting layers.0.self_attn.o_proj.o_proj.weight from torch.bfloat16 to tor +ch.float32 WARNING:Neuron:casting layers.0.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.0.mlp.up_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.0.mlp.down_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.0.input_layernorm.weight from torch.bfloat16 to torch.float -32 WARNING:Neuron:casting layers.0.post_attention_layernorm.weight from torch.bfloat16 to to -rch.float32 WARNING:Neuron:casting layers.1.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to t -orch.float32 WARNING:Neuron:casting layers.1.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to t -orch.float32 WARNING:Neuron:casting layers.1.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to t -orch.float32 WARNING:Neuron:casting layers.1.self_attn.o_proj.o_proj.weight from torch.bfloat16 to tor -ch.float32 WARNING:Neuron:casting layers.1.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 +32 WARNING:Neuron:casting layers.0.post_attention_layernorm.weight from torch.bfloat16 to to +rch.float32 WARNING:Neuron:casting layers.1.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to t +orch.float32 WARNING:Neuron:casting layers.1.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to t +orch.float32 WARNING:Neuron:casting layers.1.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to t +orch.float32 WARNING:Neuron:casting layers.1.self_attn.o_proj.o_proj.weight from torch.bfloat16 to tor +ch.float32 WARNING:Neuron:casting layers.1.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.1.mlp.up_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.1.mlp.down_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.1.input_layernorm.weight from torch.bfloat16 to torch.float -32 WARNING:Neuron:casting layers.1.post_attention_layernorm.weight from torch.bfloat16 to to -rch.float32 WARNING:Neuron:casting layers.2.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to t -orch.float32 WARNING:Neuron:casting layers.2.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to t -orch.float32 WARNING:Neuron:casting layers.2.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to t -orch.float32 WARNING:Neuron:casting layers.2.self_attn.o_proj.o_proj.weight from torch.bfloat16 to tor -ch.float32 WARNING:Neuron:casting layers.2.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 +32 WARNING:Neuron:casting layers.1.post_attention_layernorm.weight from torch.bfloat16 to to +rch.float32 WARNING:Neuron:casting layers.2.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to t +orch.float32 WARNING:Neuron:casting layers.2.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to t +orch.float32 WARNING:Neuron:casting layers.2.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to t +orch.float32 WARNING:Neuron:casting layers.2.self_attn.o_proj.o_proj.weight from torch.bfloat16 to tor +ch.float32 WARNING:Neuron:casting layers.2.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.2.mlp.up_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.2.mlp.down_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.2.input_layernorm.weight from torch.bfloat16 to torch.float -32 WARNING:Neuron:casting layers.2.post_attention_layernorm.weight from torch.bfloat16 to to -rch.float32 WARNING:Neuron:casting layers.3.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to t -orch.float32 WARNING:Neuron:casting layers.3.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to t -orch.float32 WARNING:Neuron:casting layers.3.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to t -orch.float32 WARNING:Neuron:casting layers.3.self_attn.o_proj.o_proj.weight from torch.bfloat16 to tor -ch.float32 WARNING:Neuron:casting layers.3.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 +32 WARNING:Neuron:casting layers.2.post_attention_layernorm.weight from torch.bfloat16 to to +rch.float32 WARNING:Neuron:casting layers.3.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to t +orch.float32 WARNING:Neuron:casting layers.3.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to t +orch.float32 WARNING:Neuron:casting layers.3.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to t +orch.float32 WARNING:Neuron:casting layers.3.self_attn.o_proj.o_proj.weight from torch.bfloat16 to tor +ch.float32 WARNING:Neuron:casting layers.3.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.3.mlp.up_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.3.mlp.down_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.3.input_layernorm.weight from torch.bfloat16 to torch.float -32 WARNING:Neuron:casting layers.3.post_attention_layernorm.weight from torch.bfloat16 to to -rch.float32 WARNING:Neuron:casting layers.4.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to t -orch.float32 WARNING:Neuron:casting layers.4.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to t -orch.float32 WARNING:Neuron:casting layers.4.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to t -orch.float32 WARNING:Neuron:casting layers.4.self_attn.o_proj.o_proj.weight from torch.bfloat16 to tor -ch.float32 WARNING:Neuron:casting layers.4.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 +32 WARNING:Neuron:casting layers.3.post_attention_layernorm.weight from torch.bfloat16 to to +rch.float32 WARNING:Neuron:casting layers.4.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to t +orch.float32 WARNING:Neuron:casting layers.4.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to t +orch.float32 WARNING:Neuron:casting layers.4.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to t +orch.float32 WARNING:Neuron:casting layers.4.self_attn.o_proj.o_proj.weight from torch.bfloat16 to tor +ch.float32 WARNING:Neuron:casting layers.4.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.4.mlp.up_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.4.mlp.down_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.4.input_layernorm.weight from torch.bfloat16 to torch.float -32 WARNING:Neuron:casting layers.4.post_attention_layernorm.weight from torch.bfloat16 to to -rch.float32 WARNING:Neuron:casting layers.5.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to t -orch.float32 WARNING:Neuron:casting layers.5.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to t -orch.float32 WARNING:Neuron:casting layers.5.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to t -orch.float32 WARNING:Neuron:casting layers.5.self_attn.o_proj.o_proj.weight from torch.bfloat16 to tor -ch.float32 WARNING:Neuron:casting layers.5.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 +32 WARNING:Neuron:casting layers.4.post_attention_layernorm.weight from torch.bfloat16 to to +rch.float32 WARNING:Neuron:casting layers.5.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to t +orch.float32 WARNING:Neuron:casting layers.5.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to t +orch.float32 WARNING:Neuron:casting layers.5.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to t +orch.float32 WARNING:Neuron:casting layers.5.self_attn.o_proj.o_proj.weight from torch.bfloat16 to tor +ch.float32 WARNING:Neuron:casting layers.5.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.5.mlp.up_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.5.mlp.down_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.5.input_layernorm.weight from torch.bfloat16 to torch.float -32 WARNING:Neuron:casting layers.5.post_attention_layernorm.weight from torch.bfloat16 to to -rch.float32 WARNING:Neuron:casting layers.6.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to t -orch.float32 WARNING:Neuron:casting layers.6.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to t -orch.float32 WARNING:Neuron:casting layers.6.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to t -orch.float32 WARNING:Neuron:casting layers.6.self_attn.o_proj.o_proj.weight from torch.bfloat16 to tor -ch.float32 WARNING:Neuron:casting layers.6.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 +32 WARNING:Neuron:casting layers.5.post_attention_layernorm.weight from torch.bfloat16 to to +rch.float32 WARNING:Neuron:casting layers.6.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to t +orch.float32 WARNING:Neuron:casting layers.6.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to t +orch.float32 WARNING:Neuron:casting layers.6.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to t +orch.float32 WARNING:Neuron:casting layers.6.self_attn.o_proj.o_proj.weight from torch.bfloat16 to tor +ch.float32 WARNING:Neuron:casting layers.6.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.6.mlp.up_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.6.mlp.down_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.6.input_layernorm.weight from torch.bfloat16 to torch.float -32 WARNING:Neuron:casting layers.6.post_attention_layernorm.weight from torch.bfloat16 to to -rch.float32 WARNING:Neuron:casting layers.7.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to t -orch.float32 WARNING:Neuron:casting layers.7.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to t -orch.float32 WARNING:Neuron:casting layers.7.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to t -orch.float32 WARNING:Neuron:casting layers.7.self_attn.o_proj.o_proj.weight from torch.bfloat16 to tor -ch.float32 WARNING:Neuron:casting layers.7.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 +32 WARNING:Neuron:casting layers.6.post_attention_layernorm.weight from torch.bfloat16 to to +rch.float32 WARNING:Neuron:casting layers.7.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to t +orch.float32 WARNING:Neuron:casting layers.7.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to t +orch.float32 WARNING:Neuron:casting layers.7.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to t +orch.float32 WARNING:Neuron:casting layers.7.self_attn.o_proj.o_proj.weight from torch.bfloat16 to tor +ch.float32 WARNING:Neuron:casting layers.7.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.7.mlp.up_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.7.mlp.down_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.7.input_layernorm.weight from torch.bfloat16 to torch.float -32 WARNING:Neuron:casting layers.7.post_attention_layernorm.weight from torch.bfloat16 to to -rch.float32 WARNING:Neuron:casting layers.8.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to t -orch.float32 WARNING:Neuron:casting layers.8.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to t -orch.float32 WARNING:Neuron:casting layers.8.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to t -orch.float32 WARNING:Neuron:casting layers.8.self_attn.o_proj.o_proj.weight from torch.bfloat16 to tor -ch.float32 WARNING:Neuron:casting layers.8.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 +32 WARNING:Neuron:casting layers.7.post_attention_layernorm.weight from torch.bfloat16 to to +rch.float32 WARNING:Neuron:casting layers.8.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to t +orch.float32 WARNING:Neuron:casting layers.8.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to t +orch.float32 WARNING:Neuron:casting layers.8.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to t +orch.float32 WARNING:Neuron:casting layers.8.self_attn.o_proj.o_proj.weight from torch.bfloat16 to tor +ch.float32 WARNING:Neuron:casting layers.8.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.8.mlp.up_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.8.mlp.down_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.8.input_layernorm.weight from torch.bfloat16 to torch.float -32 WARNING:Neuron:casting layers.8.post_attention_layernorm.weight from torch.bfloat16 to to -rch.float32 WARNING:Neuron:casting layers.9.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to t -orch.float32 WARNING:Neuron:casting layers.9.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to t -orch.float32 WARNING:Neuron:casting layers.9.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to t -orch.float32 WARNING:Neuron:casting layers.9.self_attn.o_proj.o_proj.weight from torch.bfloat16 to tor -ch.float32 WARNING:Neuron:casting layers.9.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 +32 WARNING:Neuron:casting layers.8.post_attention_layernorm.weight from torch.bfloat16 to to +rch.float32 WARNING:Neuron:casting layers.9.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to t +orch.float32 WARNING:Neuron:casting layers.9.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to t +orch.float32 WARNING:Neuron:casting layers.9.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to t +orch.float32 WARNING:Neuron:casting layers.9.self_attn.o_proj.o_proj.weight from torch.bfloat16 to tor +ch.float32 WARNING:Neuron:casting layers.9.mlp.gate_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.9.mlp.up_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.9.mlp.down_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.9.input_layernorm.weight from torch.bfloat16 to torch.float -32 WARNING:Neuron:casting layers.9.post_attention_layernorm.weight from torch.bfloat16 to to -rch.float32 WARNING:Neuron:casting layers.10.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to -torch.float32 WARNING:Neuron:casting layers.10.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to -torch.float32 WARNING:Neuron:casting layers.10.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to -torch.float32 WARNING:Neuron:casting layers.10.self_attn.o_proj.o_proj.weight from torch.bfloat16 to to -rch.float32 WARNING:Neuron:casting layers.10.mlp.gate_proj.weight from torch.bfloat16 to torch.float3 -2 WARNING:Neuron:casting layers.10.mlp.up_proj.weight from torch.bfloat16 to torch.float32 +32 WARNING:Neuron:casting layers.9.post_attention_layernorm.weight from torch.bfloat16 to to +rch.float32 WARNING:Neuron:casting layers.10.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to +torch.float32 WARNING:Neuron:casting layers.10.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to +torch.float32 WARNING:Neuron:casting layers.10.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to +torch.float32 WARNING:Neuron:casting layers.10.self_attn.o_proj.o_proj.weight from torch.bfloat16 to to +rch.float32 WARNING:Neuron:casting layers.10.mlp.gate_proj.weight from torch.bfloat16 to torch.float3 +2 WARNING:Neuron:casting layers.10.mlp.up_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.10.mlp.down_proj.weight from torch.bfloat16 to torch.float3 -2 WARNING:Neuron:casting layers.10.input_layernorm.weight from torch.bfloat16 to torch.floa -t32 WARNING:Neuron:casting layers.10.post_attention_layernorm.weight from torch.bfloat16 to t -orch.float32 WARNING:Neuron:casting layers.11.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to -torch.float32 WARNING:Neuron:casting layers.11.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to -torch.float32 WARNING:Neuron:casting layers.11.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to -torch.float32 WARNING:Neuron:casting layers.11.self_attn.o_proj.o_proj.weight from torch.bfloat16 to to -rch.float32 WARNING:Neuron:casting layers.11.mlp.gate_proj.weight from torch.bfloat16 to torch.float3 -2 WARNING:Neuron:casting layers.11.mlp.up_proj.weight from torch.bfloat16 to torch.float32 +2 WARNING:Neuron:casting layers.10.input_layernorm.weight from torch.bfloat16 to torch.floa +t32 WARNING:Neuron:casting layers.10.post_attention_layernorm.weight from torch.bfloat16 to t +orch.float32 WARNING:Neuron:casting layers.11.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to +torch.float32 WARNING:Neuron:casting layers.11.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to +torch.float32 WARNING:Neuron:casting layers.11.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to +torch.float32 WARNING:Neuron:casting layers.11.self_attn.o_proj.o_proj.weight from torch.bfloat16 to to +rch.float32 WARNING:Neuron:casting layers.11.mlp.gate_proj.weight from torch.bfloat16 to torch.float3 +2 WARNING:Neuron:casting layers.11.mlp.up_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.11.mlp.down_proj.weight from torch.bfloat16 to torch.float3 -2 WARNING:Neuron:casting layers.11.input_layernorm.weight from torch.bfloat16 to torch.floa -t32 WARNING:Neuron:casting layers.11.post_attention_layernorm.weight from torch.bfloat16 to t -orch.float32 WARNING:Neuron:casting layers.12.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to -torch.float32 WARNING:Neuron:casting layers.12.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to -torch.float32 WARNING:Neuron:casting layers.12.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to -torch.float32 WARNING:Neuron:casting layers.12.self_attn.o_proj.o_proj.weight from torch.bfloat16 to to -rch.float32 WARNING:Neuron:casting layers.12.mlp.gate_proj.weight from torch.bfloat16 to torch.float3 -2 WARNING:Neuron:casting layers.12.mlp.up_proj.weight from torch.bfloat16 to torch.float32 +2 WARNING:Neuron:casting layers.11.input_layernorm.weight from torch.bfloat16 to torch.floa +t32 WARNING:Neuron:casting layers.11.post_attention_layernorm.weight from torch.bfloat16 to t +orch.float32 WARNING:Neuron:casting layers.12.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to +torch.float32 WARNING:Neuron:casting layers.12.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to +torch.float32 WARNING:Neuron:casting layers.12.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to +torch.float32 WARNING:Neuron:casting layers.12.self_attn.o_proj.o_proj.weight from torch.bfloat16 to to +rch.float32 WARNING:Neuron:casting layers.12.mlp.gate_proj.weight from torch.bfloat16 to torch.float3 +2 WARNING:Neuron:casting layers.12.mlp.up_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.12.mlp.down_proj.weight from torch.bfloat16 to torch.float3 -2 WARNING:Neuron:casting layers.12.input_layernorm.weight from torch.bfloat16 to torch.floa -t32 WARNING:Neuron:casting layers.12.post_attention_layernorm.weight from torch.bfloat16 to t -orch.float32 WARNING:Neuron:casting layers.13.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to -torch.float32 WARNING:Neuron:casting layers.13.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to -torch.float32 WARNING:Neuron:casting layers.13.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to -torch.float32 WARNING:Neuron:casting layers.13.self_attn.o_proj.o_proj.weight from torch.bfloat16 to to -rch.float32 WARNING:Neuron:casting layers.13.mlp.gate_proj.weight from torch.bfloat16 to torch.float3 -2 WARNING:Neuron:casting layers.13.mlp.up_proj.weight from torch.bfloat16 to torch.float32 +2 WARNING:Neuron:casting layers.12.input_layernorm.weight from torch.bfloat16 to torch.floa +t32 WARNING:Neuron:casting layers.12.post_attention_layernorm.weight from torch.bfloat16 to t +orch.float32 WARNING:Neuron:casting layers.13.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to +torch.float32 WARNING:Neuron:casting layers.13.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to +torch.float32 WARNING:Neuron:casting layers.13.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to +torch.float32 WARNING:Neuron:casting layers.13.self_attn.o_proj.o_proj.weight from torch.bfloat16 to to +rch.float32 WARNING:Neuron:casting layers.13.mlp.gate_proj.weight from torch.bfloat16 to torch.float3 +2 WARNING:Neuron:casting layers.13.mlp.up_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.13.mlp.down_proj.weight from torch.bfloat16 to torch.float3 -2 WARNING:Neuron:casting layers.13.input_layernorm.weight from torch.bfloat16 to torch.floa -t32 WARNING:Neuron:casting layers.13.post_attention_layernorm.weight from torch.bfloat16 to t -orch.float32 WARNING:Neuron:casting layers.14.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to -torch.float32 WARNING:Neuron:casting layers.14.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to -torch.float32 WARNING:Neuron:casting layers.14.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to -torch.float32 WARNING:Neuron:casting layers.14.self_attn.o_proj.o_proj.weight from torch.bfloat16 to to -rch.float32 WARNING:Neuron:casting layers.14.mlp.gate_proj.weight from torch.bfloat16 to torch.float3 -2 WARNING:Neuron:casting layers.14.mlp.up_proj.weight from torch.bfloat16 to torch.float32 +2 WARNING:Neuron:casting layers.13.input_layernorm.weight from torch.bfloat16 to torch.floa +t32 WARNING:Neuron:casting layers.13.post_attention_layernorm.weight from torch.bfloat16 to t +orch.float32 WARNING:Neuron:casting layers.14.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to +torch.float32 WARNING:Neuron:casting layers.14.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to +torch.float32 WARNING:Neuron:casting layers.14.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to +torch.float32 WARNING:Neuron:casting layers.14.self_attn.o_proj.o_proj.weight from torch.bfloat16 to to +rch.float32 WARNING:Neuron:casting layers.14.mlp.gate_proj.weight from torch.bfloat16 to torch.float3 +2 WARNING:Neuron:casting layers.14.mlp.up_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.14.mlp.down_proj.weight from torch.bfloat16 to torch.float3 -2 WARNING:Neuron:casting layers.14.input_layernorm.weight from torch.bfloat16 to torch.floa -t32 WARNING:Neuron:casting layers.14.post_attention_layernorm.weight from torch.bfloat16 to t -orch.float32 WARNING:Neuron:casting layers.15.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to -torch.float32 WARNING:Neuron:casting layers.15.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to -torch.float32 WARNING:Neuron:casting layers.15.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to -torch.float32 WARNING:Neuron:casting layers.15.self_attn.o_proj.o_proj.weight from torch.bfloat16 to to -rch.float32 WARNING:Neuron:casting layers.15.mlp.gate_proj.weight from torch.bfloat16 to torch.float3 -2 WARNING:Neuron:casting layers.15.mlp.up_proj.weight from torch.bfloat16 to torch.float32 +2 WARNING:Neuron:casting layers.14.input_layernorm.weight from torch.bfloat16 to torch.floa +t32 WARNING:Neuron:casting layers.14.post_attention_layernorm.weight from torch.bfloat16 to t +orch.float32 WARNING:Neuron:casting layers.15.self_attn.qkv_proj.q_proj.weight from torch.bfloat16 to +torch.float32 WARNING:Neuron:casting layers.15.self_attn.qkv_proj.k_proj.weight from torch.bfloat16 to +torch.float32 WARNING:Neuron:casting layers.15.self_attn.qkv_proj.v_proj.weight from torch.bfloat16 to +torch.float32 WARNING:Neuron:casting layers.15.self_attn.o_proj.o_proj.weight from torch.bfloat16 to to +rch.float32 WARNING:Neuron:casting layers.15.mlp.gate_proj.weight from torch.bfloat16 to torch.float3 +2 WARNING:Neuron:casting layers.15.mlp.up_proj.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting layers.15.mlp.down_proj.weight from torch.bfloat16 to torch.float3 -2 WARNING:Neuron:casting layers.15.input_layernorm.weight from torch.bfloat16 to torch.floa -t32 WARNING:Neuron:casting layers.15.post_attention_layernorm.weight from torch.bfloat16 to t -orch.float32 WARNING:Neuron:casting norm.weight from torch.bfloat16 to torch.float32 +2 WARNING:Neuron:casting layers.15.input_layernorm.weight from torch.bfloat16 to torch.floa +t32 WARNING:Neuron:casting layers.15.post_attention_layernorm.weight from torch.bfloat16 to t +orch.float32 WARNING:Neuron:casting norm.weight from torch.bfloat16 to torch.float32 WARNING:Neuron:casting lm_head.weight from torch.bfloat16 to torch.float32 INFO:Neuron:Done Sharding weights in 0.5459382231347263 INFO:Neuron:Finished weights loading in 12.248568888986483 seconds INFO:Neuron:Warming up the model. INFO:Neuron:Warmup completed in 0.05255460739135742 seconds. Can't set use_return_dict with value True for PretrainedConfig { - "fused_spec_config": null, - "hidden_act": "silu", - "hidden_size": 2048, - "intermediate_size": 8192, - "max_position_embeddings": 2048, - "metadata": null, - "num_attention_heads": 32, - "num_cores_per_group": 1, - "num_hidden_layers": 16, - "num_key_value_heads": 8, - "rms_norm_eps": 1e-05, - "rope_theta": 500000.0, - "tie_word_embeddings": false, - "torch_dtype": "float32", - "transformers_version": "4.51.3", - "use_cache": true +"fused_spec_config": null, +"hidden_act": "silu", +"hidden_size": 2048, +"intermediate_size": 8192, +"max_position_embeddings": 2048, +"metadata": null, +"num_attention_heads": 32, +"num_cores_per_group": 1, +"num_hidden_layers": 16, +"num_key_value_heads": 8, +"rms_norm_eps": 1e-05, +"rope_theta": 500000.0, +"tie_word_embeddings": false, +"torch_dtype": "float32", +"transformers_version": "4.51.3", +"use_cache": true } Traceback (most recent call last): - File "/home/ec2-user/NeuronxSDK/neuronx_llama3/proper_inference.py", line 87, in main() - File "/home/ec2-user/NeuronxSDK/neuronx_llama3/proper_inference.py", line 56, in main - generation_model = HuggingFaceGenerationAdapter(model) - File "/home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inf -erence/utils/hf_adapter.py", line 101, in __init__ hf_config = to_pretrained_config(model.config) - File "/home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inf -erence/utils/hf_adapter.py", line 96, in to_pretrained_config return PretrainedConfig(**config_dict) - File "/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages/tran -sformers/configuration_utils.py", line 323, in __init__ raise err - File "/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages/tran -sformers/configuration_utils.py", line 320, in __init__ setattr(self, key, value) - File "/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages/tran -sformers/configuration_utils.py", line 205, in __setattr__ super().__setattr__(key, value) +File "/home/ec2-user/NeuronxSDK/neuronx_llama3/proper_inference.py", line 87, in main() +File "/home/ec2-user/NeuronxSDK/neuronx_llama3/proper_inference.py", line 56, in main +generation_model = HuggingFaceGenerationAdapter(model) +File "/home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inf +erence/utils/hf_adapter.py", line 101, in **init** hf_config = to_pretrained_config(model.config) +File "/home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inf +erence/utils/hf_adapter.py", line 96, in to_pretrained_config return PretrainedConfig(\*\*config_dict) +File "/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages/tran +sformers/configuration_utils.py", line 323, in **init** raise err +File "/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages/tran +sformers/configuration_utils.py", line 320, in **init** setattr(self, key, value) +File "/opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/lib/python3.10/site-packages/tran +sformers/configuration_utils.py", line 205, in **setattr** super().**setattr**(key, value) AttributeError: can't set attribute 'use_return_dict' Now I see the exact same error we had before. Let me check what the working examples do differently. Looking at the examples, I notice they use the standard NeuronLlamaForCausalLM from the framework, not our custom NeuronLlama3ForCausalLM. Let me check what the difference is. @@ -4744,39 +4715,38 @@ Created test_standard_model.py - Command python test_standard_model.py --model_path ./llama3_compiled --prompt "Hello, how are you?" --max_new_tokens 5 /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/parallel_layers/laye -rs.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +rs.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/parallel_layers/laye -rs.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +rs.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/parallel_layers/laye -rs.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( +rs.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from .mappings import ( /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwis -e.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +e.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwis -e.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +e.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwis -e.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +e.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwis -e.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +e.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwis -e.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +e.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributed/src/neuronx_distributed/modules/moe/blockwis -e.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) +e.py:42: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. component, error = import_nki(config) /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inference/m -odules/attention/utils.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.custom_calls import neuron_cumsum +odules/attention/utils.py:14: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.custom_calls import neuron_cumsum /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inference/m -odules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttentio -n_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inference/m -odules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttentio -n_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inference/m -odules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttentio -n_QKV Traceback (most recent call last): - File "/home/ec2-user/NeuronxSDK/neuronx_llama3/test_standard_model.py", line 15, in from neuronx_distributed_inference.models.llama.modeling_llama import NeuronLlamaForC -ausalLM ModuleNotFoundError: No module named 'neuronx_distributed_inference.models.llama' +odules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttentio +n_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inference/m +odules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttentio +n_QKV /home/ec2-user/NeuronxSDK/NeuronxDistributedInference/src/neuronx_distributed_inference/m +odules/lora_serving/lora_model.py:12: DeprecationWarning: torch_neuronx.nki_jit is deprecated, use nki.jit instead. from neuronx_distributed_inference.modules.attention.gqa import GQA, GroupQueryAttentio +n_QKV Traceback (most recent call last): +File "/home/ec2-user/NeuronxSDK/neuronx_llama3/test_standard_model.py", line 15, in from neuronx_distributed_inference.models.llama.modeling_llama import NeuronLlamaForC +ausalLM ModuleNotFoundError: No module named 'neuronx_distributed_inference.models.llama' given the simple_inference.py script and configurations and weights and tensor parallelism and sharding can you create an example that uses huggingface transformers and the hugging face llama3 model and running on cpu Kiro @@ -4786,18 +4756,17 @@ Created cpu_inference_example.py - Now let's test this CPU-based example: Command python cpu_inference_example.py --prompt "Hello, how are you?" --max_new_tokens 5 --do_sample -ERROR:__main__:Failed to load tokenizer: You are trying to access a gated repo. +ERROR:**main**:Failed to load tokenizer: You are trying to access a gated repo. Make sure to have access to it at https://huggingface.co/meta-llama/Llama-3.2-1B. 401 Client Error. (Request ID: Root=1-68851487-32eb741f430c6d536ec54c6b;16def46c-6af6-4ab -8-a436-0ad9c4e4b13f) +8-a436-0ad9c4e4b13f) Cannot access gated repo for url https://huggingface.co/meta-llama/Llama-3.2-1B/resolve/m -ain/config.json. Access to model meta-llama/Llama-3.2-1B is restricted. You must have access to it and be -authenticated to access it. Please log in. tokenizer_config.json: 100%|████████████████████████████| 614/614 [00:00<00:00, 6.90MB/s] +ain/config.json. Access to model meta-llama/Llama-3.2-1B is restricted. You must have access to it and be +authenticated to access it. Please log in. tokenizer_config.json: 100%|████████████████████████████| 614/614 [00:00<00:00, 6.90MB/s] vocab.json: 1.04MB [00:00, 131MB/s] merges.txt: 456kB [00:00, 147MB/s] config.json: 100%|██████████████████████████████████████| 641/641 [00:00<00:00, 7.98MB/s] @@ -4811,10 +4780,10 @@ Model: microsoft/DialoGPT-small Generation settings: max_tokens=5, temperature=0.6, top_p=0.9, do_sample=True ============================================================ -Prompt: Hello, how are you? ----------------------------------------- -Generated: Hello, how are you? Good morning everyone! ----------------------------------------- +## Prompt: Hello, how are you? + +## Generated: Hello, how are you? Good morning everyone! + Generation time: 0.08s Tokens generated: 5 Tokens/second: 59.08 @@ -4840,7 +4809,6 @@ Created neuron_vs_cpu_comparison.py - thats incorrect, firstly use the weights and config and tokenizer in llama3_neuron_checkpoint directory. second use the huggingface llama model. do not use any other model. Kiro @@ -4855,13 +4823,12 @@ Created cpu_llama3_inference.py - Command -python cpu_llama3_inference.py --prompt "Hello, how are you?" --max_new_tokens 5 -WARNING:__main__:Missing keys: ['model.embed_tokens.weight', 'model.layers.0.self_attn.q_ -proj.weight', 'model.layers.0.self_attn.k_proj.weight', 'model.layers.0.self_attn.v_proj.weight', 'model.layers.0.self_attn.o_proj.weight', 'model.layers.0.mlp.gate_proj.weight', 'model.layers.0.mlp.up_proj.weight', 'model.layers.0.mlp.down_proj.weight', 'model.layers.0.input_layernorm.weight', 'model.layers.0.post_attention_layernorm.weight', 'model.layers.1.self_attn.q_proj.weight', 'model.layers.1.self_attn.k_proj.weight', 'model.layers.1.self_attn.v_proj.weight', 'model.layers.1.self_attn.o_proj.weight', 'model.layers.1.mlp.gate_proj.weight', 'model.layers.1.mlp.up_proj.weight', 'model.layers.1.mlp.down_proj.weight', 'model.layers.1.input_layernorm.weight', 'model.layers.1.post_attention_layernorm.weight', 'model.layers.2.self_attn.q_proj.weight', 'model.layers.2.self_attn.k_proj.weight', 'model.layers.2.self_attn.v_proj.weight', 'model.layers.2.self_attn.o_proj.weight', 'model.layers.2.mlp.gate_proj.weight', 'model.layers.2.mlp.up_proj.weight', 'model.layers.2.mlp.down_proj.weight', 'model.layers.2.input_layernorm.weight', 'model.layers.2.post_attention_layernorm.weight', 'model.layers.3.self_attn.q_proj.weight', 'model.layers.3.self_attn.k_proj.weight', 'model.layers.3.self_attn.v_proj.weight', 'model.layers.3.self_attn.o_proj.weight', 'model.layers.3.mlp.gate_proj.weight', 'model.layers.3.mlp.up_proj.weight', 'model.layers.3.mlp.down_proj.weight', 'model.layers.3.input_layernorm.weight', 'model.layers.3.post_attention_layernorm.weight', 'model.layers.4.self_attn.q_proj.weight', 'model.layers.4.self_attn.k_proj.weight', 'model.layers.4.self_attn.v_proj.weight', 'model.layers.4.self_attn.o_proj.weight', 'model.layers.4.mlp.gate_proj.weight', 'model.layers.4.mlp.up_proj.weight', 'model.layers.4.mlp.down_proj.weight', 'model.layers.4.input_layernorm.weight', 'model.layers.4.post_attention_layernorm.weight', 'model.layers.5.self_attn.q_proj.weight', 'model.layers.5.self_attn.k_proj.weight', 'model.layers.5.self_attn.v_proj.weight', 'model.layers.5.self_attn.o_proj.weight', 'model.layers.5.mlp.gate_proj.weight', 'model.layers.5.mlp.up_proj.weight', 'model.layers.5.mlp.down_proj.weight', 'model.layers.5.input_layernorm.weight', 'model.layers.5.post_attention_layernorm.weight', 'model.layers.6.self_attn.q_proj.weight', 'model.layers.6.self_attn.k_proj.weight', 'model.layers.6.self_attn.v_proj.weight', 'model.layers.6.self_attn.o_proj.weight', 'model.layers.6.mlp.gate_proj.weight', 'model.layers.6.mlp.up_proj.weight', 'model.layers.6.mlp.down_proj.weight', 'model.layers.6.input_layernorm.weight', 'model.layers.6.post_attention_layernorm.weight', 'model.layers.7.self_attn.q_proj.weight', 'model.layers.7.self_attn.k_proj.weight', 'model.layers.7.self_attn.v_proj.weight', 'model.layers.7.self_attn.o_proj.weight', 'model.layers.7.mlp.gate_proj.weight', 'model.layers.7.mlp.up_proj.weight', 'model.layers.7.mlp.down_proj.weight', 'model.layers.7.input_layernorm.weight', 'model.layers.7.post_attention_layernorm.weight', 'model.layers.8.self_attn.q_proj.weight', 'model.layers.8.self_attn.k_proj.weight', 'model.layers.8.self_attn.v_proj.weight', 'model.layers.8.self_attn.o_proj.weight', 'model.layers.8.mlp.gate_proj.weight', 'model.layers.8.mlp.up_proj.weight', 'model.layers.8.mlp.down_proj.weight', 'model.layers.8.input_layernorm.weight', 'model.layers.8.post_attention_layernorm.weight', 'model.layers.9.self_attn.q_proj.weight', 'model.layers.9.self_attn.k_proj.weight', 'model.layers.9.self_attn.v_proj.weight', 'model.layers.9.self_attn.o_proj.weight', 'model.layers.9.mlp.gate_proj.weight', 'model.layers.9.mlp.up_proj.weight', 'model.layers.9.mlp.down_proj.weight', 'model.layers.9.input_layernorm.weight', 'model.layers.9.post_attention_layernorm.weight', 'model.layers.10.self_attn.q_proj.weight', 'model.layers.10.self_attn.k_proj.weight', 'model.layers.10.self_attn.v_proj.weight', 'model.layers.10.self_attn.o_proj.weight', 'model.layers.10.mlp.gate_proj.weight', 'model.layers.10.mlp.up_proj.weight', 'model.layers.10.mlp.down_proj.weight', 'model.layers.10.input_layernorm.weight', 'model.layers.10.post_attention_layernorm.weight', 'model.layers.11.self_attn.q_proj.weight', 'model.layers.11.self_attn.k_proj.weight', 'model.layers.11.self_attn.v_proj.weight', 'model.layers.11.self_attn.o_proj.weight', 'model.layers.11.mlp.gate_proj.weight', 'model.layers.11.mlp.up_proj.weight', 'model.layers.11.mlp.down_proj.weight', 'model.layers.11.input_layernorm.weight', 'model.layers.11.post_attention_layernorm.weight', 'model.layers.12.self_attn.q_proj.weight', 'model.layers.12.self_attn.k_proj.weight', 'model.layers.12.self_attn.v_proj.weight', 'model.layers.12.self_attn.o_proj.weight', 'model.layers.12.mlp.gate_proj.weight', 'model.layers.12.mlp.up_proj.weight', 'model.layers.12.mlp.down_proj.weight', 'model.layers.12.input_layernorm.weight', 'model.layers.12.post_attention_layernorm.weight', 'model.layers.13.self_attn.q_proj.weight', 'model.layers.13.self_attn.k_proj.weight', 'model.layers.13.self_attn.v_proj.weight', 'model.layers.13.self_attn.o_proj.weight', 'model.layers.13.mlp.gate_proj.weight', 'model.layers.13.mlp.up_proj.weight', 'model.layers.13.mlp.down_proj.weight', 'model.layers.13.input_layernorm.weight', 'model.layers.13.post_attention_layernorm.weight', 'model.layers.14.self_attn.q_proj.weight', 'model.layers.14.self_attn.k_proj.weight', 'model.layers.14.self_attn.v_proj.weight', 'model.layers.14.self_attn.o_proj.weight', 'model.layers.14.mlp.gate_proj.weight', 'model.layers.14.mlp.up_proj.weight', 'model.layers.14.mlp.down_proj.weight', 'model.layers.14.input_layernorm.weight', 'model.layers.14.post_attention_layernorm.weight', 'model.layers.15.self_attn.q_proj.weight', 'model.layers.15.self_attn.k_proj.weight', 'model.layers.15.self_attn.v_proj.weight', 'model.layers.15.self_attn.o_proj.weight', 'model.layers.15.mlp.gate_proj.weight', 'model.layers.15.mlp.up_proj.weight', 'model.layers.15.mlp.down_proj.weight', 'model.layers.15.input_layernorm.weight', 'model.layers.15.post_attention_layernorm.weight', 'model.norm.weight'] WARNING:__main__:Unexpected keys: ['embed_tokens.weight', 'norm.weight', 'layers.0.self_a -ttn.q_proj.weight', 'layers.0.self_attn.k_proj.weight', 'layers.0.self_attn.v_proj.weight', 'layers.0.self_attn.o_proj.weight', 'layers.0.mlp.gate_proj.weight', 'layers.0.mlp.down_proj.weight', 'layers.0.mlp.up_proj.weight', 'layers.0.input_layernorm.weight', 'layers.0.post_attention_layernorm.weight', 'layers.1.self_attn.q_proj.weight', 'layers.1.self_attn.k_proj.weight', 'layers.1.self_attn.v_proj.weight', 'layers.1.self_attn.o_proj.weight', 'layers.1.mlp.gate_proj.weight', 'layers.1.mlp.down_proj.weight', 'layers.1.mlp.up_proj.weight', 'layers.1.input_layernorm.weight', 'layers.1.post_attention_layernorm.weight', 'layers.2.self_attn.q_proj.weight', 'layers.2.self_attn.k_proj.weight', 'layers.2.self_attn.v_proj.weight', 'layers.2.self_attn.o_proj.weight', 'layers.2.mlp.gate_proj.weight', 'layers.2.mlp.down_proj.weight', 'layers.2.mlp.up_proj.weight', 'layers.2.input_layernorm.weight', 'layers.2.post_attention_layernorm.weight', 'layers.3.self_attn.q_proj.weight', 'layers.3.self_attn.k_proj.weight', 'layers.3.self_attn.v_proj.weight', 'layers.3.self_attn.o_proj.weight', 'layers.3.mlp.gate_proj.weight', 'layers.3.mlp.down_proj.weight', 'layers.3.mlp.up_proj.weight', 'layers.3.input_layernorm.weight', 'layers.3.post_attention_layernorm.weight', 'layers.4.self_attn.q_proj.weight', 'layers.4.self_attn.k_proj.weight', 'layers.4.self_attn.v_proj.weight', 'layers.4.self_attn.o_proj.weight', 'layers.4.mlp.gate_proj.weight', 'layers.4.mlp.down_proj.weight', 'layers.4.mlp.up_proj.weight', 'layers.4.input_layernorm.weight', 'layers.4.post_attention_layernorm.weight', 'layers.5.self_attn.q_proj.weight', 'layers.5.self_attn.k_proj.weight', 'layers.5.self_attn.v_proj.weight', 'layers.5.self_attn.o_proj.weight', 'layers.5.mlp.gate_proj.weight', 'layers.5.mlp.down_proj.weight', 'layers.5.mlp.up_proj.weight', 'layers.5.input_layernorm.weight', 'layers.5.post_attention_layernorm.weight', 'layers.6.self_attn.q_proj.weight', 'layers.6.self_attn.k_proj.weight', 'layers.6.self_attn.v_proj.weight', 'layers.6.self_attn.o_proj.weight', 'layers.6.mlp.gate_proj.weight', 'layers.6.mlp.down_proj.weight', 'layers.6.mlp.up_proj.weight', 'layers.6.input_layernorm.weight', 'layers.6.post_attention_layernorm.weight', 'layers.7.self_attn.q_proj.weight', 'layers.7.self_attn.k_proj.weight', 'layers.7.self_attn.v_proj.weight', 'layers.7.self_attn.o_proj.weight', 'layers.7.mlp.gate_proj.weight', 'layers.7.mlp.down_proj.weight', 'layers.7.mlp.up_proj.weight', 'layers.7.input_layernorm.weight', 'layers.7.post_attention_layernorm.weight', 'layers.8.self_attn.q_proj.weight', 'layers.8.self_attn.k_proj.weight', 'layers.8.self_attn.v_proj.weight', 'layers.8.self_attn.o_proj.weight', 'layers.8.mlp.gate_proj.weight', 'layers.8.mlp.down_proj.weight', 'layers.8.mlp.up_proj.weight', 'layers.8.input_layernorm.weight', 'layers.8.post_attention_layernorm.weight', 'layers.9.self_attn.q_proj.weight', 'layers.9.self_attn.k_proj.weight', 'layers.9.self_attn.v_proj.weight', 'layers.9.self_attn.o_proj.weight', 'layers.9.mlp.gate_proj.weight', 'layers.9.mlp.down_proj.weight', 'layers.9.mlp.up_proj.weight', 'layers.9.input_layernorm.weight', 'layers.9.post_attention_layernorm.weight', 'layers.10.self_attn.q_proj.weight', 'layers.10.self_attn.k_proj.weight', 'layers.10.self_attn.v_proj.weight', 'layers.10.self_attn.o_proj.weight', 'layers.10.mlp.gate_proj.weight', 'layers.10.mlp.down_proj.weight', 'layers.10.mlp.up_proj.weight', 'layers.10.input_layernorm.weight', 'layers.10.post_attention_layernorm.weight', 'layers.11.self_attn.q_proj.weight', 'layers.11.self_attn.k_proj.weight', 'layers.11.self_attn.v_proj.weight', 'layers.11.self_attn.o_proj.weight', 'layers.11.mlp.gate_proj.weight', 'layers.11.mlp.down_proj.weight', 'layers.11.mlp.up_proj.weight', 'layers.11.input_layernorm.weight', 'layers.11.post_attention_layernorm.weight', 'layers.12.self_attn.q_proj.weight', 'layers.12.self_attn.k_proj.weight', 'layers.12.self_attn.v_proj.weight', 'layers.12.self_attn.o_proj.weight', 'layers.12.mlp.gate_proj.weight', 'layers.12.mlp.down_proj.weight', 'layers.12.mlp.up_proj.weight', 'layers.12.input_layernorm.weight', 'layers.12.post_attention_layernorm.weight', 'layers.13.self_attn.q_proj.weight', 'layers.13.self_attn.k_proj.weight', 'layers.13.self_attn.v_proj.weight', 'layers.13.self_attn.o_proj.weight', 'layers.13.mlp.gate_proj.weight', 'layers.13.mlp.down_proj.weight', 'layers.13.mlp.up_proj.weight', 'layers.13.input_layernorm.weight', 'layers.13.post_attention_layernorm.weight', 'layers.14.self_attn.q_proj.weight', 'layers.14.self_attn.k_proj.weight', 'layers.14.self_attn.v_proj.weight', 'layers.14.self_attn.o_proj.weight', 'layers.14.mlp.gate_proj.weight', 'layers.14.mlp.down_proj.weight', 'layers.14.mlp.up_proj.weight', 'layers.14.input_layernorm.weight', 'layers.14.post_attention_layernorm.weight', 'layers.15.self_attn.q_proj.weight', 'layers.15.self_attn.k_proj.weight', 'layers.15.self_attn.v_proj.weight', 'layers.15.self_attn.o_proj.weight', 'layers.15.mlp.gate_proj.weight', 'layers.15.mlp.down_proj.weight', 'layers.15.mlp.up_proj.weight', 'layers.15.input_layernorm.weight', 'layers.15.post_attention_layernorm.weight', 'layers.0.self_attn.rank_util.rank', 'layers.1.self_attn.rank_util.rank', 'layers.2.self_attn.rank_util.rank', 'layers.3.self_attn.rank_util.rank', 'layers.4.self_attn.rank_util.rank', 'layers.5.self_attn.rank_util.rank', 'layers.6.self_attn.rank_util.rank', 'layers.7.self_attn.rank_util.rank', 'layers.8.self_attn.rank_util.rank', 'layers.9.self_attn.rank_util.rank', 'layers.10.self_attn.rank_util.rank', 'layers.11.self_attn.rank_util.rank', 'layers.12.self_attn.rank_util.rank', 'layers.13.self_attn.rank_util.rank', 'layers.14.self_attn.rank_util.rank', 'layers.15.self_attn.rank_util.rank', 'rank_util.rank'] `generation_config` default values have been modified to match model-specific defaults: { -'bos_token_id': 1}. If this is not desired, please set these values explicitly. +python cpu*llama3_inference.py --prompt "Hello, how are you?" --max_new_tokens 5 +WARNING:**main**:Missing keys: ['model.embed_tokens.weight', 'model.layers.0.self_attn.q* +proj.weight', 'model.layers.0.self_attn.k_proj.weight', 'model.layers.0.self_attn.v_proj.weight', 'model.layers.0.self_attn.o_proj.weight', 'model.layers.0.mlp.gate_proj.weight', 'model.layers.0.mlp.up_proj.weight', 'model.layers.0.mlp.down_proj.weight', 'model.layers.0.input_layernorm.weight', 'model.layers.0.post_attention_layernorm.weight', 'model.layers.1.self_attn.q_proj.weight', 'model.layers.1.self_attn.k_proj.weight', 'model.layers.1.self_attn.v_proj.weight', 'model.layers.1.self_attn.o_proj.weight', 'model.layers.1.mlp.gate_proj.weight', 'model.layers.1.mlp.up_proj.weight', 'model.layers.1.mlp.down_proj.weight', 'model.layers.1.input_layernorm.weight', 'model.layers.1.post_attention_layernorm.weight', 'model.layers.2.self_attn.q_proj.weight', 'model.layers.2.self_attn.k_proj.weight', 'model.layers.2.self_attn.v_proj.weight', 'model.layers.2.self_attn.o_proj.weight', 'model.layers.2.mlp.gate_proj.weight', 'model.layers.2.mlp.up_proj.weight', 'model.layers.2.mlp.down_proj.weight', 'model.layers.2.input_layernorm.weight', 'model.layers.2.post_attention_layernorm.weight', 'model.layers.3.self_attn.q_proj.weight', 'model.layers.3.self_attn.k_proj.weight', 'model.layers.3.self_attn.v_proj.weight', 'model.layers.3.self_attn.o_proj.weight', 'model.layers.3.mlp.gate_proj.weight', 'model.layers.3.mlp.up_proj.weight', 'model.layers.3.mlp.down_proj.weight', 'model.layers.3.input_layernorm.weight', 'model.layers.3.post_attention_layernorm.weight', 'model.layers.4.self_attn.q_proj.weight', 'model.layers.4.self_attn.k_proj.weight', 'model.layers.4.self_attn.v_proj.weight', 'model.layers.4.self_attn.o_proj.weight', 'model.layers.4.mlp.gate_proj.weight', 'model.layers.4.mlp.up_proj.weight', 'model.layers.4.mlp.down_proj.weight', 'model.layers.4.input_layernorm.weight', 'model.layers.4.post_attention_layernorm.weight', 'model.layers.5.self_attn.q_proj.weight', 'model.layers.5.self_attn.k_proj.weight', 'model.layers.5.self_attn.v_proj.weight', 'model.layers.5.self_attn.o_proj.weight', 'model.layers.5.mlp.gate_proj.weight', 'model.layers.5.mlp.up_proj.weight', 'model.layers.5.mlp.down_proj.weight', 'model.layers.5.input_layernorm.weight', 'model.layers.5.post_attention_layernorm.weight', 'model.layers.6.self_attn.q_proj.weight', 'model.layers.6.self_attn.k_proj.weight', 'model.layers.6.self_attn.v_proj.weight', 'model.layers.6.self_attn.o_proj.weight', 'model.layers.6.mlp.gate_proj.weight', 'model.layers.6.mlp.up_proj.weight', 'model.layers.6.mlp.down_proj.weight', 'model.layers.6.input_layernorm.weight', 'model.layers.6.post_attention_layernorm.weight', 'model.layers.7.self_attn.q_proj.weight', 'model.layers.7.self_attn.k_proj.weight', 'model.layers.7.self_attn.v_proj.weight', 'model.layers.7.self_attn.o_proj.weight', 'model.layers.7.mlp.gate_proj.weight', 'model.layers.7.mlp.up_proj.weight', 'model.layers.7.mlp.down_proj.weight', 'model.layers.7.input_layernorm.weight', 'model.layers.7.post_attention_layernorm.weight', 'model.layers.8.self_attn.q_proj.weight', 'model.layers.8.self_attn.k_proj.weight', 'model.layers.8.self_attn.v_proj.weight', 'model.layers.8.self_attn.o_proj.weight', 'model.layers.8.mlp.gate_proj.weight', 'model.layers.8.mlp.up_proj.weight', 'model.layers.8.mlp.down_proj.weight', 'model.layers.8.input_layernorm.weight', 'model.layers.8.post_attention_layernorm.weight', 'model.layers.9.self_attn.q_proj.weight', 'model.layers.9.self_attn.k_proj.weight', 'model.layers.9.self_attn.v_proj.weight', 'model.layers.9.self_attn.o_proj.weight', 'model.layers.9.mlp.gate_proj.weight', 'model.layers.9.mlp.up_proj.weight', 'model.layers.9.mlp.down_proj.weight', 'model.layers.9.input_layernorm.weight', 'model.layers.9.post_attention_layernorm.weight', 'model.layers.10.self_attn.q_proj.weight', 'model.layers.10.self_attn.k_proj.weight', 'model.layers.10.self_attn.v_proj.weight', 'model.layers.10.self_attn.o_proj.weight', 'model.layers.10.mlp.gate_proj.weight', 'model.layers.10.mlp.up_proj.weight', 'model.layers.10.mlp.down_proj.weight', 'model.layers.10.input_layernorm.weight', 'model.layers.10.post_attention_layernorm.weight', 'model.layers.11.self_attn.q_proj.weight', 'model.layers.11.self_attn.k_proj.weight', 'model.layers.11.self_attn.v_proj.weight', 'model.layers.11.self_attn.o_proj.weight', 'model.layers.11.mlp.gate_proj.weight', 'model.layers.11.mlp.up_proj.weight', 'model.layers.11.mlp.down_proj.weight', 'model.layers.11.input_layernorm.weight', 'model.layers.11.post_attention_layernorm.weight', 'model.layers.12.self_attn.q_proj.weight', 'model.layers.12.self_attn.k_proj.weight', 'model.layers.12.self_attn.v_proj.weight', 'model.layers.12.self_attn.o_proj.weight', 'model.layers.12.mlp.gate_proj.weight', 'model.layers.12.mlp.up_proj.weight', 'model.layers.12.mlp.down_proj.weight', 'model.layers.12.input_layernorm.weight', 'model.layers.12.post_attention_layernorm.weight', 'model.layers.13.self_attn.q_proj.weight', 'model.layers.13.self_attn.k_proj.weight', 'model.layers.13.self_attn.v_proj.weight', 'model.layers.13.self_attn.o_proj.weight', 'model.layers.13.mlp.gate_proj.weight', 'model.layers.13.mlp.up_proj.weight', 'model.layers.13.mlp.down_proj.weight', 'model.layers.13.input_layernorm.weight', 'model.layers.13.post_attention_layernorm.weight', 'model.layers.14.self_attn.q_proj.weight', 'model.layers.14.self_attn.k_proj.weight', 'model.layers.14.self_attn.v_proj.weight', 'model.layers.14.self_attn.o_proj.weight', 'model.layers.14.mlp.gate_proj.weight', 'model.layers.14.mlp.up_proj.weight', 'model.layers.14.mlp.down_proj.weight', 'model.layers.14.input_layernorm.weight', 'model.layers.14.post_attention_layernorm.weight', 'model.layers.15.self_attn.q_proj.weight', 'model.layers.15.self_attn.k_proj.weight', 'model.layers.15.self_attn.v_proj.weight', 'model.layers.15.self_attn.o_proj.weight', 'model.layers.15.mlp.gate_proj.weight', 'model.layers.15.mlp.up_proj.weight', 'model.layers.15.mlp.down_proj.weight', 'model.layers.15.input_layernorm.weight', 'model.layers.15.post_attention_layernorm.weight', 'model.norm.weight'] WARNING:**main**:Unexpected keys: ['embed_tokens.weight', 'norm.weight', 'layers.0.self_a +ttn.q_proj.weight', 'layers.0.self_attn.k_proj.weight', 'layers.0.self_attn.v_proj.weight', 'layers.0.self_attn.o_proj.weight', 'layers.0.mlp.gate_proj.weight', 'layers.0.mlp.down_proj.weight', 'layers.0.mlp.up_proj.weight', 'layers.0.input_layernorm.weight', 'layers.0.post_attention_layernorm.weight', 'layers.1.self_attn.q_proj.weight', 'layers.1.self_attn.k_proj.weight', 'layers.1.self_attn.v_proj.weight', 'layers.1.self_attn.o_proj.weight', 'layers.1.mlp.gate_proj.weight', 'layers.1.mlp.down_proj.weight', 'layers.1.mlp.up_proj.weight', 'layers.1.input_layernorm.weight', 'layers.1.post_attention_layernorm.weight', 'layers.2.self_attn.q_proj.weight', 'layers.2.self_attn.k_proj.weight', 'layers.2.self_attn.v_proj.weight', 'layers.2.self_attn.o_proj.weight', 'layers.2.mlp.gate_proj.weight', 'layers.2.mlp.down_proj.weight', 'layers.2.mlp.up_proj.weight', 'layers.2.input_layernorm.weight', 'layers.2.post_attention_layernorm.weight', 'layers.3.self_attn.q_proj.weight', 'layers.3.self_attn.k_proj.weight', 'layers.3.self_attn.v_proj.weight', 'layers.3.self_attn.o_proj.weight', 'layers.3.mlp.gate_proj.weight', 'layers.3.mlp.down_proj.weight', 'layers.3.mlp.up_proj.weight', 'layers.3.input_layernorm.weight', 'layers.3.post_attention_layernorm.weight', 'layers.4.self_attn.q_proj.weight', 'layers.4.self_attn.k_proj.weight', 'layers.4.self_attn.v_proj.weight', 'layers.4.self_attn.o_proj.weight', 'layers.4.mlp.gate_proj.weight', 'layers.4.mlp.down_proj.weight', 'layers.4.mlp.up_proj.weight', 'layers.4.input_layernorm.weight', 'layers.4.post_attention_layernorm.weight', 'layers.5.self_attn.q_proj.weight', 'layers.5.self_attn.k_proj.weight', 'layers.5.self_attn.v_proj.weight', 'layers.5.self_attn.o_proj.weight', 'layers.5.mlp.gate_proj.weight', 'layers.5.mlp.down_proj.weight', 'layers.5.mlp.up_proj.weight', 'layers.5.input_layernorm.weight', 'layers.5.post_attention_layernorm.weight', 'layers.6.self_attn.q_proj.weight', 'layers.6.self_attn.k_proj.weight', 'layers.6.self_attn.v_proj.weight', 'layers.6.self_attn.o_proj.weight', 'layers.6.mlp.gate_proj.weight', 'layers.6.mlp.down_proj.weight', 'layers.6.mlp.up_proj.weight', 'layers.6.input_layernorm.weight', 'layers.6.post_attention_layernorm.weight', 'layers.7.self_attn.q_proj.weight', 'layers.7.self_attn.k_proj.weight', 'layers.7.self_attn.v_proj.weight', 'layers.7.self_attn.o_proj.weight', 'layers.7.mlp.gate_proj.weight', 'layers.7.mlp.down_proj.weight', 'layers.7.mlp.up_proj.weight', 'layers.7.input_layernorm.weight', 'layers.7.post_attention_layernorm.weight', 'layers.8.self_attn.q_proj.weight', 'layers.8.self_attn.k_proj.weight', 'layers.8.self_attn.v_proj.weight', 'layers.8.self_attn.o_proj.weight', 'layers.8.mlp.gate_proj.weight', 'layers.8.mlp.down_proj.weight', 'layers.8.mlp.up_proj.weight', 'layers.8.input_layernorm.weight', 'layers.8.post_attention_layernorm.weight', 'layers.9.self_attn.q_proj.weight', 'layers.9.self_attn.k_proj.weight', 'layers.9.self_attn.v_proj.weight', 'layers.9.self_attn.o_proj.weight', 'layers.9.mlp.gate_proj.weight', 'layers.9.mlp.down_proj.weight', 'layers.9.mlp.up_proj.weight', 'layers.9.input_layernorm.weight', 'layers.9.post_attention_layernorm.weight', 'layers.10.self_attn.q_proj.weight', 'layers.10.self_attn.k_proj.weight', 'layers.10.self_attn.v_proj.weight', 'layers.10.self_attn.o_proj.weight', 'layers.10.mlp.gate_proj.weight', 'layers.10.mlp.down_proj.weight', 'layers.10.mlp.up_proj.weight', 'layers.10.input_layernorm.weight', 'layers.10.post_attention_layernorm.weight', 'layers.11.self_attn.q_proj.weight', 'layers.11.self_attn.k_proj.weight', 'layers.11.self_attn.v_proj.weight', 'layers.11.self_attn.o_proj.weight', 'layers.11.mlp.gate_proj.weight', 'layers.11.mlp.down_proj.weight', 'layers.11.mlp.up_proj.weight', 'layers.11.input_layernorm.weight', 'layers.11.post_attention_layernorm.weight', 'layers.12.self_attn.q_proj.weight', 'layers.12.self_attn.k_proj.weight', 'layers.12.self_attn.v_proj.weight', 'layers.12.self_attn.o_proj.weight', 'layers.12.mlp.gate_proj.weight', 'layers.12.mlp.down_proj.weight', 'layers.12.mlp.up_proj.weight', 'layers.12.input_layernorm.weight', 'layers.12.post_attention_layernorm.weight', 'layers.13.self_attn.q_proj.weight', 'layers.13.self_attn.k_proj.weight', 'layers.13.self_attn.v_proj.weight', 'layers.13.self_attn.o_proj.weight', 'layers.13.mlp.gate_proj.weight', 'layers.13.mlp.down_proj.weight', 'layers.13.mlp.up_proj.weight', 'layers.13.input_layernorm.weight', 'layers.13.post_attention_layernorm.weight', 'layers.14.self_attn.q_proj.weight', 'layers.14.self_attn.k_proj.weight', 'layers.14.self_attn.v_proj.weight', 'layers.14.self_attn.o_proj.weight', 'layers.14.mlp.gate_proj.weight', 'layers.14.mlp.down_proj.weight', 'layers.14.mlp.up_proj.weight', 'layers.14.input_layernorm.weight', 'layers.14.post_attention_layernorm.weight', 'layers.15.self_attn.q_proj.weight', 'layers.15.self_attn.k_proj.weight', 'layers.15.self_attn.v_proj.weight', 'layers.15.self_attn.o_proj.weight', 'layers.15.mlp.gate_proj.weight', 'layers.15.mlp.down_proj.weight', 'layers.15.mlp.up_proj.weight', 'layers.15.input_layernorm.weight', 'layers.15.post_attention_layernorm.weight', 'layers.0.self_attn.rank_util.rank', 'layers.1.self_attn.rank_util.rank', 'layers.2.self_attn.rank_util.rank', 'layers.3.self_attn.rank_util.rank', 'layers.4.self_attn.rank_util.rank', 'layers.5.self_attn.rank_util.rank', 'layers.6.self_attn.rank_util.rank', 'layers.7.self_attn.rank_util.rank', 'layers.8.self_attn.rank_util.rank', 'layers.9.self_attn.rank_util.rank', 'layers.10.self_attn.rank_util.rank', 'layers.11.self_attn.rank_util.rank', 'layers.12.self_attn.rank_util.rank', 'layers.13.self_attn.rank_util.rank', 'layers.14.self_attn.rank_util.rank', 'layers.15.self_attn.rank_util.rank', 'rank_util.rank'] `generation_config` default values have been modified to match model-specific defaults: { +'bos_token_id': 1}. If this is not desired, please set these values explicitly. ============================================================ CPU LLAMA3 INFERENCE RESULTS (Using Actual Weights) ============================================================ @@ -4870,10 +4837,10 @@ Model config: 16 layers, 2048 hidden_size Generation settings: max_tokens=5, temperature=0.6, top_p=0.9, do_sample=False ============================================================ -Prompt: Hello, how are you? ----------------------------------------- -Generated: Hello, how are you? Constantoise cover Otherwiseunei ----------------------------------------- +## Prompt: Hello, how are you? + +## Generated: Hello, how are you? Constantoise cover Otherwiseunei + Generation time: 0.36s Tokens generated: 5 Tokens/second: 13.87 @@ -4883,12 +4850,11 @@ Accepted edits to cpu_llama3_inference.py - Command python cpu_llama3_inference.py --prompt "Hello, how are you?" --max_new_tokens 5 -WARNING:__main__:Unexpected keys: ['model.layers.0.self_attn.rank_util.rank', 'model.layers.1.se -lf_attn.rank_util.rank', 'model.layers.2.self_attn.rank_util.rank', 'model.layers.3.self_attn.rank_util.rank', 'model.layers.4.self_attn.rank_util.rank', 'model.layers.5.self_attn.rank_util.rank', 'model.layers.6.self_attn.rank_util.rank', 'model.layers.7.self_attn.rank_util.rank', 'model.layers.8.self_attn.rank_util.rank', 'model.layers.9.self_attn.rank_util.rank', 'model.layers.10.self_attn.rank_util.rank', 'model.layers.11.self_attn.rank_util.rank', 'model.layers.12.self_attn.rank_util.rank', 'model.layers.13.self_attn.rank_util.rank', 'model.layers.14.self_attn.rank_util.rank', 'model.layers.15.self_attn.rank_util.rank'] `generation_config` default values have been modified to match model-specific defaults: {'bos_to -ken_id': 1}. If this is not desired, please set these values explicitly. +WARNING:**main**:Unexpected keys: ['model.layers.0.self_attn.rank_util.rank', 'model.layers.1.se +lf_attn.rank_util.rank', 'model.layers.2.self_attn.rank_util.rank', 'model.layers.3.self_attn.rank_util.rank', 'model.layers.4.self_attn.rank_util.rank', 'model.layers.5.self_attn.rank_util.rank', 'model.layers.6.self_attn.rank_util.rank', 'model.layers.7.self_attn.rank_util.rank', 'model.layers.8.self_attn.rank_util.rank', 'model.layers.9.self_attn.rank_util.rank', 'model.layers.10.self_attn.rank_util.rank', 'model.layers.11.self_attn.rank_util.rank', 'model.layers.12.self_attn.rank_util.rank', 'model.layers.13.self_attn.rank_util.rank', 'model.layers.14.self_attn.rank_util.rank', 'model.layers.15.self_attn.rank_util.rank'] `generation_config` default values have been modified to match model-specific defaults: {'bos_to +ken_id': 1}. If this is not desired, please set these values explicitly. ============================================================ CPU LLAMA3 INFERENCE RESULTS (Using Actual Weights) ============================================================ @@ -4897,10 +4863,10 @@ Model config: 16 layers, 2048 hidden_size Generation settings: max_tokens=5, temperature=0.6, top_p=0.9, do_sample=False ============================================================ -Prompt: Hello, how are you? ----------------------------------------- -Generated: Hello, how are you? I am I am I ----------------------------------------- +## Prompt: Hello, how are you? + +## Generated: Hello, how are you? I am I am I + Generation time: 0.34s Tokens generated: 5 Tokens/second: 14.80 @@ -4916,7 +4882,6 @@ Accepted edits to INFERENCE_TROUBLESHOOTING_GUIDE.md - Now we have: Working Neuron inference (simple_inference.py) - Uses our custom model class with manual generation loop @@ -4933,17 +4898,10 @@ Created compare_inference.py - Session Too Long Continue Summarize and continue in a new session. - - - - - - Claude Sonnet 4.0 Autopilot @@ -4952,10 +4910,6 @@ New Session New Session (Continu... - - - - Kiro Understood. I have the context from our previous conversation and am ready to continue helping you. @@ -4996,95 +4950,95 @@ NEURONX vs CPU INFERENCE COMPARISON Testing 10 prompts with max_new_tokens=10 ================================================================================ -[1/10] Testing prompt: 'Hello, how are you?' ------------------------------------------------------------- -INFO:__main__:Running NeuronX inference for: 'Hello, how are you?...' -INFO:__main__:Running CPU inference for: 'Hello, how are you?...' +## [1/10] Testing prompt: 'Hello, how are you?' + +INFO:**main**:Running NeuronX inference for: 'Hello, how are you?...' +INFO:**main**:Running CPU inference for: 'Hello, how are you?...' NeuronX: Hello, how are you? I am I am I am I am I am (22.36s) -CPU: Hello, how are you? I am I am I am I am I am (25.01s) - ✅ Outputs are identical - 🚀 NeuronX is 1.12x faster - -[2/10] Testing prompt: 'The meaning of life is' ------------------------------------------------------------- -INFO:__main__:Running NeuronX inference for: 'The meaning of life is...' -INFO:__main__:Running CPU inference for: 'The meaning of life is...' +CPU: Hello, how are you? I am I am I am I am I am (25.01s) +✅ Outputs are identical +🚀 NeuronX is 1.12x faster + +## [2/10] Testing prompt: 'The meaning of life is' + +INFO:**main**:Running NeuronX inference for: 'The meaning of life is...' +INFO:**main**:Running CPU inference for: 'The meaning of life is...' NeuronX: The meaning of life is a word is a word that you can be used (22.25s) -CPU: The meaning of life is a word is a word that you can be used (25.27s) - ✅ Outputs are identical - 🚀 NeuronX is 1.14x faster - -[3/10] Testing prompt: 'Once upon a time' ------------------------------------------------------------- -INFO:__main__:Running NeuronX inference for: 'Once upon a time...' -INFO:__main__:Running CPU inference for: 'Once upon a time...' +CPU: The meaning of life is a word is a word that you can be used (25.27s) +✅ Outputs are identical +🚀 NeuronX is 1.14x faster + +## [3/10] Testing prompt: 'Once upon a time' + +INFO:**main**:Running NeuronX inference for: 'Once upon a time...' +INFO:**main**:Running CPU inference for: 'Once upon a time...' NeuronX: Once upon a time, a long time, the other day, I (21.95s) -CPU: Once upon a time, a long time, the other day, I (25.19s) - ✅ Outputs are identical - 🚀 NeuronX is 1.15x faster - -[4/10] Testing prompt: 'What is artificial intelligence?' ------------------------------------------------------------- -INFO:__main__:Running NeuronX inference for: 'What is artificial intelligenc...' -INFO:__main__:Running CPU inference for: 'What is artificial intelligenc...' +CPU: Once upon a time, a long time, the other day, I (25.19s) +✅ Outputs are identical +🚀 NeuronX is 1.15x faster + +## [4/10] Testing prompt: 'What is artificial intelligence?' + +INFO:**main**:Running NeuronX inference for: 'What is artificial intelligenc...' +INFO:**main**:Running CPU inference for: 'What is artificial intelligenc...' NeuronX: What is artificial intelligence? Artificial intelligence? Artificial intelligence? (22.03s) -CPU: What is artificial intelligence? Artificial intelligence? Artificial intelligence? (25.31s) - ✅ Outputs are identical - 🚀 NeuronX is 1.15x faster - -[5/10] Testing prompt: 'Python is a programming language' ------------------------------------------------------------- -INFO:__main__:Running NeuronX inference for: 'Python is a programming langua...' -INFO:__main__:Running CPU inference for: 'Python is a programming langua...' +CPU: What is artificial intelligence? Artificial intelligence? Artificial intelligence? (25.31s) +✅ Outputs are identical +🚀 NeuronX is 1.15x faster + +## [5/10] Testing prompt: 'Python is a programming language' + +INFO:**main**:Running NeuronX inference for: 'Python is a programming langua...' +INFO:**main**:Running CPU inference for: 'Python is a programming langua...' NeuronX: Python is a programming language that is used to create a program to create a (22.15s) -CPU: Python is a programming language that is used to create a program to create a (25.08s) - ✅ Outputs are identical - 🚀 NeuronX is 1.13x faster - -[6/10] Testing prompt: 'The weather today is' ------------------------------------------------------------- -INFO:__main__:Running NeuronX inference for: 'The weather today is...' -INFO:__main__:Running CPU inference for: 'The weather today is...' +CPU: Python is a programming language that is used to create a program to create a (25.08s) +✅ Outputs are identical +🚀 NeuronX is 1.13x faster + +## [6/10] Testing prompt: 'The weather today is' + +INFO:**main**:Running NeuronX inference for: 'The weather today is...' +INFO:**main**:Running CPU inference for: 'The weather today is...' NeuronX: The weather today is a little bit cloudy and the rain, but the (22.02s) -CPU: The weather today is a little bit cloudy and the rain, but the (25.37s) - ✅ Outputs are identical - 🚀 NeuronX is 1.15x faster - -[7/10] Testing prompt: 'I love to eat' ------------------------------------------------------------- -INFO:__main__:Running NeuronX inference for: 'I love to eat...' -INFO:__main__:Running CPU inference for: 'I love to eat...' +CPU: The weather today is a little bit cloudy and the rain, but the (25.37s) +✅ Outputs are identical +🚀 NeuronX is 1.15x faster + +## [7/10] Testing prompt: 'I love to eat' + +INFO:**main**:Running NeuronX inference for: 'I love to eat...' +INFO:**main**:Running CPU inference for: 'I love to eat...' NeuronX: I love to eat, but I love to eat my food, I (22.14s) -CPU: I love to eat, but I love to eat my food, I (25.14s) - ✅ Outputs are identical - 🚀 NeuronX is 1.14x faster - -[8/10] Testing prompt: 'The capital of France is' ------------------------------------------------------------- -INFO:__main__:Running NeuronX inference for: 'The capital of France is...' -INFO:__main__:Running CPU inference for: 'The capital of France is...' +CPU: I love to eat, but I love to eat my food, I (25.14s) +✅ Outputs are identical +🚀 NeuronX is 1.14x faster + +## [8/10] Testing prompt: 'The capital of France is' + +INFO:**main**:Running NeuronX inference for: 'The capital of France is...' +INFO:**main**:Running CPU inference for: 'The capital of France is...' NeuronX: The capital of France is a country of the world’s most beautiful city of (21.99s) -CPU: The capital of France is a country of the world’s most beautiful city of (25.30s) - ✅ Outputs are identical - 🚀 NeuronX is 1.15x faster - -[9/10] Testing prompt: 'Machine learning is' ------------------------------------------------------------- -INFO:__main__:Running NeuronX inference for: 'Machine learning is...' -INFO:__main__:Running CPU inference for: 'Machine learning is...' +CPU: The capital of France is a country of the world’s most beautiful city of (25.30s) +✅ Outputs are identical +🚀 NeuronX is 1.15x faster + +## [9/10] Testing prompt: 'Machine learning is' + +INFO:**main**:Running NeuronX inference for: 'Machine learning is...' +INFO:**main**:Running CPU inference for: 'Machine learning is...' NeuronX: Machine learning is a branch of the 3D (22.02s) -CPU: Machine learning is a branch of the 3D (25.21s) - ✅ Outputs are identical - 🚀 NeuronX is 1.14x faster - -[10/10] Testing prompt: 'In the future, technology will' ------------------------------------------------------------- -INFO:__main__:Running NeuronX inference for: 'In the future, technology will...' -INFO:__main__:Running CPU inference for: 'In the future, technology will...' +CPU: Machine learning is a branch of the 3D (25.21s) +✅ Outputs are identical +🚀 NeuronX is 1.14x faster + +## [10/10] Testing prompt: 'In the future, technology will' + +INFO:**main**:Running NeuronX inference for: 'In the future, technology will...' +INFO:**main**:Running CPU inference for: 'In the future, technology will...' NeuronX: In the future, technology will be able to which is the (21.80s) -CPU: In the future, technology will be able to which is the (25.42s) - ✅ Outputs are identical - 🚀 NeuronX is 1.17x faster +CPU: In the future, technology will be able to which is the (25.42s) +✅ Outputs are identical +🚀 NeuronX is 1.17x faster ================================================================================ SUMMARY STATISTICS @@ -5142,7 +5096,6 @@ Created INFERENCE_COMPARISON_RESULTS.md - Perfect! I've created a comprehensive markdown document that captures all the testing details and results. The document includes: 📋 Complete Documentation: @@ -5168,11 +5121,5 @@ Warning messages and their significance File structure and implementation details The document serves as both a test report and a reference for future development, clearly demonstrating that the NeuronX implementation is working correctly and provides meaningful performance benefits while maintaining perfect functional equivalence with the CPU baseline. - - - - - - Claude Sonnet 4.0 Autopilot diff --git a/skills/neuron-framework-autoport/references/systemPrompts/systemPrompt.md b/skills/neuron-framework-autoport/references/systemPrompts/systemPrompt.md index 6e4fa96..2417952 100644 --- a/skills/neuron-framework-autoport/references/systemPrompts/systemPrompt.md +++ b/skills/neuron-framework-autoport/references/systemPrompts/systemPrompt.md @@ -9,46 +9,55 @@ Dependencies are resolved in SKILL.md's "Resolve Dependencies" step before this ```bash neuron-ls ``` + If 0 cores are detected and the user did not specify dry-run mode, tell the user to allocate a compute node with Neuron hardware and STOP. Clear any stale compile cache: + ```bash rm -rf /var/tmp/neuron-compile-cache ``` - ## Project Guidelines ### Prohibited Packages + - Do not import, reference, or run any code from `transformers_neuronx`. It is an old API library. ### PYTHONPATH issues + - If you run into issues with imports and PYTHONPATH, do not make changes to the script — change PYTHONPATH instead. When you test, do the same. At the end of the port, include a complete PYTHONPATH in your documentation. ### Error Handling + - Do not generate any `try/except` statements. - Let errors surface directly without catching them. - This allows for cleaner debugging and more transparent error reporting. ### File Organization + - Store all temporary files defined as those files that are not the final product and do not contain the modeling or configuration file, for instance those that are used to compile, or test the model (both Python and Markdown) in a sub-directory of the project root called `agent_artifacts/tmp/` - Store all model files that you have generated that contain the modeling or configuration in a directory called neuron_port in a sub-directory of the project root. - This keeps the workspace clean and organizes generated artifacts that can then be deleted later ### Tracing + - For every major step checkpoint prompts, completions, and tool use into a sub-directory of the project root called `agent_artifacts/traces` - This provides an audit trail of agent interactions and decisions ### Weights + - Store all weights, checkpoints and other downloaded artifacts in a sub-directory of the project root called `agent_artifacts/data - Do not store weights, checkpoints and other downloaded artifacts anywhere else except the above directory. ### Critical Knowledge Base to be consulted anytime an issue comes up you cant solve + - A knowledge base exists curated by expert Neuron SDK subject matter experts that have answers to many common, and also unique issues - Location is in NeuroborosFoundations/knowledge_base - Leverage this as it will reduce the number of steps, context you need, and power you consume ### Hardware + - You are on a trn1.32xlarge with 32 NeuronCores and 16GB per core. Use `neuron-ls` to verify if unsure. ## Reference @@ -56,7 +65,8 @@ rm -rf /var/tmp/neuron-compile-cache ### Tool Documentation #### Compile Tool -* **compile_neuron_model**: Compile your ported model for NeuronX hardware + +- **compile_neuron_model**: Compile your ported model for NeuronX hardware - Parameters: model_class_path, config_class_path, neuron_config_class_path, model_path, output_path, batch_size, seq_len, tp_degree, use_fp16 - Returns compilation status and output path - **Run this FIRST** to verify your model compiles successfully @@ -64,7 +74,8 @@ rm -rf /var/tmp/neuron-compile-cache - **HOW TO RUN** always capture output into a file in agent_artifacts/tmp #### Inference Tool -* **run_neuron_inference**: Run inference on a compiled NeuronX model + +- **run_neuron_inference**: Run inference on a compiled NeuronX model - Parameters: model_class_path, config_class_path, model_path, compiled_path, prompt, max_new_tokens, temperature, top_p - Returns generated text and performance metrics - **Run this SECOND** to verify your compiled model can generate text @@ -73,7 +84,9 @@ rm -rf /var/tmp/neuron-compile-cache - **HOW TO RUN** always capture output into a file in agent_artifacts/tmp. If you believe it is actually a correct file make sure you call the output file agent_artifacts/tmp/correct_inference.log #### Validation Tool + Located at `scripts/validate_model.py`. Compares Neuron output against HuggingFace reference. + - Modes: `token` (default, greedy match), `logit` (distribution comparison for debugging), `comprehensive` (both + extra metrics for final validation) - **Success criteria: >= 95% greedy token match rate** - **Run this LAST** after inference works @@ -86,17 +99,21 @@ Located at `scripts/validate_model.py`. Compares Neuron output against HuggingFa ### Debugging Support #### Compiler Issues + - IF you get a JSON error like `[NLA001] Unhandled exception with message: [json.exception.parse_error.101]` THEN delete the compiler cache at `/var/tmp/neuron-compile-cache` and retry - IF you get `FileNotFoundError` on neff_output paths THEN delete the compiler cache at `/var/tmp/neuron-compile-cache` and retry - Use the logs in `agent_artifacts/data/neff_output/context_encoding_model/`, specifically `log-neuron-cc.txt`. Use bash to read these logs. #### Ignorable Warnings + - `WARNING:Neuron:TP degree (XX) and KV heads (YY) are not divisible. Overriding attention sharding strategy to GQA.CONVERT_TO_MHA!` — ignore this, it is not important. ### Codebase Navigation #### NeuronxDistributed (NxD) + Located at `${NXD_SRC}` — resolved during "Resolve Dependencies" in SKILL.md. + - `${NXD_SRC}/src/neuronx_distributed/modules/` — Main transformer modules: attention, LoRA, MoE - `${NXD_SRC}/src/neuronx_distributed/operators/` — Model-specific operators (e.g., argmax) - `${NXD_SRC}/src/neuronx_distributed/overrides/` — Transformer-specific features such as RoPE @@ -104,7 +121,9 @@ Located at `${NXD_SRC}` — resolved during "Resolve Dependencies" in SKILL.md. - Ignore: `kernels/`, `lightning/`, `optimizer/`, `pipeline/`, `scripts/`, `trainer/`, `utils/` #### NeuronxDistributedInference (NxDI) + Located at `${NXDI_SRC}` — resolved during "Resolve Dependencies" in SKILL.md. + - `${NXDI_SRC}/src/neuronx_distributed_inference/modules/` — High-level building blocks - `attention/` — All attention types except sliding window - `moe_v2.py` — MoE architecture (ignore `moe.py`) diff --git a/skills/neuron-framework-equivalence/SKILL.md b/skills/neuron-framework-equivalence/SKILL.md index aaad3ae..c59858b 100644 --- a/skills/neuron-framework-equivalence/SKILL.md +++ b/skills/neuron-framework-equivalence/SKILL.md @@ -11,17 +11,17 @@ Verify and diagnose functional equivalence between a **source** (reference) and Before starting, collect these from the user. Ask for any missing ones. -| Input | Description | Example | -|-------|-------------|---------| -| `SOURCE_MODEL_PATH` | Path to source model weights (HF format) | `/path/to/hf_models/Qwen3-0.6B` | -| `COMPILED_MODEL_PATH` | Path to compiled target model | `/path/to/neuron_models/Qwen3-0.6B` | -| `TARGET_MODELING_FILE` | Path to target's modeling .py file | `/path/to/modeling_qwen3.py` | -| `TARGET_INNER_CLASS` | Inner model class (extends NeuronBaseModel) | `NeuronQwen3Model` | -| `TARGET_CAUSAL_CLASS` | ForCausalLM wrapper class | `NeuronQwen3ForCausalLM` | -| `TARGET_CONFIG_CLASS` | InferenceConfig class | `Qwen3InferenceConfig` | -| `VENV` | Path to Python venv with torch + neuronx | `/opt/aws_neuronx_venv_pytorch_2_8_nxd_inference` | -| `EXP_DIR` | Experiment output directory | `agent_artifacts/equiv_qwen3` | -| `VLLM_NEURON_DIR` | **vLLM-Neuron targets only.** Project root of the `vllm-neuron` editable install | `/path/to/vllm-neuron` | +| Input | Description | Example | +| ---------------------- | -------------------------------------------------------------------------------- | ------------------------------------------------- | +| `SOURCE_MODEL_PATH` | Path to source model weights (HF format) | `/path/to/hf_models/Qwen3-0.6B` | +| `COMPILED_MODEL_PATH` | Path to compiled target model | `/path/to/neuron_models/Qwen3-0.6B` | +| `TARGET_MODELING_FILE` | Path to target's modeling .py file | `/path/to/modeling_qwen3.py` | +| `TARGET_INNER_CLASS` | Inner model class (extends NeuronBaseModel) | `NeuronQwen3Model` | +| `TARGET_CAUSAL_CLASS` | ForCausalLM wrapper class | `NeuronQwen3ForCausalLM` | +| `TARGET_CONFIG_CLASS` | InferenceConfig class | `Qwen3InferenceConfig` | +| `VENV` | Path to Python venv with torch + neuronx | `/opt/aws_neuronx_venv_pytorch_2_8_nxd_inference` | +| `EXP_DIR` | Experiment output directory | `agent_artifacts/equiv_qwen3` | +| `VLLM_NEURON_DIR` | **vLLM-Neuron targets only.** Project root of the `vllm-neuron` editable install | `/path/to/vllm-neuron` | Set `SCRIPTS_DIR` to the absolute path of this skill's `scripts/` directory. @@ -58,10 +58,10 @@ Before proceeding past Stage 0, confirm: ### Which stages need hardware -| Stages | Mode | Needs compiled model + Neuron device? | -|--------|------|----------------------------------------| -| 0, 2, 3, 4 | CPU (`NXD_CPU_MODE=1`, TP=1) | No | -| 1, 5, 6, 7 | Device | **Yes** | +| Stages | Mode | Needs compiled model + Neuron device? | +| ---------- | ---------------------------- | ------------------------------------- | +| 0, 2, 3, 4 | CPU (`NXD_CPU_MODE=1`, TP=1) | No | +| 1, 5, 6, 7 | Device | **Yes** | Stage 1 is a **device** stage despite its low stage number — `run_stage1.py` calls the adapter's `device_inference()` and requires `COMPILED_MODEL_PATH`. Do not plan on running @@ -202,10 +202,10 @@ See [references/report-template.md](references/report-template.md) for the struc R = ||target - source_fp32||_F / (||source_lowprec - source_fp32||_F + ε) ``` -| R | Meaning | -|---|---------| -| ≈ 1.0 | Healthy — matches precision baseline | -| > 1.2 | Bug — excess divergence | +| R | Meaning | +| ----- | --------------------------------------- | +| ≈ 1.0 | Healthy — matches precision baseline | +| > 1.2 | Bug — excess divergence | | < 1.0 | Over-precision — extra `.float()` calls | ## Verdict diff --git a/skills/neuron-framework-equivalence/STAGE0.md b/skills/neuron-framework-equivalence/STAGE0.md index 0cec613..ba1ee95 100644 --- a/skills/neuron-framework-equivalence/STAGE0.md +++ b/skills/neuron-framework-equivalence/STAGE0.md @@ -39,7 +39,7 @@ The script generates all 8 tree artifacts and prints both trees. The target tree **vLLM-Neuron targets — automatic environment check (by design, not a bug):** `run_stage0.py` builds the target tree through the stack adapter, and `get_adapter()` runs `check_environment()` first. For the `vllm_neuron` adapter this **deliberately exits early with an `EnvironmentError`** if any of these hold: -1. `vllm_neuron` is not importable — fix by prepending `{VLLM_NEURON_DIR}` to `PYTHONPATH` (see SKILL.md → *vLLM-Neuron Targets*). +1. `vllm_neuron` is not importable — fix by prepending `{VLLM_NEURON_DIR}` to `PYTHONPATH` (see SKILL.md → _vLLM-Neuron Targets_). 2. The installed `vllm` framework is off the pinned `0.24` line. 3. The installed `vllm-neuron` plugin is off the pinned `0.24` line. @@ -52,6 +52,7 @@ This is intentional fail-fast behavior — it prevents a cryptic mid-stage `Asse Compare the two printed trees and build `{EXP_DIR}/component_mapping.json`. Rules: + 1. Start from leaf modules, work upward 2. Inspect source code to verify semantic equivalence (same name ≠ same function) 3. Support one-to-one and one-to-many mappings (e.g., fused QKV → split Q/K/V) @@ -71,6 +72,7 @@ python3 {SCRIPTS_DIR}/detect_class_divergence.py \ ``` The script detects three patterns: + 1. **Factory functions** (`get_rmsnorm_cls()`, `get_attn_cls()`) that return different classes based on `NXD_CPU_MODE` 2. **Conditional assignments** (`self.norm = ClassA() if cpu else ClassB()`) 3. **NKI kernel imports** (`CustomRMSNorm` on device vs `LlamaRMSNorm` on CPU) @@ -85,11 +87,11 @@ For mapped components with shape/layout differences (fused operators, transposed ## Troubleshooting -| Issue | Cause | Solution | -|-------|-------|----------| -| `'NoneType' has no attribute 'windowed_context_encoding_size'` | Config validation requires `neuron_config` | Pass `neuron_config` to `from_pretrained()` | -| `intra_layer_model parallel group is not initialized` | Parallel state not initialized | `run_stage0.py` handles this — if running manually, call `init_process_group("gloo")` then `initialize_model_parallel(tp=1)` | -| `Please initialize parallel processing via 'torchrun'` | `world_size > 1` without torchrun | Use `tp_degree=1, world_size=1` for structure inspection | -| `No module named 'modeling_xxx'` | Missing sys.path entry | Check `--target-module-file` path is correct | -| HF model type not recognized | Transformers version too old | Check `transformers.__version__` supports the model | -| `from_pretrained` fails on config class | Config class uses non-standard constructor | Script falls back to two-arg constructor automatically | +| Issue | Cause | Solution | +| -------------------------------------------------------------- | ------------------------------------------ | ---------------------------------------------------------------------------------------------------------------------------- | +| `'NoneType' has no attribute 'windowed_context_encoding_size'` | Config validation requires `neuron_config` | Pass `neuron_config` to `from_pretrained()` | +| `intra_layer_model parallel group is not initialized` | Parallel state not initialized | `run_stage0.py` handles this — if running manually, call `init_process_group("gloo")` then `initialize_model_parallel(tp=1)` | +| `Please initialize parallel processing via 'torchrun'` | `world_size > 1` without torchrun | Use `tp_degree=1, world_size=1` for structure inspection | +| `No module named 'modeling_xxx'` | Missing sys.path entry | Check `--target-module-file` path is correct | +| HF model type not recognized | Transformers version too old | Check `transformers.__version__` supports the model | +| `from_pretrained` fails on config class | Config class uses non-standard constructor | Script falls back to two-arg constructor automatically | diff --git a/skills/neuron-framework-equivalence/STAGE2.md b/skills/neuron-framework-equivalence/STAGE2.md index 36d4ec8..25367a3 100644 --- a/skills/neuron-framework-equivalence/STAGE2.md +++ b/skills/neuron-framework-equivalence/STAGE2.md @@ -66,6 +66,7 @@ If a leaf passes but its composite fails → bug is in composition logic. Use [templates/conftest_template.py](templates/conftest_template.py) for scaffolding and [templates/test_template.py](templates/test_template.py) for the per-component pattern. Key rules: + 1. **Shared FP32 weights** across all three module instances 2. **`nn.Parameter()` replacement** for `ColumnParallelLinear` (not `copy_()` — it silently preserves dtype) 3. **`.eval()` mode** on Neuron modules with `pad=True` @@ -84,19 +85,20 @@ The runner imports test files in-process, executes `test_*` functions, captures ## Interpreting Results -| R-ratio | Meaning | -|---------|---------| -| ≈ 1.0 | Port matches precision baseline. No bug. | +| R-ratio | Meaning | +| ------------- | -------------------------------------------------------------------- | +| ≈ 1.0 | Port matches precision baseline. No bug. | | 1.0 < R < 1.2 | Slight excess — may be TP rounding or kernel difference. Acceptable. | -| R >> 1.2 | Porting bug. Proceed to Stage 3. | -| R < 1.0 | Over-precision. Check for extra `.float()` calls. | -| R >> 10 | Missing algorithm (YaRN, MoE routing, etc.) | -| R >> 100 | Completely wrong computation. | +| R >> 1.2 | Porting bug. Proceed to Stage 3. | +| R < 1.0 | Over-precision. Check for extra `.float()` calls. | +| R >> 10 | Missing algorithm (YaRN, MoE routing, etc.) | +| R >> 100 | Completely wrong computation. | ## Visual Analysis (QQ Plots) Beyond the R-ratio, examine the error distribution via QQ plots and histograms. See [references/example_plots/](references/example_plots/) for examples: + - `positive_samples/` — PASS cases: error distributions overlap, QQ plot on 45° line - `negative_samples/` — FAIL cases: divergent distributions, off-diagonal QQ plots @@ -116,9 +118,9 @@ neuron_linear.weight = torch.nn.Parameter(weight.to(torch.bfloat16)) # neuron_linear.weight.dtype == torch.bfloat16 ``` -| Module type | Method | Why | -|-------------|--------|-----| -| `nn.Linear`, `nn.Embedding`, `nn.LayerNorm` | `copy_()` is fine | These don't enforce a fixed dtype on their weight tensors | +| Module type | Method | Why | +| --------------------------------------------------------------------------------------------- | ----------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `nn.Linear`, `nn.Embedding`, `nn.LayerNorm` | `copy_()` is fine | These don't enforce a fixed dtype on their weight tensors | | `ColumnParallelLinear`, `RowParallelLinear`, or any module with `dtype=torch.float32` default | **Must use `nn.Parameter()` replacement** | `copy_()` silently converts bf16→fp32. The forward pass then fails with `mat1 and mat2 must have the same dtype` or silently computes in fp32 producing R-ratios in the 2–5x range. | ## Self-Reflection: Test Correctness Verification @@ -137,11 +139,11 @@ After writing all tests and before declaring results, verify each test against t This table complements the "Interpreting Results" table above. It maps component categories to expected R-ratio behaviors: -| Component type | Expected R | Guidance | -|---|---|---| -| Leaf components (norm, embedding, linear, lm_head) | ≈ 1.0 | If > 1.2, formula-level bug in the port | -| Components with missing features (e.g., RoPE without YaRN) | >> 10 | Target uses a different or incomplete algorithm | -| Components with logic bugs (e.g., MoE routing ignored) | >> 100 | Target has incorrect forward-pass logic | -| Components with known precision differences | 1.2–2.0 | Document the difference; consider `tolerance_ratio=2.0` | +| Component type | Expected R | Guidance | +| ---------------------------------------------------------- | ---------- | ------------------------------------------------------- | +| Leaf components (norm, embedding, linear, lm_head) | ≈ 1.0 | If > 1.2, formula-level bug in the port | +| Components with missing features (e.g., RoPE without YaRN) | >> 10 | Target uses a different or incomplete algorithm | +| Components with logic bugs (e.g., MoE routing ignored) | >> 100 | Target has incorrect forward-pass logic | +| Components with known precision differences | 1.2–2.0 | Document the difference; consider `tolerance_ratio=2.0` | **Key diagnostic:** When a leaf component PASSES but a composite FAILS, the bug is in the composition logic (e.g., routing, weight application, residual connections) — not in the individual subcomponents. diff --git a/skills/neuron-framework-equivalence/STAGE3.md b/skills/neuron-framework-equivalence/STAGE3.md index b012850..91d9017 100644 --- a/skills/neuron-framework-equivalence/STAGE3.md +++ b/skills/neuron-framework-equivalence/STAGE3.md @@ -17,11 +17,11 @@ The earliest step-pattern point is the primary fault candidate. ## Root-Cause Classification -| R magnitude | Likely cause | Examples | -|-------------|-------------|---------| -| R >> 10 | Missing algorithm or wrong formula | YaRN scaling absent from RoPE, MoE routing ignored, wrong masking | -| 1.2 < R < 3 | Precision ordering or missing multiplier | Variance in BF16 instead of FP32, attention scaling omitted | -| R < 1 | Over-precision (unintended FP32 upcast) | Extra `.float()` call not in reference | +| R magnitude | Likely cause | Examples | +| ----------- | ---------------------------------------- | ----------------------------------------------------------------- | +| R >> 10 | Missing algorithm or wrong formula | YaRN scaling absent from RoPE, MoE routing ignored, wrong masking | +| 1.2 < R < 3 | Precision ordering or missing multiplier | Variance in BF16 instead of FP32, attention scaling omitted | +| R < 1 | Over-precision (unintended FP32 upcast) | Extra `.float()` call not in reference | ## Output diff --git a/skills/neuron-framework-equivalence/STAGE4.md b/skills/neuron-framework-equivalence/STAGE4.md index b9d0853..8f4809e 100644 --- a/skills/neuron-framework-equivalence/STAGE4.md +++ b/skills/neuron-framework-equivalence/STAGE4.md @@ -10,12 +10,12 @@ Fix failing components with standalone monkey patches. Use this table for immediate triage when a component fails. (`run_stage3.py` automates this classification, but this table is essential for manual debugging context.) -| R-ratio range | Likely cause | Examples | -|---|---|---| -| 100x+ | Formula-level bug — wrong function or missing operation | YaRN scaling absent from RoPE, MoE routing ignored | -| 1–3x | Precision issue — dtype casting or operation ordering | Variance computed in BF16 instead of FP32, scaling applied after cast | -| < 1.0 | Over-precision — extra `.float()` calls not in reference | Target artificially closer to FP32 than the BF16 baseline | -| 1000x+ | Routing problem — wrong code path or CPU class running device logic | Factory function returning wrong class, dispatch path mismatch | +| R-ratio range | Likely cause | Examples | +| ------------- | ------------------------------------------------------------------- | --------------------------------------------------------------------- | +| 100x+ | Formula-level bug — wrong function or missing operation | YaRN scaling absent from RoPE, MoE routing ignored | +| 1–3x | Precision issue — dtype casting or operation ordering | Variance computed in BF16 instead of FP32, scaling applied after cast | +| < 1.0 | Over-precision — extra `.float()` calls not in reference | Target artificially closer to FP32 than the BF16 baseline | +| 1000x+ | Routing problem — wrong code path or CPU class running device logic | Factory function returning wrong class, dispatch path mismatch | ## Workflow @@ -61,7 +61,7 @@ Repeat Steps 1-5 until all R < τ_R. If a patch fixes one module but breaks a do ## Detailed Debugging Guides - [references/cpu-component-debugging.md](references/cpu-component-debugging.md) — Full CPU debugging workflow with pitfalls, examples, and patterns from real debugging sessions -- [references/device-component-debugging.md](references/device-component-debugging.md) — Device-specific debugging: XLA-compatible patch patterns (SPMDRank, index_select, _reduce), pre_shard_weights_hook injection, compiler log analysis (`log-neuron-cc.txt`), escalation to compiler debugging +- [references/device-component-debugging.md](references/device-component-debugging.md) — Device-specific debugging: XLA-compatible patch patterns (SPMDRank, index_select, \_reduce), pre_shard_weights_hook injection, compiler log analysis (`log-neuron-cc.txt`), escalation to compiler debugging - [references/device-e2e-debugging.md](references/device-e2e-debugging.md) — Device E2E debugging: 1-layer isolation technique, fix-compile-verify cycle, full model validation - [references/cpu-e2e-debugging.md](references/cpu-e2e-debugging.md) — CPU E2E debugging: TP=1 FP32 baseline, mp.spawn patch inheritance, weight sharding pipeline, bias restoration - [references/debugging-case-study-gptoss.md](references/debugging-case-study-gptoss.md) — Complete worked example from GPT-OSS 20B with specific error ratios, root causes, and patches diff --git a/skills/neuron-framework-equivalence/STAGE5.md b/skills/neuron-framework-equivalence/STAGE5.md index 219a5d6..b17b975 100644 --- a/skills/neuron-framework-equivalence/STAGE5.md +++ b/skills/neuron-framework-equivalence/STAGE5.md @@ -16,6 +16,7 @@ Verify the assembled model with real weights under teacher forcing. Covers Stage ## What the Design Doc Requires For each prompt and each **teacher-forced position t**: + - R-ratio on output logits (three-tensor: source FP32, source BF16, target BF16) - Cosine similarity cos(v_source, v_target) ≥ θ (Condition B) - KL divergence D_KL(P_source ∥ P_target) ≤ δ (Condition C) @@ -53,6 +54,7 @@ PYTHONPATH={SCRIPTS_DIR} python3 {SCRIPTS_DIR}/run_teacher_forced_comparison.py For per-layer intermediate comparison (beyond final logits), use `templates/diagnostic_forward_template.py` which provides hook-based capture at named module boundaries. Key rules: + - Strip the `model.` prefix from source model names to get consistent names across both sides - Device captures pad input to full `seq_len` — slice to match: `device_tensor[:, :ref_seq_len, :]` - `self_attn` on device captures `cos_cache`, not hidden_states — use `post_attention_layernorm` as attention quality proxy diff --git a/skills/neuron-framework-equivalence/STAGE6.md b/skills/neuron-framework-equivalence/STAGE6.md index 2a64006..35dc097 100644 --- a/skills/neuron-framework-equivalence/STAGE6.md +++ b/skills/neuron-framework-equivalence/STAGE6.md @@ -9,12 +9,14 @@ See [STAGE5.md](STAGE5.md) for the combined instructions. ## Condition B: Semantic Consistency Per-position cosine similarity: `cos(v_source, v_target) ≥ θ` + - p5 percentile must be ≥ θ (default 0.95) - Tail fraction below θ must be ≤ ρ (default 0.05) ## Condition C: Distributional Equivalence Per-position KL divergence: `D_KL(P_source ∥ P_target) ≤ δ` + - p95 percentile must be ≤ δ - Maximum must be ≤ δ_max diff --git a/skills/neuron-framework-equivalence/references/adapter-contract.md b/skills/neuron-framework-equivalence/references/adapter-contract.md index 9e7f2e1..5f47842 100644 --- a/skills/neuron-framework-equivalence/references/adapter-contract.md +++ b/skills/neuron-framework-equivalence/references/adapter-contract.md @@ -17,54 +17,54 @@ The equivalence methodology (R-ratio, 3-tensor, 8 stages) is platform-agnostic. Set up distributed process groups for CPU-mode testing. Called once before `create_model()`. -| Stack | Implementation | -|---|---| -| NxDI | `torch.distributed.init_process_group("gloo")` + `neuronx_distributed.parallel_state.initialize_model_parallel(tp)` | +| Stack | Implementation | +| -------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| NxDI | `torch.distributed.init_process_group("gloo")` + `neuronx_distributed.parallel_state.initialize_model_parallel(tp)` | | vLLM-Neuron (0.24.0) | `torch.distributed.init_process_group("gloo")` + `vllm_neuron.parallel.neuron_parallel_state.initialize_neuron_parallel_state(tp_global_ranks=..., local_rank=0)` — delegates `VllmConfig`/`set_current_vllm_config`/`initialize_model_parallel`. Do NOT call vLLM's `init_distributed_environment` first (pre-creates `_WORLD`). | ### `create_model(target_module_file, target_class_name, target_config_name, hf_model_path)` Instantiate the target model in CPU mode. Returns an unweighted model for tree building and component mapping. -| Stack | Implementation | -|---|---| -| NxDI | `NeuronConfig(on_cpu=True)` → `ConfigClass.from_pretrained(path, neuron_config=...)` → `InnerClass(config)` | +| Stack | Implementation | +| ----------- | ---------------------------------------------------------------------------------------------------------------- | +| NxDI | `NeuronConfig(on_cpu=True)` → `ConfigClass.from_pretrained(path, neuron_config=...)` → `InnerClass(config)` | | vLLM-Neuron | `NXD_CPU_MODE=1` → `AutoConfig.from_pretrained()` → `ConfigClass.from_configs(hf_config)` → `ModelClass(config)` | ### `load_weights(model, hf_model_path, dtype)` Load HuggingFace weights into the target model with all stack-specific transforms. -| Stack | Key transforms | -|---|---| -| NxDI | `model.load(path)` or standard state_dict loading | +| Stack | Key transforms | +| ----------- | -------------------------------------------------------------------------------------------------------------------------------- | +| NxDI | `model.load(path)` or standard state_dict loading | | vLLM-Neuron | All linear weights transposed (`.t()`), Q/K/V fused (`cat([Q.t(), K.t(), V.t()], dim=-1)`), weight names `_weight` not `.weight` | ### `forward(model, input_ids)` Run a forward pass, return logits as float32 tensor. -| Stack | Signature handled internally | -|---|---| -| NxDI | `model(input_ids, attention_mask, position_ids)` → logits | +| Stack | Signature handled internally | +| ----------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | +| NxDI | `model(input_ids, attention_mask, position_ids)` → logits | | vLLM-Neuron | `model(input_ids, positions, attn_metadata, sampling_positions)` → logits. Adapter constructs `attn_metadata` dict-of-dicts, allocates KV caches, passes `sampling_positions`. | ### `device_inference(model_id, tp_size, prompts, max_tokens)` Run inference on actual Neuron hardware. Returns list of `{"text": ..., "tokens": [...]}`. -| Stack | Implementation | -|---|---| -| NxDI | `HuggingFaceGenerationAdapter` + compiled-model loading via `scripts/nxdi_compiled_loader.py` | -| vLLM-Neuron | `vllm.LLM(model=id, tensor_parallel_size=tp)` + `SamplingParams(temperature=0.0)` | +| Stack | Implementation | +| ----------- | --------------------------------------------------------------------------------------------- | +| NxDI | `HuggingFaceGenerationAdapter` + compiled-model loading via `scripts/nxdi_compiled_loader.py` | +| vLLM-Neuron | `vllm.LLM(model=id, tensor_parallel_size=tp)` + `SamplingParams(temperature=0.0)` | ## Environment Check (`check_environment()`) Optional, no-op by default. `get_adapter()` calls it right after constructing the adapter. Override it to pin dependency versions and fail fast with a clear, actionable message + early exit instead of crashing mid-stage. -| Stack | Implementation | -|---|---| -| NxDI | inherits no-op default | +| Stack | Implementation | +| ----------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| NxDI | inherits no-op default | | vLLM-Neuron | verifies `vllm_neuron` is importable and that both `vllm` (`PINNED_VLLM_VERSION`) and the `vllm-neuron` plugin (`PINNED_VLLM_NEURON_VERSION`) are on the pinned `0.24` line; raises `EnvironmentError` otherwise | Pass `check_environment=False` to `get_adapter()` only for pure introspection that exercises no stack APIs. @@ -72,6 +72,7 @@ Pass `check_environment=False` to `get_adapter()` only for pure introspection th ## Auto-Detection If `target_stack` is not specified, the registry reads the target modeling file's imports: + - `from vllm_neuron.*` → `vllm_neuron` adapter - `from neuronx_distributed_inference.*` → `nxdi` adapter - Neither → defaults to `nxdi` @@ -121,7 +122,7 @@ Full prompt-level diagnosis combining multiple plugins. Wraps `accuracy_debugger ## Existing Adapters -| Adapter | File | Stack | Diagnostics | -|---|---|---|---| -| `NxDIAdapter` | `scripts/adapters/nxdi.py` | NeuronX Distributed Inference | Core 5 only | -| `VLLMNeuronAdapter` | `scripts/adapters/vllm_neuron.py` | vLLM-Neuron | Core 5 + all diagnostic methods | +| Adapter | File | Stack | Diagnostics | +| ------------------- | --------------------------------- | ----------------------------- | ------------------------------- | +| `NxDIAdapter` | `scripts/adapters/nxdi.py` | NeuronX Distributed Inference | Core 5 only | +| `VLLMNeuronAdapter` | `scripts/adapters/vllm_neuron.py` | vLLM-Neuron | Core 5 + all diagnostic methods | diff --git a/skills/neuron-framework-equivalence/references/cpu-component-debugging.md b/skills/neuron-framework-equivalence/references/cpu-component-debugging.md index d8dc3d3..1bd073d 100644 --- a/skills/neuron-framework-equivalence/references/cpu-component-debugging.md +++ b/skills/neuron-framework-equivalence/references/cpu-component-debugging.md @@ -28,7 +28,7 @@ When a component equivalence test (from the `component-testing` skill) fails **o 3. **Note the error magnitude.** - 100x+ → formula/algorithm mismatch - 1.2–3x → precision ordering issue - - < 1.0 → patch computes at *higher* precision than the reference (dtype over-precision) + - < 1.0 → patch computes at _higher_ precision than the reference (dtype over-precision) ### Phase 2: Compare Implementations Side-by-Side @@ -41,14 +41,14 @@ When a component equivalence test (from the `component-testing` skill) fails **o 6. **Identify the root cause category:** - | Category | Symptoms | Example | - |----------|----------|---------| - | **Missing algorithm** | Error ratio 50x–1000x+. Target uses a simpler formula. | YaRN scaling omitted from RoPE | - | **Missing multiplier/scaling** | Error ratio 1.3x–2x. Values are proportionally off. | `attention_scaling` factor not applied | - | **Config parameter gap** | Target config missing a field the algorithm needs. | `original_max_position_embeddings` absent | - | **Precision ordering** | Error ratio 1.2x–15x. One output (e.g., cos) passes but another (e.g., sin) fails. | Scaling applied after bf16 cast instead of before | - | **Shape mismatch** | Comparison invalid. Shapes differ between reference and target. | `[bs, seq, dim/2]` vs `[bs, seq, dim]` | - | **Routing/logic ignored** | Error ratio 1000x+. Target produces structurally different output. | MoE routing weights ignored | + | Category | Symptoms | Example | + | ------------------------------ | ---------------------------------------------------------------------------------- | ------------------------------------------------- | + | **Missing algorithm** | Error ratio 50x–1000x+. Target uses a simpler formula. | YaRN scaling omitted from RoPE | + | **Missing multiplier/scaling** | Error ratio 1.3x–2x. Values are proportionally off. | `attention_scaling` factor not applied | + | **Config parameter gap** | Target config missing a field the algorithm needs. | `original_max_position_embeddings` absent | + | **Precision ordering** | Error ratio 1.2x–15x. One output (e.g., cos) passes but another (e.g., sin) fails. | Scaling applied after bf16 cast instead of before | + | **Shape mismatch** | Comparison invalid. Shapes differ between reference and target. | `[bs, seq, dim/2]` vs `[bs, seq, dim]` | + | **Routing/logic ignored** | Error ratio 1000x+. Target produces structurally different output. | MoE routing weights ignored | 7. **Run numerical diagnostics inside Docker** to confirm. Compare intermediate values (e.g., `inv_freq`, `freqs`, `cos`, `sin`) between reference and target: @@ -186,10 +186,10 @@ The patch must use the **same dtype strategy** as the reference. Read every `.fl **Two failure modes:** -| Mode | Symptom | Cause | -|------|---------|-------| -| **Under-precision** | Error ratio 1.2–2x | Target computes in bf16 where reference uses fp32 | -| **Over-precision** | Error ratio < 1.0 | Patch adds `.float()` calls the reference doesn't have | +| Mode | Symptom | Cause | +| ------------------- | ------------------ | ------------------------------------------------------ | +| **Under-precision** | Error ratio 1.2–2x | Target computes in bf16 where reference uses fp32 | +| **Over-precision** | Error ratio < 1.0 | Patch adds `.float()` calls the reference doesn't have | Error ratio < 1.0 is a bug: it means the target is closer to the fp32 ground truth than the bf16 reference is, which can only happen if the patch computes at higher precision. @@ -219,7 +219,7 @@ output = torch.einsum('nei,eih->neh', x, w) **Case C — Reference upcasts for specific operations only:** -YaRN RoPE: HF computes `inv_freq @ position_ids` in fp32, applies `attention_scaling` in fp32, then casts to input dtype. Applying scaling *after* a bf16 cast caused sin error ratio 13.7x. Fix: replicate the same fp32 → scale → cast sequence. +YaRN RoPE: HF computes `inv_freq @ position_ids` in fp32, applies `attention_scaling` in fp32, then casts to input dtype. Applying scaling _after_ a bf16 cast caused sin error ratio 13.7x. Fix: replicate the same fp32 → scale → cast sequence. ### Phase 5: Integrate and Verify @@ -268,9 +268,9 @@ When a patch doesn't produce the expected result, verify in order: For each fixed component, produce: -| File | Location | Purpose | -|------|----------|---------| -| `.py` | `{EXP_DIR}/patches/` | Standalone monkey-patch file | +| File | Location | Purpose | +| -------------------------------- | -------------------------- | ---------------------------------------------------------------- | +| `.py` | `{EXP_DIR}/patches/` | Standalone monkey-patch file | | Updated `test_NN_.py` | `experiments//tests/` | Test imports and applies patch, expected-failure markers removed | --- @@ -280,14 +280,17 @@ For each fixed component, produce: **Failing test:** `test_02_rotary_emb.py` — error ratio 130x (cos) and 133x (sin). **Root cause diagnosis:** + - HF `GptOssRotaryEmbedding` uses `ROPE_INIT_FUNCTIONS["yarn"]` which blends interpolated/extrapolated `inv_freq` and applies `attention_scaling = 1.3466` to cos/sin. - Neuron `NeuronGptOssRotaryEmbedding` wraps framework's basic `RotaryEmbedding(dim, max_pos, base)` which only computes standard RoPE. The `rope_scaling` config is stored but never used. **Patch:** `yarn_rotary_patch.py` — `apply_yarn_rotary_patch()`: + 1. `_patched_init`: computes YaRN `inv_freq` and `attention_scaling`, stores as `self._yarn_inv_freq` and `self.attention_scaling` 2. `_patched_forward`: computes rotary embeddings in fp32 using YaRN inv_freq, applies attention_scaling before dtype cast, emits `cat((freqs, freqs), dim=-1)` format **Pitfalls encountered during development:** + 1. Neuron config's `rope_scaling` dict was missing `original_max_position_embeddings` (4096), so the correction range was computed against 131072, making YaRN inv_freq nearly identical to default. Fixed by deriving: `original = max_position_embeddings / factor`. 2. Framework's `register_buffer("inv_freq", None)` resisted direct assignment. Fixed by storing `_yarn_inv_freq` on the wrapper instead. 3. Applying `attention_scaling` after the framework cast cos/sin to bf16 caused sin error ratio of 13.7x. Fixed by reimplementing forward inline, applying scaling in fp32 before the cast. @@ -315,11 +318,13 @@ For each fixed component, produce: **Failing test:** `test_06_mlp_moe.py` — error ratio 1809x. `test_05_experts.py` — error ratio 1774x. **Root cause:** Three bugs: + 1. `NeuronGptOssExperts.forward()` ignores `router_indices` and `routing_weights` — all tokens go through all 32 experts instead of top-4. 2. `NeuronGptOssMLP.forward()` uses `router_scores.sum(dim=-1)` as multiplier — softmax scores sum to ~1.0, making it a no-op. 3. Down-proj weight layout mismatch (`i*E+e` vs `e*I+i` indexing) and per-expert biases collapsed to single sum. **Patch:** `mlp_moe_patch.py`: + 1. Extracts per-expert weights from flattened RowParallelLinear via reshape + permute 2. Computes per-expert down projections via `einsum('nei,eih->neh')` 3. Applies routing weights per expert, then sums across experts diff --git a/skills/neuron-framework-equivalence/references/cpu-e2e-debugging.md b/skills/neuron-framework-equivalence/references/cpu-e2e-debugging.md index 05e823d..fd7d955 100644 --- a/skills/neuron-framework-equivalence/references/cpu-e2e-debugging.md +++ b/skills/neuron-framework-equivalence/references/cpu-e2e-debugging.md @@ -20,11 +20,11 @@ Verify the Neuron model's forward-pass logic is correct at FP32 with TP=1. **Pass criterion:** `error_ratio < 1.2` -| Tensor | Role | -|--------|------| -| HF FP32 | Ground truth | -| HF BF16 | Baseline dtype error | -| Neuron BF16 | Target | +| Tensor | Role | +| ----------- | -------------------- | +| HF FP32 | Ground truth | +| HF BF16 | Baseline dtype error | +| Neuron BF16 | Target | ``` error_ratio = rel_fro_norm(Neuron_BF16, HF_FP32) / rel_fro_norm(HF_BF16, HF_FP32) @@ -35,6 +35,7 @@ error_ratio = rel_fro_norm(Neuron_BF16, HF_FP32) / rel_fro_norm(HF_BF16, HF_FP32 ## Phase 3: TP>1 Validation **Pass criteria:** + - Fast test: `rel_fro(TP=1, TP=N) < 1e-2` - Real weights: 3-tensor `error_ratio < 1.2` at each TP degree @@ -68,6 +69,7 @@ This was the **primary root cause** of TP>1 divergence for GPT-OSS 20B. Use a tiny 1-layer model with random weights for 5-10 second iteration cycles. See `templates/fast_tp_equiv_test_template.py`. **TINY_CONFIG design rules:** + - `num_key_value_heads` must be divisible by all TP degrees you test - `num_local_experts` must be divisible by all TP degrees - `num_hidden_layers=1` keeps it fast while exercising the full forward path @@ -80,37 +82,37 @@ The **first checkpoint that shows FAIL** (rel_fro >= 1e-1) localizes the bug to ### Common TP>1 Issues -| Issue | Symptom | Fix | -|-------|---------|-----| -| **Patches not applied in worker** | `patched_flag=0` in TP>1; all intermediates diverge | Add `apply_all_patches()` as first action in `_tp_worker` | -| **Weight layout mismatch** | MoE output FAIL; all later stages FAIL | `fix_()` before `get_sharded_checkpoint` | -| **Biases removed by sharding** | Attention output diverges; biases are zero | Manual bias restoration with TP-aware slicing | -| **CONVERT_TO_MHA** | KV bias shape mismatch when `tp_degree % kv_heads != 0` | Replicate bias via `repeat_interleave` then shard | -| **Parameter not TP-sharded** | Full-size parameter vs local head count → shape error | Manual TP slicing: `param[rank*local:(rank+1)*local]` | -| **Missing all-reduce** | MoE output is 1/N of correct value | Add `torch.distributed.all_reduce()` after partial computation | -| **Port conflict** | `EADDRINUSE` error | Use different `MASTER_PORT` for TP=1 (8080) and TP>1 (29501) | +| Issue | Symptom | Fix | +| --------------------------------- | ------------------------------------------------------- | -------------------------------------------------------------- | +| **Patches not applied in worker** | `patched_flag=0` in TP>1; all intermediates diverge | Add `apply_all_patches()` as first action in `_tp_worker` | +| **Weight layout mismatch** | MoE output FAIL; all later stages FAIL | `fix_()` before `get_sharded_checkpoint` | +| **Biases removed by sharding** | Attention output diverges; biases are zero | Manual bias restoration with TP-aware slicing | +| **CONVERT_TO_MHA** | KV bias shape mismatch when `tp_degree % kv_heads != 0` | Replicate bias via `repeat_interleave` then shard | +| **Parameter not TP-sharded** | Full-size parameter vs local head count → shape error | Manual TP slicing: `param[rank*local:(rank+1)*local]` | +| **Missing all-reduce** | MoE output is 1/N of correct value | Add `torch.distributed.all_reduce()` after partial computation | +| **Port conflict** | `EADDRINUSE` error | Use different `MASTER_PORT` for TP=1 (8080) and TP>1 (29501) | ### Bias Restoration (TP-Aware) `get_sharded_checkpoint` removes biases it considers "redundant". Three cases for restoration: -| Bias size vs local heads | Action | -|-------------------------|--------| -| Equal | Use as-is | -| Greater (full-size bias) | Chunk and shard: `bias.chunk(tp_degree)[rank]` | -| Less (CONVERT_TO_MHA) | Replicate via `repeat_interleave(repeats)` then chunk-shard | +| Bias size vs local heads | Action | +| ------------------------ | ----------------------------------------------------------- | +| Equal | Use as-is | +| Greater (full-size bias) | Chunk and shard: `bias.chunk(tp_degree)[rank]` | +| Less (CONVERT_TO_MHA) | Replicate via `repeat_interleave(repeats)` then chunk-shard | --- ## Phase 5: Full Model Validation -| Test | Metric | Threshold | -|------|--------|-----------| -| FP32 Direct TP=1 | rel_fro_norm | < 1e-5 | -| 3-Tensor BF16 TP=1 | error_ratio | < 1.2 | -| 3-Tensor BF16 TP=4 | error_ratio | < 1.2 | -| 3-Tensor BF16 TP=8 | error_ratio | < 1.2 | -| Token coherence | All same next token | Match HF FP32 | +| Test | Metric | Threshold | +| ------------------ | ------------------- | ------------- | +| FP32 Direct TP=1 | rel_fro_norm | < 1e-5 | +| 3-Tensor BF16 TP=1 | error_ratio | < 1.2 | +| 3-Tensor BF16 TP=4 | error_ratio | < 1.2 | +| 3-Tensor BF16 TP=8 | error_ratio | < 1.2 | +| Token coherence | All same next token | Match HF FP32 | ### Key Insight from GPT-OSS @@ -139,6 +141,7 @@ When `tp_degree % num_kv_heads != 0`, KV heads are replicated to match Q heads. ### Pitfall 5: Parameters Not Marked tensor_model_parallel (MEDIUM) Some parameters (e.g., attention sinks) stay full-size at TP>1 but `self.num_heads` is the LOCAL count. Add TP-aware slicing: + ```python if total_heads != self.num_heads: tp_rank = parallel_state.get_tensor_model_parallel_rank() @@ -148,6 +151,7 @@ if total_heads != self.num_heads: ### Pitfall 6: Missing All-Reduce After Partial Computation (MEDIUM) At TP>1, partial results must be summed across ranks: + ```python if tp_size > 1: torch.distributed.all_reduce(output, group=get_tensor_model_parallel_group()) diff --git a/skills/neuron-framework-equivalence/references/debug-orchestration.md b/skills/neuron-framework-equivalence/references/debug-orchestration.md index f8aa852..da39746 100644 --- a/skills/neuron-framework-equivalence/references/debug-orchestration.md +++ b/skills/neuron-framework-equivalence/references/debug-orchestration.md @@ -40,26 +40,26 @@ Stage 5: Re-run validation for clean report ## Decision Table -| Symptom | Go to | -|---------|-------| -| Component test fails on CPU | Stage 1: cpu-component-debugging | -| Component passes CPU, fails device | Stage 2: device-component-debugging | -| Component code matches reference, device still diverges | Suspected compiler issue | -| CPU E2E fails at TP>1 | Stage 3: cpu-e2e-debugging (check mp.spawn, weight layout) | -| Device E2E diverges from CPU E2E | Stage 4: device-e2e-debugging (1-layer isolation) | -| Compilation fails | Suspected compiler issue | -| All stages pass | Re-run validation, generate clean report | +| Symptom | Go to | +| ------------------------------------------------------- | ---------------------------------------------------------- | +| Component test fails on CPU | Stage 1: cpu-component-debugging | +| Component passes CPU, fails device | Stage 2: device-component-debugging | +| Component code matches reference, device still diverges | Suspected compiler issue | +| CPU E2E fails at TP>1 | Stage 3: cpu-e2e-debugging (check mp.spawn, weight layout) | +| Device E2E diverges from CPU E2E | Stage 4: device-e2e-debugging (1-layer isolation) | +| Compilation fails | Suspected compiler issue | +| All stages pass | Re-run validation, generate clean report | --- ## Stage Gates -| Stage | Pass criterion | -|-------|---------------| -| 1. CPU Components | All component R ≤ 1.2 on CPU | -| 2. Device Components | All component R ≤ 1.2 on device | -| 3. CPU E2E | TP=1 FP32 rel_fro < 1e-5, TP=1 BF16 R < 1.2, TP>1 BF16 R < 1.2 | -| 4. Device E2E | Full model R ≤ 1.2, top-1 token match, coherent 20-token generation | +| Stage | Pass criterion | +| -------------------- | ------------------------------------------------------------------- | +| 1. CPU Components | All component R ≤ 1.2 on CPU | +| 2. Device Components | All component R ≤ 1.2 on device | +| 3. CPU E2E | TP=1 FP32 rel_fro < 1e-5, TP=1 BF16 R < 1.2, TP>1 BF16 R < 1.2 | +| 4. Device E2E | Full model R ≤ 1.2, top-1 token match, coherent 20-token generation | --- @@ -85,12 +85,12 @@ For Stage 4 debugging, which can span dozens of turns per failing component, use ## Reference: GPT-OSS 20B Timeline -| Stage | Duration | What was found | -|-------|----------|----------------| -| CPU Components | 2 weeks | 4 patches: rmsnorm precision, YaRN rotary, MoE routing + weight layout, attention sinks | -| CPU E2E | 1 week | mp.spawn patch inheritance, weight layout fix, bias restoration | -| Device E2E | 2 weeks | SPMDRank for ParallelEmbedding (590→127), windowed attention path (127→0.97) | -| **Final result** | | R = 1.0005, top-1 match, coherent generation | +| Stage | Duration | What was found | +| ---------------- | -------- | --------------------------------------------------------------------------------------- | +| CPU Components | 2 weeks | 4 patches: rmsnorm precision, YaRN rotary, MoE routing + weight layout, attention sinks | +| CPU E2E | 1 week | mp.spawn patch inheritance, weight layout fix, bias restoration | +| Device E2E | 2 weeks | SPMDRank for ParallelEmbedding (590→127), windowed attention path (127→0.97) | +| **Final result** | | R = 1.0005, top-1 match, coherent generation | --- diff --git a/skills/neuron-framework-equivalence/references/debugging-case-study-gptoss.md b/skills/neuron-framework-equivalence/references/debugging-case-study-gptoss.md index e2b02f0..f21edda 100644 --- a/skills/neuron-framework-equivalence/references/debugging-case-study-gptoss.md +++ b/skills/neuron-framework-equivalence/references/debugging-case-study-gptoss.md @@ -5,37 +5,37 @@ Use this to pattern-match on failures you encounter. ## Timeline -| Stage | Duration | What Was Found | -|-------|----------|----------------| -| Component debugging (CPU) | 2 weeks | 4 patches: rmsnorm precision, YaRN rotary scaling, MoE routing + weight layout, attention sinks | -| CPU E2E debugging | 1 week | mp.spawn patch inheritance, weight layout fix, bias restoration | -| Device E2E debugging | 2 weeks | Root cause #1: ParallelEmbedding missing SPMDRank (590→127). Root cause #2: Windowed attention path missing sink injection (127→0.97) | +| Stage | Duration | What Was Found | +| ------------------------- | -------- | ------------------------------------------------------------------------------------------------------------------------------------- | +| Component debugging (CPU) | 2 weeks | 4 patches: rmsnorm precision, YaRN rotary scaling, MoE routing + weight layout, attention sinks | +| CPU E2E debugging | 1 week | mp.spawn patch inheritance, weight layout fix, bias restoration | +| Device E2E debugging | 2 weeks | Root cause #1: ParallelEmbedding missing SPMDRank (590→127). Root cause #2: Windowed attention path missing sink injection (127→0.97) | ## Final Result -| Metric | Value | -|--------|-------| -| Full model error_ratio | 1.0005 | -| Top-1 token match | Yes | -| 20-token generation | Coherent ("a topic of much debate and speculation...") | -| Total patches | 5 (3 unified CPU+device, 2 device-only) | +| Metric | Value | +| ---------------------- | ------------------------------------------------------ | +| Full model error_ratio | 1.0005 | +| Top-1 token match | Yes | +| 20-token generation | Coherent ("a topic of much debate and speculation...") | +| Total patches | 5 (3 unified CPU+device, 2 device-only) | --- ## Component-Level Results (Stage 2) -| Component | Error Ratio | Threshold | Result | Root Cause | -|-----------|-------------|-----------|--------|------------| -| RMSNorm | 1.85 | 2.0 | PASS* | Neuron computes variance in bf16 instead of fp32 | -| Embedding | 1.00 | 1.1 | PASS | Identical bf16 lookup | -| Rotary Embedding | 130–133 | 1.2 | FAIL | Missing YaRN scaling entirely | -| Linear Projections (Q/K/V/O) | 1.15–1.17 | 1.2 | PASS | Identical linear math | -| Router | 0.57 | 1.2 | PASS | Same linear + topk + softmax | -| Experts (with routing) | 1774 | 1.2 | FAIL | Neuron ignores routing entirely | -| Full MLP/MoE | 1809 | 1.2 | FAIL | Cascaded routing + experts bugs | -| LM Head | 1.00 | 1.2 | PASS | Identical linear math | +| Component | Error Ratio | Threshold | Result | Root Cause | +| ---------------------------- | ----------- | --------- | ------ | ------------------------------------------------ | +| RMSNorm | 1.85 | 2.0 | PASS\* | Neuron computes variance in bf16 instead of fp32 | +| Embedding | 1.00 | 1.1 | PASS | Identical bf16 lookup | +| Rotary Embedding | 130–133 | 1.2 | FAIL | Missing YaRN scaling entirely | +| Linear Projections (Q/K/V/O) | 1.15–1.17 | 1.2 | PASS | Identical linear math | +| Router | 0.57 | 1.2 | PASS | Same linear + topk + softmax | +| Experts (with routing) | 1774 | 1.2 | FAIL | Neuron ignores routing entirely | +| Full MLP/MoE | 1809 | 1.2 | FAIL | Cascaded routing + experts bugs | +| LM Head | 1.00 | 1.2 | PASS | Identical linear math | -*RMSNorm was hidden by a relaxed 2.0 tolerance. After patching, it dropped to 1.0. +\*RMSNorm was hidden by a relaxed 2.0 tolerance. After patching, it dropped to 1.0. --- @@ -44,14 +44,17 @@ Use this to pattern-match on failures you encounter. **Symptom:** `test_02_rotary_emb.py` — error ratio 130x (cos) and 133x (sin). **Diagnosis:** + - HF `GptOssRotaryEmbedding` uses `ROPE_INIT_FUNCTIONS["yarn"]` which blends interpolated/extrapolated `inv_freq` and applies `attention_scaling = 1.3466` to cos/sin. - Neuron `NeuronGptOssRotaryEmbedding` wraps the framework's basic `RotaryEmbedding(dim, max_pos, base)` which only computes standard RoPE. The `rope_scaling` config is stored but never used. **Patch:** `yarn_rotary_patch.py`: + 1. `_patched_init`: computes YaRN `inv_freq` and `attention_scaling`, stores as `self._yarn_inv_freq` and `self.attention_scaling` 2. `_patched_forward`: computes rotary embeddings in fp32 using YaRN inv_freq, applies attention_scaling before dtype cast **Pitfalls encountered:** + 1. Neuron config's `rope_scaling` dict was missing `original_max_position_embeddings` (4096). Derived: `original = max_position_embeddings / factor`. 2. Framework's `register_buffer("inv_freq", None)` resisted direct assignment. Stored `_yarn_inv_freq` on the wrapper instead. 3. Applying `attention_scaling` after the framework cast cos/sin to bf16 caused sin error ratio of 13.7x. Fixed by reimplementing forward inline, applying scaling in fp32 before the cast. @@ -79,11 +82,13 @@ Use this to pattern-match on failures you encounter. **Symptom:** `test_05_experts.py` — error ratio 1774x. `test_06_mlp_moe.py` — error ratio 1809x. **Diagnosis:** Three bugs: + 1. `NeuronGptOssExperts.forward()` ignores `router_indices` and `routing_weights` — all tokens go through all 32 experts instead of top-4. 2. `NeuronGptOssMLP.forward()` uses `router_scores.sum(dim=-1)` as multiplier — softmax scores sum to ~1.0, making it a no-op. 3. Down-proj weight layout mismatch (`i*E+e` vs `e*I+i` indexing) and per-expert biases collapsed to single sum. **Patch:** `mlp_moe_patch.py`: + 1. Extracts per-expert weights from flattened RowParallelLinear via reshape + permute 2. Computes per-expert down projections via `einsum('nei,eih->neh')` 3. Applies routing weights per expert, then sums across experts @@ -106,20 +111,20 @@ Use this to pattern-match on failures you encounter. ### Secondary Issues at TP>1 -| Issue | Fix | -|-------|-----| +| Issue | Fix | +| ----------------------------------------------------- | -------------------------------------------------------- | | Down-proj weight layout incompatible with TP sharding | `fix_down_proj_layout()` before `get_sharded_checkpoint` | -| Attention sink parameters not TP-sharded | Manual slicing: `sinks[rank*local:(rank+1)*local]` | -| Biases removed by `get_sharded_checkpoint` | Manual bias restoration with TP-aware slicing | +| Attention sink parameters not TP-sharded | Manual slicing: `sinks[rank*local:(rank+1)*local]` | +| Biases removed by `get_sharded_checkpoint` | Manual bias restoration with TP-aware slicing | ### CPU E2E Final Results -| Test | Metric | Value | Threshold | Result | -|------|--------|-------|-----------|--------| -| FP32 Direct TP=1 | rel_fro_norm | 3.986e-7 | 1e-5 | PASS | -| 3-Tensor BF16 TP=1 | error_ratio | 0.878 | 1.2 | PASS | -| 3-Tensor BF16 TP=4 | error_ratio | 0.850 | 1.2 | PASS | -| 3-Tensor BF16 TP=8 | error_ratio | 0.976 | 1.2 | PASS | +| Test | Metric | Value | Threshold | Result | +| ------------------ | ------------ | -------- | --------- | ------ | +| FP32 Direct TP=1 | rel_fro_norm | 3.986e-7 | 1e-5 | PASS | +| 3-Tensor BF16 TP=1 | error_ratio | 0.878 | 1.2 | PASS | +| 3-Tensor BF16 TP=4 | error_ratio | 0.850 | 1.2 | PASS | +| 3-Tensor BF16 TP=8 | error_ratio | 0.976 | 1.2 | PASS | --- @@ -140,6 +145,7 @@ GPT-OSS layer 0 is `sliding_attention`, which dispatches to `perform_prefill_win ### Key Insight The two device-specific root causes would never have been found by CPU testing alone: + - **SPMDRank:** `parallel_state.get_rank()` returns correct values on CPU but bakes as constant 0 during XLA tracing - **Windowed attention:** CPU used the patched path by coincidence (TP=1 always uses `perform_prefill`), while device layer 0 used the unpatched windowed path @@ -147,11 +153,11 @@ The two device-specific root causes would never have been found by CPU testing a ## Final Patch Inventory -| Patch | What It Fixes | CPU/Device | -|-------|--------------|------------| -| `yarn_rotary_patch.py` | YaRN scaling in rotary embeddings | Both | -| `rmsnorm_patch.py` | FP32 variance computation | Both | -| `mlp_moe_patch.py` | MoE routing + weight layout | Both | -| `attention_sink_patch.py` | Sinks on both attention paths | Both | -| `embedding_patch_device.py` | SPMDRank for ParallelEmbedding | Device | -| `attention_bias_patch_device.py` | Enable attention biases in projections | Device | +| Patch | What It Fixes | CPU/Device | +| -------------------------------- | -------------------------------------- | ---------- | +| `yarn_rotary_patch.py` | YaRN scaling in rotary embeddings | Both | +| `rmsnorm_patch.py` | FP32 variance computation | Both | +| `mlp_moe_patch.py` | MoE routing + weight layout | Both | +| `attention_sink_patch.py` | Sinks on both attention paths | Both | +| `embedding_patch_device.py` | SPMDRank for ParallelEmbedding | Device | +| `attention_bias_patch_device.py` | Enable attention biases in projections | Device | diff --git a/skills/neuron-framework-equivalence/references/device-component-debugging.md b/skills/neuron-framework-equivalence/references/device-component-debugging.md index dbf81e5..8530b4e 100644 --- a/skills/neuron-framework-equivalence/references/device-component-debugging.md +++ b/skills/neuron-framework-equivalence/references/device-component-debugging.md @@ -47,13 +47,13 @@ result = compare_3tensors(ref_fp32_out, ref_bf16_out, device_bf16_out) ## Device-Specific Root Cause Categories -| Root Cause | Error Magnitude | Symptom | Diagnostic | -|-----------|----------------|---------|------------| -| **SPMDRank** (Python ints bake as constants) | 100x-600x | Embeddings/routing use rank-0 values on all TP cores | Check if component uses `parallel_state.get_rank()` for slicing | -| **Code path divergence** | 100x+ | Device uses different dispatch path than CPU | Trace framework dispatch logic; check `sliding_window`, `is_prefill`, etc. | -| **Missing Parameters** (not in compiled NEFF) | 10x+ | Bias or weight tensor is zero/absent on device | Check for tensors assigned but not registered as `nn.Parameter` | -| **Weight preprocessing** (pre_shard_weights_hook) | 1000x+ | Weights laid out differently than expected | Compare weight checksums before/after shard hook | -| **Missing bias flags** | 10x+ | `preprocess_checkpoint` removes biases as "redundant" | Check if base class constructor receives `has_bias=True` | +| Root Cause | Error Magnitude | Symptom | Diagnostic | +| ------------------------------------------------- | --------------- | ----------------------------------------------------- | -------------------------------------------------------------------------- | +| **SPMDRank** (Python ints bake as constants) | 100x-600x | Embeddings/routing use rank-0 values on all TP cores | Check if component uses `parallel_state.get_rank()` for slicing | +| **Code path divergence** | 100x+ | Device uses different dispatch path than CPU | Trace framework dispatch logic; check `sliding_window`, `is_prefill`, etc. | +| **Missing Parameters** (not in compiled NEFF) | 10x+ | Bias or weight tensor is zero/absent on device | Check for tensors assigned but not registered as `nn.Parameter` | +| **Weight preprocessing** (pre_shard_weights_hook) | 1000x+ | Weights laid out differently than expected | Compare weight checksums before/after shard hook | +| **Missing bias flags** | 10x+ | `preprocess_checkpoint` removes biases as "redundant" | Check if base class constructor receives `has_bias=True` | --- @@ -93,7 +93,7 @@ indices = torch.arange(local_size, device=tensor.device) + (rank * local_size).t local_slice = torch.index_select(tensor, 0, indices) ``` -### Pattern 3: Framework _reduce Instead of torch.distributed +### Pattern 3: Framework \_reduce Instead of torch.distributed **Problem:** `torch.distributed.all_reduce()` is not executable during XLA tracing. @@ -197,7 +197,7 @@ Revise the patch, return to diagnosis ## Pitfalls 1. **TensorCaptureConfig requires OnDeviceSamplingConfig** — without it, captured tensors are silently not returned. -2. **TP-gathered outputs are N*vocab_size** — take first `vocab_size` entries: `logits[:, :, :vocab_size]`. +2. **TP-gathered outputs are N\*vocab_size** — take first `vocab_size` entries: `logits[:, :, :vocab_size]`. 3. **preprocess_checkpoint removes "redundant" biases** — ensure base class receives bias flags (`qkv_bias=True`, `o_bias=True`). 4. **self_attn captures cos_cache, not hidden_states** — use `post_attention_layernorm` as attention quality proxy. 5. **"Removing redundant keys" warning is normal** — framework's preshard hook remaps individual q/k/v into combined qkv. @@ -211,6 +211,7 @@ If framework code matches the reference (verified by manual reconstruction on CP **Before escalating**, analyze the compiler log `log-neuron-cc.txt` for errors, warnings, or unexpected optimization passes that may explain the divergence. Escalate when: + - `log-neuron-cc.txt` has been reviewed and does not reveal an actionable fix - All patches verified correct on CPU - Code paths confirmed identical between CPU and device modes diff --git a/skills/neuron-framework-equivalence/references/device-e2e-debugging.md b/skills/neuron-framework-equivalence/references/device-e2e-debugging.md index d73d213..0bb0e50 100644 --- a/skills/neuron-framework-equivalence/references/device-e2e-debugging.md +++ b/skills/neuron-framework-equivalence/references/device-e2e-debugging.md @@ -14,12 +14,12 @@ Measure how far device output is from the HF reference: 2. Run device E2E with the same prompt, capture logits 3. Compute error_ratio = `||device - fp32|| / ||bf16 - fp32||` -| error_ratio | Interpretation | Action | -|-------------|---------------|--------| -| <= 1.2 | Within BF16 precision | **PASS** — done | -| 1.2 - 10 | Moderate divergence | Precision or scaling bug | -| 10 - 100 | Significant divergence | Missing operation or wrong code path | -| 100+ | Catastrophic divergence | Fundamental issue (wrong weights, missing algorithm) | +| error_ratio | Interpretation | Action | +| ----------- | ----------------------- | ---------------------------------------------------- | +| <= 1.2 | Within BF16 precision | **PASS** — done | +| 1.2 - 10 | Moderate divergence | Precision or scaling bug | +| 10 - 100 | Significant divergence | Missing operation or wrong code path | +| 100+ | Catastrophic divergence | Fundamental issue (wrong weights, missing algorithm) | --- @@ -29,16 +29,16 @@ Override `config.num_hidden_layers = 1` when creating the model. Generate 1-laye ### Choosing Modules to Capture -| # | Module | What It Captures | Sharded? | -|---|--------|-----------------|----------| -| 1 | `embed_tokens` | Token embeddings | Gathered | -| 2 | `layers.0.input_layernorm` | Pre-attention norm | No | -| 3 | `layers.0.self_attn` | Attention output | No (all-reduced) | -| 4 | `layers.0.post_attention_layernorm` | Pre-MLP norm | No | -| 5 | `layers.0.mlp` | MLP output | No (all-reduced) | -| 6 | `layers.0` | Full layer output | No | -| 7 | `norm` | Final norm | No | -| 8 | `lm_head` | Logits | Gathered | +| # | Module | What It Captures | Sharded? | +| --- | ----------------------------------- | ------------------ | ---------------- | +| 1 | `embed_tokens` | Token embeddings | Gathered | +| 2 | `layers.0.input_layernorm` | Pre-attention norm | No | +| 3 | `layers.0.self_attn` | Attention output | No (all-reduced) | +| 4 | `layers.0.post_attention_layernorm` | Pre-MLP norm | No | +| 5 | `layers.0.mlp` | MLP output | No (all-reduced) | +| 6 | `layers.0` | Full layer output | No | +| 7 | `norm` | Final norm | No | +| 8 | `lm_head` | Logits | Gathered | ### Reading the Comparison Table @@ -58,12 +58,12 @@ The **first FAIL** is where the bug originates. Everything after may cascade fro ### Strategy 1: Read the First Failure -| First failing module | Likely cause | -|---------------------|-------------| -| `embed_tokens` | Embedding vocabulary partition issue (SPMDRank) | +| First failing module | Likely cause | +| -------------------------- | ------------------------------------------------------------------ | +| `embed_tokens` | Embedding vocabulary partition issue (SPMDRank) | | `post_attention_layernorm` | Attention computation bug (missing algorithm, wrong dispatch path) | -| `layers.0.mlp` / `experts` | MoE routing, weight layout, or bias issue | -| `norm` / `lm_head` | Cascaded error from earlier layers | +| `layers.0.mlp` / `experts` | MoE routing, weight layout, or bias issue | +| `norm` / `lm_head` | Cascaded error from earlier layers | ### Strategy 2: Manual Component Reconstruction @@ -125,12 +125,12 @@ Once all 1-layer modules pass: 4. Verify top-1 token match between HF and device 5. Run multi-token generation (20 tokens) — check for coherent output -| Metric | Threshold | -|--------|-----------| -| error_ratio (last position logits) | <= 1.2 | -| Top-1 token match | Yes | -| Top-5 token overlap | >= 4/5 | -| 20-token generation | Coherent, no repetition/garbage | +| Metric | Threshold | +| ---------------------------------- | ------------------------------- | +| error_ratio (last position logits) | <= 1.2 | +| Top-1 token match | Yes | +| Top-5 token overlap | >= 4/5 | +| 20-token generation | Coherent, no repetition/garbage | --- @@ -144,22 +144,23 @@ Once all 1-layer modules pass: ## Common Root Causes -| Root Cause | Error Magnitude | Fix Pattern | -|-----------|----------------|-------------| -| Missing SPMDRank | 100x-600x | Add SPMDRank, inject in pre_shard_weights_hook | -| Missing algorithm on alternate code path | 100x+ | Trace code paths, patch the missed path | -| Weight layout mismatch | 1000x+ | Fix layout in pre_shard_weights_hook before sharding | -| Missing bias Parameters | 10x+ | Create bias Parameter in __init__, pass bias flags to base class | -| Patches not applied in mp.spawn workers | matches no-patch | Re-apply all patches inside each worker function | +| Root Cause | Error Magnitude | Fix Pattern | +| ---------------------------------------- | ---------------- | ---------------------------------------------------------------- | +| Missing SPMDRank | 100x-600x | Add SPMDRank, inject in pre_shard_weights_hook | +| Missing algorithm on alternate code path | 100x+ | Trace code paths, patch the missed path | +| Weight layout mismatch | 1000x+ | Fix layout in pre_shard_weights_hook before sharding | +| Missing bias Parameters | 10x+ | Create bias Parameter in **init**, pass bias flags to base class | +| Patches not applied in mp.spawn workers | matches no-patch | Re-apply all patches inside each worker function | --- ## XLA-Compatible Patch Patterns See [device-component-debugging.md](device-component-debugging.md) for the five key patterns: + 1. SPMDRank instead of parallel_state 2. torch.index_select instead of torch.narrow -3. Framework _reduce instead of torch.distributed +3. Framework \_reduce instead of torch.distributed 4. Parameters must exist before tracing 5. Unified CPU+device patches diff --git a/skills/neuron-framework-equivalence/references/dump-tensors.md b/skills/neuron-framework-equivalence/references/dump-tensors.md index 7cb30f9..6d1abf7 100644 --- a/skills/neuron-framework-equivalence/references/dump-tensors.md +++ b/skills/neuron-framework-equivalence/references/dump-tensors.md @@ -28,6 +28,7 @@ MODULES_TO_CAPTURE = [ ``` captured_tensors_{phase}_step_{step}_module_{module_name}_output.pt ``` + - `phase`: `"cte"` for context encoding (prefill), `"tg"` for token generation - `step`: generation step number (1 for CTE) - Tuple outputs: element 0 saved as `_output_0.pt` @@ -43,12 +44,12 @@ device_aligned = device_tensor[:, :ref_seq_len, :] ### Known Quirks -| Module | Quirk | Workaround | -|--------|-------|------------| -| `self_attn` | Device captures `cos_cache` (3rd field of NeuronAttentionBaseOutput) instead of hidden_states | Use `post_attention_layernorm` as proxy for attention quality | -| `embed_tokens` | FP32 and BF16 embeddings identical (lookup, no computation) → baseline_err = 0 → error_ratio = inf | Check cosine similarity instead; cosine = 1.0 means PASS | -| `lm_head` | Device with on-device sampling outputs `[1, 1, vocab]` (last position only), HF outputs `[1, seq_len, vocab]` | Compare only last position: `hf[:, -1:, :]` vs `device[:, :, :]` | -| Sharded modules | Device captures local TP shard, not global tensor | Mark as sharded; interpret with care or skip in layer comparison | +| Module | Quirk | Workaround | +| --------------- | ------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------- | +| `self_attn` | Device captures `cos_cache` (3rd field of NeuronAttentionBaseOutput) instead of hidden_states | Use `post_attention_layernorm` as proxy for attention quality | +| `embed_tokens` | FP32 and BF16 embeddings identical (lookup, no computation) → baseline_err = 0 → error_ratio = inf | Check cosine similarity instead; cosine = 1.0 means PASS | +| `lm_head` | Device with on-device sampling outputs `[1, 1, vocab]` (last position only), HF outputs `[1, seq_len, vocab]` | Compare only last position: `hf[:, -1:, :]` vs `device[:, :, :]` | +| Sharded modules | Device captures local TP shard, not global tensor | Mark as sharded; interpret with care or skip in layer comparison | --- diff --git a/skills/neuron-framework-equivalence/references/enable-model-run.md b/skills/neuron-framework-equivalence/references/enable-model-run.md index 4c9f5d2..22fa31f 100644 --- a/skills/neuron-framework-equivalence/references/enable-model-run.md +++ b/skills/neuron-framework-equivalence/references/enable-model-run.md @@ -4,10 +4,10 @@ Workflow for compiling a target model and enabling it to run on both Neuron devi ## Templates -| Template | Purpose | -|----------|---------| +| Template | Purpose | +| ----------------------------------------------- | ---------------------------------------------------------- | | `templates/neuron_device_validator_template.py` | Validates model runs on Neuron hardware, checks throughput | -| `templates/run_inference_cpu_template.py` | CPU inference with monkey patches for E2E comparison | +| `templates/run_inference_cpu_template.py` | CPU inference with monkey patches for E2E comparison | --- @@ -48,6 +48,7 @@ config = ModelInferenceConfig(neuron_config, load_config=load_pretrained_config( ``` **Expected output:** + ``` {COMPILED_MODEL_PATH}/ ├── model.pt # NEFF binary @@ -58,6 +59,7 @@ config = ModelInferenceConfig(neuron_config, load_config=load_pretrained_config( ### 2. Validate Neuron Execution Run `templates/neuron_device_validator_template.py` and verify: + - `neuron-ls` shows available devices - Logs show `CPU Mode: False` - Throughput meets threshold (small models: >20 tok/s, medium: >10 tok/s) @@ -65,6 +67,7 @@ Run `templates/neuron_device_validator_template.py` and verify: ### 3. Enable CPU Execution Run `templates/run_inference_cpu_template.py` with `cpu_mode=True`: + - Bypasses NEFF, loads from HF weights - Use same dtype (bfloat16) for fair comparison - CPU is typically 10–20x slower — this validates the difference @@ -73,19 +76,20 @@ Run `templates/run_inference_cpu_template.py` with `cpu_mode=True`: ## Compilation Troubleshooting -| Issue | Solution | -|-------|----------| -| HLO verification fails | Set `NEURON_CC_FLAGS='--internal-hlo2tensorizer-options=--verify-hlo=false'` | -| `get_program_sharding_info` missing | Add fallback function in affected files | -| Compilation OOM | Reduce `seq_len` or `batch_size` | -| NEFF not found at inference | Check `output_path` matches `compiled_model_path` | -| Low throughput (< 5 tok/s) | Verify not running on CPU fallback — check `neuron-ls` and logs | +| Issue | Solution | +| ----------------------------------- | ---------------------------------------------------------------------------- | +| HLO verification fails | Set `NEURON_CC_FLAGS='--internal-hlo2tensorizer-options=--verify-hlo=false'` | +| `get_program_sharding_info` missing | Add fallback function in affected files | +| Compilation OOM | Reduce `seq_len` or `batch_size` | +| NEFF not found at inference | Check `output_path` matches `compiled_model_path` | +| Low throughput (< 5 tok/s) | Verify not running on CPU fallback — check `neuron-ls` and logs | --- ## Logging Requirements Always capture full output: + ```bash # Compilation — the snippet from "1. Compile the Model" above python3 .py 2>&1 | tee logs/compilation.log @@ -98,6 +102,7 @@ python3 .py 2>&1 | tee logs/inference_neuron.log ## Known Issue: BF16 with Gloo Backend (CPU TP>1) **Problem:** CPU inference with `tp_degree > 1` fails with: + ``` "The gloo backend does not natively support bfloat16" ``` @@ -107,6 +112,7 @@ python3 .py 2>&1 | tee logs/inference_neuron.log **Fix:** Two changes required: 1. **`comm.py`** (NeuronxDistributed) — upcast BF16→FP32 before reduction, cast back after: + ```python def all_reduce(...): if cpu_mode(): diff --git a/skills/neuron-framework-equivalence/references/equiv-concept.md b/skills/neuron-framework-equivalence/references/equiv-concept.md index 615166c..f396743 100644 --- a/skills/neuron-framework-equivalence/references/equiv-concept.md +++ b/skills/neuron-framework-equivalence/references/equiv-concept.md @@ -28,11 +28,11 @@ Two neural network implementations are **numerically equivalent** when they prod Equivalence is evaluated under matching precision configurations. Common data types: -| Data Type | Description | -|-----------|-------------| -| fp32 | 32-bit IEEE float — used as high-precision ground truth | -| bf16 | bfloat16 — 16-bit with 8-bit exponent, 7-bit mantissa | -| mxfp8 | Microscaling FP8 — 8-bit with block scaling factors | +| Data Type | Description | +| --------- | ------------------------------------------------------- | +| fp32 | 32-bit IEEE float — used as high-precision ground truth | +| bf16 | bfloat16 — 16-bit with 8-bit exponent, 7-bit mantissa | +| mxfp8 | Microscaling FP8 — 8-bit with block scaling factors | The reference and target are compared under the **same** target precision (e.g., both bf16). The fp32 reference serves as a separate high-precision anchor. @@ -42,11 +42,11 @@ The reference and target are compared under the **same** target precision (e.g., The comparison method requires three outputs. These can come from live execution or from pre-computed tensor files stored on disk. -| Run | Implementation | Precision | Purpose | -|-----|---------------|-----------|---------| -| **out_1** | Reference | fp32 | High-precision ground truth (usually on CPU or GPU) | -| **out_2** | Reference | Target precision (bf16, mxfp8, etc.) | Quantization baseline (usually on CPU or GPU) | -| **out_3** | Target | Target precision (bf16, mxfp8, etc.) | The implementation under test, on a user-specified platform (CPU, Trainium 1/2/3, etc.) | +| Run | Implementation | Precision | Purpose | +| --------- | -------------- | ------------------------------------ | --------------------------------------------------------------------------------------- | +| **out_1** | Reference | fp32 | High-precision ground truth (usually on CPU or GPU) | +| **out_2** | Reference | Target precision (bf16, mxfp8, etc.) | Quantization baseline (usually on CPU or GPU) | +| **out_3** | Target | Target precision (bf16, mxfp8, etc.) | The implementation under test, on a user-specified platform (CPU, Trainium 1/2/3, etc.) | The method determines whether **out_3** is numerically equivalent to **out_2**, using **out_1** as the high-precision anchor. The specific hardware platform for run 3 is a user-specified parameter — the method itself is hardware-agnostic. @@ -100,11 +100,11 @@ When fp32 execution is unavailable, we cannot directly measure the inherent prec Run the reference implementation twice with the same precision: -| Run | Input | Output | -|-----|-------|--------| -| **out_1** | Original input X | Reference @ target precision | +| Run | Input | Output | +| --------- | ----------------------- | ---------------------------- | +| **out_1** | Original input X | Reference @ target precision | | **out_2** | Perturbed input X + eps | Reference @ target precision | -| **out_3** | Original input X | Target @ target precision | +| **out_3** | Original input X | Target @ target precision | ### Perturbation Details @@ -133,23 +133,24 @@ Beyond the scalar error ratio, we examine the **distribution** of elementwise er For any 3-way comparison (Method 1 or Method 2): -- **err_{2,1}** = out_2 - out_1 (elementwise difference tensor between the baseline and ground truth) -- **err_{3,1}** = out_3 - out_1 (elementwise difference tensor between the target and ground truth) +- **err\_{2,1}** = out_2 - out_1 (elementwise difference tensor between the baseline and ground truth) +- **err\_{3,1}** = out_3 - out_1 (elementwise difference tensor between the target and ground truth) -Each element in these tensors is treated as a **sample of numerical error**. If out_2 and out_3 are meaningfully close to each other, err_{2,1} and err_{3,1} should follow the **same distribution**. +Each element in these tensors is treated as a **sample of numerical error**. If out*2 and out_3 are meaningfully close to each other, err*{2,1} and err\_{3,1} should follow the **same distribution**. ### Visual Tools -| Tool | What to look for | -|------|-----------------| -| **QQ Plot** | Plot quantiles of err_{3,1} against quantiles of err_{2,1}. Points should fall on the **45-degree line** if the distributions match. | -| **Histogram** | Overlay both error distributions. The shapes should **overlap closely**. | +| Tool | What to look for | +| ------------- | ------------------------------------------------------------------------------------------------------------------------------------ | +| **QQ Plot** | Plot quantiles of err*{3,1} against quantiles of err*{2,1}. Points should fall on the **45-degree line** if the distributions match. | +| **Histogram** | Overlay both error distributions. The shapes should **overlap closely**. | The QQ plot is particularly informative: systematic deviations from the diagonal indicate a distributional shift (e.g., the target has heavier tails, a different mean, or a different variance), which points to a specific type of implementation error. ### Implementation Reference `tensor_compare.py` in `.claude/skills/equiv-concept/scripts/` provides: + - `compare_3tensors(out_1, out_2, out_3)` — returns a 12-key dictionary of normwise and elementwise metrics with `_2_1` and `_3_1` suffixes - `compare_2tensors(tensor1, tensor2)` — returns a 6-key dictionary for pairwise comparison - `_visualize_differences_two_series(elem_diff1, elem_diff2, ...)` — generates overlaid histograms and QQ plots @@ -160,17 +161,17 @@ The QQ plot is particularly informative: systematic deviations from the diagonal ### Error Ratio Thresholds -| Error Ratio | Interpretation | -|-------------|---------------| -| ~ 1.0 | Excellent — target matches reference precision error | -| <= 1.1 - 1.2 | Good — within acceptable tolerance | -| 1.2 - 2.0 | Marginal — may be acceptable with documented justification (e.g., known precision ordering differences at higher tensor parallelism) | -| >> 1.2 | Fail — implementation bug likely | +| Error Ratio | Interpretation | +| ------------ | ------------------------------------------------------------------------------------------------------------------------------------ | +| ~ 1.0 | Excellent — target matches reference precision error | +| <= 1.1 - 1.2 | Good — within acceptable tolerance | +| 1.2 - 2.0 | Marginal — may be acceptable with documented justification (e.g., known precision ordering differences at higher tensor parallelism) | +| >> 1.2 | Fail — implementation bug likely | ### Distribution Criteria - QQ plot points should lie on the 45-degree line -- Histograms of err_{2,1} and err_{3,1} should overlap +- Histograms of err*{2,1} and err*{3,1} should overlap A passing error ratio with a failing QQ plot (or vice versa) warrants further investigation — the normwise metric can mask localized outliers that the distributional analysis reveals. @@ -180,13 +181,13 @@ A passing error ratio with a failing QQ plot (or vice versa) warrants further in This concepts skill provides the theoretical foundation. The following execution skills implement the workflow: -| Step | Skill | Purpose | -|------|-------|---------| -| Environment setup | `env-setup` | Docker container with proper mounts and dependencies | -| Model structure analysis | `build-model-tree`, `component-mapping` | Understand and map module hierarchies between reference and target | -| Component-level testing | `component-testing` | Build bottom-up equivalence tests using the 3-way comparison | -| CPU component debugging | `cpu-component-debugging` | Diagnose and fix failing component tests on CPU via monkey patches | -| Device component debugging | `device-component-debugging` | Diagnose and fix failing component tests on device using XLA-compatible patches | -| Intermediate tensor capture | `tensor-capture` | Capture tensors at specific layers for targeted comparison | -| Device execution | `enable-model-run` | Compile and run the target on Neuron devices | -| Compiler issues | (analyze `log-neuron-cc.txt`, then escalate externally) | Analyze `log-neuron-cc.txt` for errors/warnings first, then file a [Neuron SDK GitHub issue](https://github.com/aws-neuron/aws-neuron-sdk/issues) with reproduction steps | +| Step | Skill | Purpose | +| --------------------------- | ------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| Environment setup | `env-setup` | Docker container with proper mounts and dependencies | +| Model structure analysis | `build-model-tree`, `component-mapping` | Understand and map module hierarchies between reference and target | +| Component-level testing | `component-testing` | Build bottom-up equivalence tests using the 3-way comparison | +| CPU component debugging | `cpu-component-debugging` | Diagnose and fix failing component tests on CPU via monkey patches | +| Device component debugging | `device-component-debugging` | Diagnose and fix failing component tests on device using XLA-compatible patches | +| Intermediate tensor capture | `tensor-capture` | Capture tensors at specific layers for targeted comparison | +| Device execution | `enable-model-run` | Compile and run the target on Neuron devices | +| Compiler issues | (analyze `log-neuron-cc.txt`, then escalate externally) | Analyze `log-neuron-cc.txt` for errors/warnings first, then file a [Neuron SDK GitHub issue](https://github.com/aws-neuron/aws-neuron-sdk/issues) with reproduction steps | diff --git a/skills/neuron-framework-equivalence/references/expected_structural_diffs.md b/skills/neuron-framework-equivalence/references/expected_structural_diffs.md index ef804e0..441772f 100644 --- a/skills/neuron-framework-equivalence/references/expected_structural_diffs.md +++ b/skills/neuron-framework-equivalence/references/expected_structural_diffs.md @@ -6,50 +6,52 @@ pairs, adapt accordingly. ## Module Type Changes (TP Sharding) -| HuggingFace | Neuron Port | Reason | -|-------------|-------------|--------| -| `nn.Linear` | `ColumnParallelLinear` | Tensor parallel sharding (column split) | -| `nn.Linear` | `RowParallelLinear` | Tensor parallel sharding (row split) | -| `nn.Embedding` | `ParallelEmbedding` | Embedding sharded across TP ranks | +| HuggingFace | Neuron Port | Reason | +| -------------- | ---------------------- | --------------------------------------- | +| `nn.Linear` | `ColumnParallelLinear` | Tensor parallel sharding (column split) | +| `nn.Linear` | `RowParallelLinear` | Tensor parallel sharding (row split) | +| `nn.Embedding` | `ParallelEmbedding` | Embedding sharded across TP ranks | ## Structural Wrappers (Framework) -| HuggingFace | Neuron Port | Reason | -|-------------|-------------|--------| -| Flat `q_proj, k_proj, v_proj` | Wrapped in `GroupQueryAttention_QKV` | NxDI attention framework | -| Flat `o_proj` | Wrapped in `GroupQueryAttention_O` | NxDI attention framework | -| (none) | `SPMDRank` (rank_util) | TP rank tracking | -| (none) | `KVCacheManager` (kv_mgr) | Inference KV cache management | +| HuggingFace | Neuron Port | Reason | +| ----------------------------- | ------------------------------------ | ----------------------------- | +| Flat `q_proj, k_proj, v_proj` | Wrapped in `GroupQueryAttention_QKV` | NxDI attention framework | +| Flat `o_proj` | Wrapped in `GroupQueryAttention_O` | NxDI attention framework | +| (none) | `SPMDRank` (rank_util) | TP rank tracking | +| (none) | `KVCacheManager` (kv_mgr) | Inference KV cache management | ## Normalization Variants -| HuggingFace | Neuron Port (CPU mode) | Neuron Port (Neuron HW) | -|-------------|----------------------|------------------------| -| `XxxRMSNorm` | `LlamaRMSNorm` or same class | `CustomRMSNorm` (NKI kernel) | -| `nn.LayerNorm` | `nn.LayerNorm` | `nn.LayerNorm` | +| HuggingFace | Neuron Port (CPU mode) | Neuron Port (Neuron HW) | +| -------------- | ---------------------------- | ---------------------------- | +| `XxxRMSNorm` | `LlamaRMSNorm` or same class | `CustomRMSNorm` (NKI kernel) | +| `nn.LayerNorm` | `nn.LayerNorm` | `nn.LayerNorm` | ## Operator Fusion/Split -| HuggingFace | Neuron Port | Reason | -|-------------|-------------|--------| -| Fused `gate_up_proj` [2*inter, hidden] | Split `gate_proj` + `up_proj` | TP requires separate sharding | -| Fused `qkv_proj` [3*hidden, hidden] | Split `q_proj` + `k_proj` + `v_proj` | TP requires separate sharding | -| Single `RotaryEmbedding` at model level | Per-layer `RotaryEmbedding` | Implementation choice | +| HuggingFace | Neuron Port | Reason | +| --------------------------------------- | ------------------------------------ | ----------------------------- | +| Fused `gate_up_proj` [2*inter, hidden] | Split `gate_proj` + `up_proj` | TP requires separate sharding | +| Fused `qkv_proj` [3*hidden, hidden] | Split `q_proj` + `k_proj` + `v_proj` | TP requires separate sharding | +| Single `RotaryEmbedding` at model level | Per-layer `RotaryEmbedding` | Implementation choice | ## Activation Functions -| HuggingFace | Neuron Port | Notes | -|-------------|-------------|-------| +| HuggingFace | Neuron Port | Notes | +| ------------------- | ------------------------------- | -------------------------- | | `NewGELUActivation` | `F.gelu(x, approximate='tanh')` | Same math, different class | -| `SiLU` | `SiLU` | Identical | +| `SiLU` | `SiLU` | Identical | ## Modules with No Counterpart **HF-only (no Neuron equivalent):** + - `Dropout` layers — disabled during inference (set to 0.0) - `RotaryEmbedding` at model level (Neuron uses per-layer) **Neuron-only (no HF equivalent):** + - `SPMDRank` — distributed rank utilities - `KVCacheManager` — inference KV caching - `LogitsProcessor` — sampling-time logit manipulation (some models) diff --git a/skills/neuron-framework-equivalence/references/report-template.md b/skills/neuron-framework-equivalence/references/report-template.md index c43d756..2676c25 100644 --- a/skills/neuron-framework-equivalence/references/report-template.md +++ b/skills/neuron-framework-equivalence/references/report-template.md @@ -38,11 +38,11 @@ If a phase genuinely cannot complete (e.g., compilation fails), report the block ## Stage 2: Component-Level Results -| Component | R-ratio | Threshold | Result | -|-----------|---------|-----------|--------| -| rmsnorm | {r} | 1.2 | {PASS/FAIL} | -| embedding | {r} | 1.2 | {PASS/FAIL} | -| ... | ... | ... | ... | +| Component | R-ratio | Threshold | Result | +| --------- | ------- | --------- | ----------- | +| rmsnorm | {r} | 1.2 | {PASS/FAIL} | +| embedding | {r} | 1.2 | {PASS/FAIL} | +| ... | ... | ... | ... | - Components tested: {N} - Passed: {P}/{N} @@ -51,6 +51,7 @@ If a phase genuinely cannot complete (e.g., compilation fails), report the block ## Stage 3: Fault Localization {If Stage 2 had failures:} + - Primary fault: {component} (R={r}, pattern={spike/step}) - Root cause classification: {cause} @@ -63,12 +64,12 @@ If a phase genuinely cannot complete (e.g., compilation fails), report the block ## Stages 5+6: E2E Teacher-Forced Comparison -| Metric | Value | Threshold | Result | -|--------|-------|-----------|--------| -| R-ratio (p95) | {r} | 1.2 | {PASS/FAIL} | -| Cosine sim (p5) | {cos} | 0.95 | {PASS/FAIL} | -| KL divergence (p95) | {kl} | 0.1 | {PASS/FAIL} | -| Top-1 agreement | {pct}% | 50% | {PASS/FAIL} | +| Metric | Value | Threshold | Result | +| ------------------- | ------ | --------- | ----------- | +| R-ratio (p95) | {r} | 1.2 | {PASS/FAIL} | +| Cosine sim (p5) | {cos} | 0.95 | {PASS/FAIL} | +| KL divergence (p95) | {kl} | 0.1 | {PASS/FAIL} | +| Top-1 agreement | {pct}% | 50% | {PASS/FAIL} | ## Stage 7: Downstream Evaluation diff --git a/skills/neuron-framework-equivalence/references/vllm-neuron-adaptation.md b/skills/neuron-framework-equivalence/references/vllm-neuron-adaptation.md index b52ba62..0f61d4f 100644 --- a/skills/neuron-framework-equivalence/references/vllm-neuron-adaptation.md +++ b/skills/neuron-framework-equivalence/references/vllm-neuron-adaptation.md @@ -6,19 +6,19 @@ Critical details for writing equivalence tests against vLLM-Neuron models. Based All linear weights in vLLM-Neuron are **transposed** relative to HuggingFace. This is the single most important difference. -| Weight | HF Shape | vLLM-Neuron Shape | Transform | -|--------|----------|-------------------|-----------| -| gate_proj | `[I, H]` | `[H, I]` | `.t()` | -| up_proj | `[I, H]` | `[H, I]` | `.t()` | -| down_proj | `[H, I]` | `[I, H]` | `.t()` | -| q_proj | `[q, H]` | fused into QKV | see below | -| k_proj | `[kv, H]` | fused into QKV | see below | -| v_proj | `[kv, H]` | fused into QKV | see below | -| o_proj | `[H, q]` | `[q, H]` | `.t()` | -| QKV (fused) | N/A | `[H, q+2kv]` | `cat([Q.t(), K.t(), V.t()], dim=-1)` | -| Norms | `[H]` | `[H]` | direct copy | -| Embedding | `[V, H]` | `[V, H]` | direct copy | -| LM head | `[V, H]` | `[V, H]` | direct copy | +| Weight | HF Shape | vLLM-Neuron Shape | Transform | +| ----------- | --------- | ----------------- | ------------------------------------ | +| gate_proj | `[I, H]` | `[H, I]` | `.t()` | +| up_proj | `[I, H]` | `[H, I]` | `.t()` | +| down_proj | `[H, I]` | `[I, H]` | `.t()` | +| q_proj | `[q, H]` | fused into QKV | see below | +| k_proj | `[kv, H]` | fused into QKV | see below | +| v_proj | `[kv, H]` | fused into QKV | see below | +| o_proj | `[H, q]` | `[q, H]` | `.t()` | +| QKV (fused) | N/A | `[H, q+2kv]` | `cat([Q.t(), K.t(), V.t()], dim=-1)` | +| Norms | `[H]` | `[H]` | direct copy | +| Embedding | `[V, H]` | `[V, H]` | direct copy | +| LM head | `[V, H]` | `[V, H]` | direct copy | ## Weight Naming @@ -53,20 +53,20 @@ vLLM-Neuron's `load_weights()` calls `get_current_vllm_config()` which only work ## Environment Variables -| Context | Required | -|---------|----------| -| CPU testing | `NXD_CPU_MODE=1`, `WORLD_SIZE=1`, `MASTER_ADDR=localhost`, `MASTER_PORT=8099`, `RANK=0` | -| Device testing | `NEURON_SKIP_EFA_AFFINITY=1`, `TOKENIZERS_PARALLELISM=false` | -| EP models | Add `NXDI_SWITCH_CC=1` | +| Context | Required | +| -------------- | --------------------------------------------------------------------------------------- | +| CPU testing | `NXD_CPU_MODE=1`, `WORLD_SIZE=1`, `MASTER_ADDR=localhost`, `MASTER_PORT=8099`, `RANK=0` | +| Device testing | `NEURON_SKIP_EFA_AFFINITY=1`, `TOKENIZERS_PARALLELISM=false` | +| EP models | Add `NXDI_SWITCH_CC=1` | ## Component Test Differences -| Component | HF Forward | vLLM-Neuron Forward | Weight Setup | -|-----------|-----------|---------------------|--------------| -| MLP | `forward(x)` | `forward(x, is_prefill=True)` | `.t()` on gate/up/down | -| Q/K/V | 3 separate `F.linear()` | Single `NF.qkv_proj()` with fused weight | `cat([Q.t(), K.t(), V.t()], dim=-1)` | -| O Projection | `F.linear(x, o_proj.weight)` | `NF.o_proj(x, o_proj_weight)` | `.t()` | -| Full Attention | Independently testable | NOT independently testable (needs KV cache + attn_metadata) | Test QKV + O projections separately | +| Component | HF Forward | vLLM-Neuron Forward | Weight Setup | +| -------------- | ---------------------------- | ----------------------------------------------------------- | ------------------------------------ | +| MLP | `forward(x)` | `forward(x, is_prefill=True)` | `.t()` on gate/up/down | +| Q/K/V | 3 separate `F.linear()` | Single `NF.qkv_proj()` with fused weight | `cat([Q.t(), K.t(), V.t()], dim=-1)` | +| O Projection | `F.linear(x, o_proj.weight)` | `NF.o_proj(x, o_proj_weight)` | `.t()` | +| Full Attention | Independently testable | NOT independently testable (needs KV cache + attn_metadata) | Test QKV + O projections separately | ## Shape Alignment diff --git a/skills/neuron-nki-debugging/SKILL.md b/skills/neuron-nki-debugging/SKILL.md index 6936eb6..af9cd66 100644 --- a/skills/neuron-nki-debugging/SKILL.md +++ b/skills/neuron-nki-debugging/SKILL.md @@ -61,6 +61,7 @@ Before running kernels on device, resolve the NKI virtual environment path: 3. If still not found, report: "NKI_VENV_PATH not configured. Set the environment variable or create .claude/nki-dev-suite.local.md with nki_venv_path in frontmatter." Activate before running any device tests: + ```bash source $NKI_VENV_PATH/bin/activate ``` @@ -74,13 +75,13 @@ Before compilation, detect the current hardware platform: ### Platform Target Mapping -| Hardware | Instance | Target Flag | Generation | -|----------|----------|-------------|------------| -| Trainium 1 | trn1 | `--target trn1` | gen2 | -| Trainium 1n | trn1n | `--target trn1n` | gen2 | -| Inferentia 2 | inf2 | `--target inf2` | gen2 | -| Trainium 2 | trn2 | `--target trn2` | gen3 | -| Trainium 3 | trn3 | `--target trn3` | gen4 | +| Hardware | Instance | Target Flag | Generation | +| ------------ | -------- | ---------------- | ---------- | +| Trainium 1 | trn1 | `--target trn1` | gen2 | +| Trainium 1n | trn1n | `--target trn1n` | gen2 | +| Inferentia 2 | inf2 | `--target inf2` | gen2 | +| Trainium 2 | trn2 | `--target trn2` | gen3 | +| Trainium 3 | trn3 | `--target trn3` | gen4 | Match the `--target` flag and `platform_target` decorator argument to your detected hardware. @@ -99,10 +100,10 @@ os.environ["NEURON_PLATFORM_TARGET_OVERRIDE"] = "trn2" os.environ["NEURON_RT_VISIBLE_CORES"] = "0" ``` -| Flag | Purpose | -|------|---------| -| `--target` | Hardware platform (trn1, trn2, trn3, inf2) | -| `--lnc 1` | Single NeuronCore (simplifies debugging) | +| Flag | Purpose | +| ------------------------- | ----------------------------------------------------------------------------------- | +| `--target` | Hardware platform (trn1, trn2, trn3, inf2) | +| `--lnc 1` | Single NeuronCore (simplifies debugging) | | `NEURON_RT_VISIBLE_CORES` | Pin to specific core(s) — prevents contention when multiple agents run concurrently | See `references/compiler-flags.md` for complete flag reference. @@ -117,7 +118,6 @@ def my_kernel(input_tensor): The `platform_target` environment variable MUST match the `--target` in NEURON_CC_FLAGS. - ### Step 3: Create Test Script ```python @@ -182,11 +182,12 @@ reference_output = reference_implementation(input_data) # Compiles to separate ## Compiler Artifacts Mode -For advanced debugging that preserves compiler outputs for inspection, use when you need to understand detailed compilation behavior. +For advanced debugging that preserves compiler outputs for inspection, use when you need to understand detailed compilation behavior. **When to use:** "compiler artifacts", "compiler flags", "inspect compiler log" See `references/compiler-artifacts.md` for: + - Compiler debug flag configuration (`--verbose`, `--target`, `--lnc`) - Finding the compiler temp folder - Understanding generated artifacts (`*.neff`, `log-neuron-cc.txt`) @@ -195,26 +196,26 @@ See `references/compiler-artifacts.md` for: ### Error Categories -| Error Pattern | Category | Reference | -|--------------|----------|-----------| -| `NCC_EVRF*` | Verification error | See `references/ncc-verification-errors.md` | -| `NCC_EOOM*` | Out of memory | See `references/ncc-memory-resource-errors.md` | -| `NCC_E*` (other) | Type/operation error | See `references/ncc-type-operation-errors.md` | +| Error Pattern | Category | Reference | +| ---------------- | -------------------- | ---------------------------------------------- | +| `NCC_EVRF*` | Verification error | See `references/ncc-verification-errors.md` | +| `NCC_EOOM*` | Out of memory | See `references/ncc-memory-resource-errors.md` | +| `NCC_E*` (other) | Type/operation error | See `references/ncc-type-operation-errors.md` | ### Quick Reference -See `references/compiler-error-codes.md` for the complete index of all 28 NCC_* error codes. +See `references/compiler-error-codes.md` for the complete index of all 28 NCC\_\* error codes. ### Common Error Quick Fixes -| Error Code | Category | Quick Fix | -|------------|----------|-----------| +| Error Code | Category | Quick Fix | +| ------------- | -------------------- | --------------------------------------------------------- | | `NCC_EVRF001` | Unsupported operator | Use alternative operator from `neuronx-cc list-operators` | -| `NCC_EOOM001` | Memory exceeded | Reduce batch size, use tensor/pipeline parallelism | -| `NCC_EVRF007` | Instruction limit | Apply model parallelism | -| `NCC_EVRF005` | Unsupported FP8 type | Convert to float16/bfloat16 or use gen3+ hardware | -| `NCC_EARG001` | LNC configuration | Use supported LNC count for target hardware | -| `NCC_EVRF024` | Output tensor > 4GB | Reduce tensor size or use tensor parallelism | +| `NCC_EOOM001` | Memory exceeded | Reduce batch size, use tensor/pipeline parallelism | +| `NCC_EVRF007` | Instruction limit | Apply model parallelism | +| `NCC_EVRF005` | Unsupported FP8 type | Convert to float16/bfloat16 or use gen3+ hardware | +| `NCC_EARG001` | LNC configuration | Use supported LNC count for target hardware | +| `NCC_EVRF024` | Output tensor > 4GB | Reduce tensor size or use tensor parallelism | ## Profiling (Optional) @@ -290,14 +291,13 @@ print("Validation passed!") **Required settings:** -| Setting | Source | Description | -|---------|--------|-------------| +| Setting | Source | Description | +| --------------- | --------------------------------------------------- | --------------------------------- | | `nki_venv_path` | `.claude/nki-dev-suite.local.md` or `NKI_VENV_PATH` | Python venv with neuronx packages | **Related skills:** -| Skill | Use When | -|-------|----------| -| `/neuron-nki-profiling` | Profile kernel performance | -| `/neuron-nki-docs` | Look up API documentation and error codes | - +| Skill | Use When | +| ----------------------- | ----------------------------------------- | +| `/neuron-nki-profiling` | Profile kernel performance | +| `/neuron-nki-docs` | Look up API documentation and error codes | diff --git a/skills/neuron-nki-debugging/references/compiler-artifacts.md b/skills/neuron-nki-debugging/references/compiler-artifacts.md index 7a71552..5efcf10 100644 --- a/skills/neuron-nki-debugging/references/compiler-artifacts.md +++ b/skills/neuron-nki-debugging/references/compiler-artifacts.md @@ -21,11 +21,11 @@ os.environ["NEURON_CC_FLAGS"] = ( ) ``` -| Flag | Purpose | -|------|---------| -| `--target ` | Target platform (`trn1`, `trn2`, `inf2`) | -| `--lnc ` | Logical NeuronCore config (1 or 2, default 2 on trn2) | -| `--verbose ` | Output verbosity: `info`, `warning`, `error`, `critical`, `debug` | +| Flag | Purpose | +| --------------------- | ----------------------------------------------------------------- | +| `--target ` | Target platform (`trn1`, `trn2`, `inf2`) | +| `--lnc ` | Logical NeuronCore config (1 or 2, default 2 on trn2) | +| `--verbose ` | Output verbosity: `info`, `warning`, `error`, `critical`, `debug` | ## Finding the Compiler Temp Folder @@ -43,18 +43,19 @@ ls -lt /tmp/$USER/neuroncc_compile_workdir/ | head -5 ## Generated Artifacts -| File | Description | -|------|-------------| -| `*.neff` | Compiled Neuron Executable File Format (binary) | -| `log-neuron-cc.txt` | Detailed compiler log | +| File | Description | +| ------------------- | ----------------------------------------------- | +| `*.neff` | Compiled Neuron Executable File Format (binary) | +| `log-neuron-cc.txt` | Detailed compiler log | -### *.neff +### \*.neff The compiled binary executed on Neuron hardware. Use `neuron-explorer` tools to analyze performance. ### log-neuron-cc.txt Complete compiler log including: + - Compilation phases and timing - Warnings and diagnostics - Memory allocation decisions diff --git a/skills/neuron-nki-debugging/references/compiler-error-codes.md b/skills/neuron-nki-debugging/references/compiler-error-codes.md index ce9063e..19fc2df 100644 --- a/skills/neuron-nki-debugging/references/compiler-error-codes.md +++ b/skills/neuron-nki-debugging/references/compiler-error-codes.md @@ -1,12 +1,13 @@ -# Neuron Compiler Error Codes (NCC_*) +# Neuron Compiler Error Codes (NCC\_\*) Quick reference index to Neuron Compiler (neuronx-cc) error codes. For detailed fixes and code examples, see the linked reference files. ## Overview -Neuron Compiler error codes (NCC_* prefix) indicate issues during NEFF generation from NKI kernels. These errors occur during the compilation phase after NKI kernel code has been successfully parsed but before executable NEFF files can be generated. +Neuron Compiler error codes (NCC\_\* prefix) indicate issues during NEFF generation from NKI kernels. These errors occur during the compilation phase after NKI kernel code has been successfully parsed but before executable NEFF files can be generated. **Error Code Format**: `NCC_` + - Category: 4-letter code indicating error type - Number: 3-digit identifier within category @@ -20,25 +21,25 @@ Unsupported operations, data types, or configurations detected during compilatio **Detailed reference**: `ncc-verification-errors.md` -| Error Code | Error Message | Quick Fix | -|------------|---------------|-----------| -| NCC_EVRF001 | An unsupported operator was used | Use alternative operator from `neuronx-cc list-operators` | -| NCC_EVRF004 | Complex data types are not supported | Use real-valued tensors or emulate complex arithmetic | -| NCC_EVRF005 | Unsupported F8E4M3FNUZ/F8E5M2FNUZ data type | Convert to float16/bfloat16 | -| NCC_EVRF006 | Unsupported RNG algorithm | Use default RNG via standard APIs | -| NCC_EVRF007 | Instruction count exceeds limit | Apply model parallelism | -| NCC_EVRF009 | Activation memory exceeds HBM limit | Reduce batch size or use parallelism | -| NCC_EVRF010 | Simultaneous input and kernel dilation | Use only input OR kernel dilation | -| NCC_EVRF011 | Strided convolution with dilated input | Remove stride or input dilation | -| NCC_EVRF013 | TopK does not support int32/int64 | Cast to float before TopK | -| NCC_EVRF015 | Unrecognized custom call target | Use supported custom call target | -| NCC_EVRF016 | Scatter-reduce with integer/boolean types | Cast to float types | -| NCC_EVRF017 | Reduce-window with base dilation > 1 | Set base_dilation to (1,1,1,1) | -| NCC_EVRF018 | Reduce-window with window dilation > 1 | Set window_dilation to (1,1,1,1) | -| NCC_EVRF019 | Reduce-window wrong operand count | Split into single-operand operations | -| NCC_EVRF022 | Shift-right-arithmetic on non-32-bit | Cast first argument to 32-bit | -| NCC_EVRF024 | Output tensor size exceeds 4GB | Reduce tensor size or use parallelism | -| NCC_EVRF031 | Scatter out-of-bounds (iota size mismatch) | Match iota size to operand dimension | +| Error Code | Error Message | Quick Fix | +| ----------- | ------------------------------------------- | --------------------------------------------------------- | +| NCC_EVRF001 | An unsupported operator was used | Use alternative operator from `neuronx-cc list-operators` | +| NCC_EVRF004 | Complex data types are not supported | Use real-valued tensors or emulate complex arithmetic | +| NCC_EVRF005 | Unsupported F8E4M3FNUZ/F8E5M2FNUZ data type | Convert to float16/bfloat16 | +| NCC_EVRF006 | Unsupported RNG algorithm | Use default RNG via standard APIs | +| NCC_EVRF007 | Instruction count exceeds limit | Apply model parallelism | +| NCC_EVRF009 | Activation memory exceeds HBM limit | Reduce batch size or use parallelism | +| NCC_EVRF010 | Simultaneous input and kernel dilation | Use only input OR kernel dilation | +| NCC_EVRF011 | Strided convolution with dilated input | Remove stride or input dilation | +| NCC_EVRF013 | TopK does not support int32/int64 | Cast to float before TopK | +| NCC_EVRF015 | Unrecognized custom call target | Use supported custom call target | +| NCC_EVRF016 | Scatter-reduce with integer/boolean types | Cast to float types | +| NCC_EVRF017 | Reduce-window with base dilation > 1 | Set base_dilation to (1,1,1,1) | +| NCC_EVRF018 | Reduce-window with window dilation > 1 | Set window_dilation to (1,1,1,1) | +| NCC_EVRF019 | Reduce-window wrong operand count | Split into single-operand operations | +| NCC_EVRF022 | Shift-right-arithmetic on non-32-bit | Cast first argument to 32-bit | +| NCC_EVRF024 | Output tensor size exceeds 4GB | Reduce tensor size or use parallelism | +| NCC_EVRF031 | Scatter out-of-bounds (iota size mismatch) | Match iota size to operand dimension | ### NCC_EOOM - Out of Memory Errors (2 codes) @@ -46,12 +47,13 @@ Memory requirements exceed hardware limits. **Detailed reference**: `ncc-memory-resource-errors.md` -| Error Code | Error Message | Quick Fix | -|------------|---------------|-----------| +| Error Code | Error Message | Quick Fix | +| ----------- | ------------------------------------- | -------------------------------------------------- | | NCC_EOOM001 | Model tensor memory exceeds HBM limit | Reduce batch size, use tensor/pipeline parallelism | -| NCC_EOOM002 | Memory limit exceeded | Reduce batch size, use tensor/pipeline parallelism | +| NCC_EOOM002 | Memory limit exceeded | Reduce batch size, use tensor/pipeline parallelism | **Hardware HBM Limits**: + - Trn1/Trn2/Trn3: 32 GB per device - Inf2: 32 GB per device @@ -59,41 +61,41 @@ Memory requirements exceed hardware limits. **Detailed reference**: `ncc-type-operation-errors.md` -| Error Code | Error Message | Quick Fix | -|------------|---------------|-----------| +| Error Code | Error Message | Quick Fix | +| ----------- | ----------------------------- | ------------------------------------------- | | NCC_EARG001 | Unsupported LNC configuration | Use supported LNC count for target hardware | **Supported LNC Configurations by Hardware**: -| Hardware | Supported LNC Values | -|----------|---------------------| -| Trn1 (gen2) | 1 | -| Inf2 (gen2) | 1, 2 | -| Trn2 (gen3) | 1, 2, 4 | -| Trn3 (gen4) | 1, 2, 4, 8 | +| Hardware | Supported LNC Values | +| ----------- | -------------------- | +| Trn1 (gen2) | 1 | +| Inf2 (gen2) | 1, 2 | +| Trn2 (gen3) | 1, 2, 4 | +| Trn3 (gen4) | 1, 2, 4, 8 | ### NCC_EBVF - Buffer Verification Errors (1 code) **Detailed reference**: `ncc-memory-resource-errors.md` -| Error Code | Error Message | Quick Fix | -|------------|---------------|-----------| +| Error Code | Error Message | Quick Fix | +| ----------- | ------------------------------- | ----------------------- | | NCC_EBVF030 | Instruction count exceeds limit | Apply model parallelism | ### NCC_EHCA - Custom Call Errors (1 code) **Detailed reference**: `ncc-type-operation-errors.md` -| Error Code | Error Message | Quick Fix | -|------------|---------------|-----------| +| Error Code | Error Message | Quick Fix | +| ----------- | ------------------------------- | -------------------------------- | | NCC_EHCA005 | Unrecognized custom call target | Use supported custom call target | ### NCC_ESFH - Safe Float Handling Errors (1 code) **Detailed reference**: `ncc-type-operation-errors.md` -| Error Code | Error Message | Quick Fix | -|------------|---------------|-----------| +| Error Code | Error Message | Quick Fix | +| ----------- | ---------------------------------------- | -------------------- | | NCC_ESFH002 | 64-bit constant cannot convert to 32-bit | Use 32-bit constants | ### NCC_ESPP - Shape Parser Errors (2 codes) @@ -102,45 +104,46 @@ Data type support and shape parsing issues. **Detailed reference**: `ncc-type-operation-errors.md` -| Error Code | Error Message | Quick Fix | -|------------|---------------|-----------| -| NCC_ESPP004 | Unsupported data type for codegen | Use fp32/fp16/bf16 | -| NCC_ESPP047 | Unsupported FP8 data type | Convert to float16 or use gen3+ hardware | +| Error Code | Error Message | Quick Fix | +| ----------- | --------------------------------- | ---------------------------------------- | +| NCC_ESPP004 | Unsupported data type for codegen | Use fp32/fp16/bf16 | +| NCC_ESPP047 | Unsupported FP8 data type | Convert to float16 or use gen3+ hardware | **Supported Dtypes by Hardware**: -| Dtype | gen2 (Trn1/Inf2) | gen3+ (Trn2/Trn3) | -|-------|------------------|-------------------| -| fp32, fp16, bf16 | Yes | Yes | -| fp8_e4m3, fp8_e5m2 | No | Yes | +| Dtype | gen2 (Trn1/Inf2) | gen3+ (Trn2/Trn3) | +| ------------------ | ---------------- | ----------------- | +| fp32, fp16, bf16 | Yes | Yes | +| fp8_e4m3, fp8_e5m2 | No | Yes | ### NCC_EUOC - Unsupported Operation Errors (1 code) **Detailed reference**: `ncc-type-operation-errors.md` -| Error Code | Error Message | Quick Fix | -|------------|---------------|-----------| +| Error Code | Error Message | Quick Fix | +| ----------- | -------------------- | ------------------------ | | NCC_EUOC002 | Unsupported operator | Use alternative operator | ### NCC_EXSP - Expansion Errors (1 code) **Detailed reference**: `ncc-type-operation-errors.md` -| Error Code | Error Message | Quick Fix | -|------------|---------------|-----------| +| Error Code | Error Message | Quick Fix | +| ----------- | ----------------------------------- | ---------------------------- | | NCC_EXSP001 | Activation memory exceeds HBM limit | Reduce size, use parallelism | ### NCC_EXTP - Expansion Tensor Errors (1 code) **Detailed reference**: `ncc-memory-resource-errors.md` -| Error Code | Error Message | Quick Fix | -|------------|---------------|-----------| +| Error Code | Error Message | Quick Fix | +| ----------- | ------------------------------- | ----------------------- | | NCC_EXTP004 | Instruction count exceeds limit | Apply model parallelism | ## Hardware Compatibility All error codes apply to: + - **Inf1**: Inferentia 1 - **Inf2**: Inferentia 2 - **Trn1**: Trainium 1 @@ -153,28 +156,28 @@ Check individual error code documentation for hardware-specific notes. ### By Symptom -| Symptom | Likely Error Code | Quick Fix | -|---------|-------------------|-----------| -| Unsupported operator | NCC_EVRF001, NCC_EUOC002 | Find alternative operator | -| Out of memory | NCC_EOOM001, NCC_EOOM002, NCC_EXSP001 | Reduce batch size, use parallelism | -| FP8 type error | NCC_EVRF005, NCC_ESPP047 | Check hardware generation, convert type | -| Instruction count too high | NCC_EVRF007, NCC_EBVF030, NCC_EXTP004 | Apply model parallelism | -| Convolution dilation error | NCC_EVRF010, NCC_EVRF011 | Use single dilation type | -| Data type not supported | NCC_EVRF004, NCC_ESPP004 | Use supported dtype | -| Tensor size exceeds limit | NCC_EVRF024 | Reduce tensor dimensions | -| Custom call error | NCC_EVRF015, NCC_EHCA005 | Use supported target name | -| LNC configuration error | NCC_EARG001 | Use supported LNC for hardware | +| Symptom | Likely Error Code | Quick Fix | +| -------------------------- | ------------------------------------- | --------------------------------------- | +| Unsupported operator | NCC_EVRF001, NCC_EUOC002 | Find alternative operator | +| Out of memory | NCC_EOOM001, NCC_EOOM002, NCC_EXSP001 | Reduce batch size, use parallelism | +| FP8 type error | NCC_EVRF005, NCC_ESPP047 | Check hardware generation, convert type | +| Instruction count too high | NCC_EVRF007, NCC_EBVF030, NCC_EXTP004 | Apply model parallelism | +| Convolution dilation error | NCC_EVRF010, NCC_EVRF011 | Use single dilation type | +| Data type not supported | NCC_EVRF004, NCC_ESPP004 | Use supported dtype | +| Tensor size exceeds limit | NCC_EVRF024 | Reduce tensor dimensions | +| Custom call error | NCC_EVRF015, NCC_EHCA005 | Use supported target name | +| LNC configuration error | NCC_EARG001 | Use supported LNC for hardware | ### By Operation Type -| Operation | Common Errors | Notes | -|-----------|---------------|-------| -| Matrix operations | NCC_EVRF001, NCC_EUOC002 | Some ops not supported | -| Convolutions | NCC_EVRF010, NCC_EVRF011 | Dilation restrictions | -| Reductions | NCC_EVRF017, NCC_EVRF018 | Window dilation limits | -| Scatter/Gather | NCC_EVRF016, NCC_EVRF031 | Type and bounds checks | -| Type casting | NCC_EVRF004, NCC_EVRF005 | Limited dtype support | -| Random number generation | NCC_EVRF006 | Algorithm restrictions | +| Operation | Common Errors | Notes | +| ------------------------ | ------------------------ | ---------------------- | +| Matrix operations | NCC_EVRF001, NCC_EUOC002 | Some ops not supported | +| Convolutions | NCC_EVRF010, NCC_EVRF011 | Dilation restrictions | +| Reductions | NCC_EVRF017, NCC_EVRF018 | Window dilation limits | +| Scatter/Gather | NCC_EVRF016, NCC_EVRF031 | Type and bounds checks | +| Type casting | NCC_EVRF004, NCC_EVRF005 | Limited dtype support | +| Random number generation | NCC_EVRF006 | Algorithm restrictions | ## Error Resolution Workflow @@ -230,11 +233,11 @@ neuronx-cc ← NCC_* errors occur here NEFF File ``` -NCC_* errors occur during the neuronx-cc phase when generating NEFF files from the intermediate representation. +NCC\_\* errors occur during the neuronx-cc phase when generating NEFF files from the intermediate representation. ## Best Practices -**When encountering NCC_* errors**: +**When encountering NCC\_\* errors**: 1. **Read the full error message** - Contains context and file/line info 2. **Check detailed reference** - See linked files for code examples @@ -243,6 +246,7 @@ NCC_* errors occur during the neuronx-cc phase when generating NEFF files from t 5. **Simplify if complex** - Break large kernels into smaller pieces **Prevention**: + - Use supported operations for target hardware - Check dtype compatibility before compilation - Monitor memory usage for large models diff --git a/skills/neuron-nki-debugging/references/compiler-flags.md b/skills/neuron-nki-debugging/references/compiler-flags.md index f15f649..78eb486 100644 --- a/skills/neuron-nki-debugging/references/compiler-flags.md +++ b/skills/neuron-nki-debugging/references/compiler-flags.md @@ -13,32 +13,32 @@ os.environ["NEURON_CC_FLAGS"] = "--target trn2 --lnc 1" ## Core Flags -| Flag | Values | Required | Description | -|------|--------|----------|-------------| -| `--target` | trn1, trn1n, trn2, trn3, inf2 | Yes | Target hardware platform | -| `--lnc` | 1, 2 | Recommended | Logical NeuronCore count | -| `--verbose` | info, warning, error, debug | No | Logging verbosity level | +| Flag | Values | Required | Description | +| ----------- | ----------------------------- | ----------- | ------------------------ | +| `--target` | trn1, trn1n, trn2, trn3, inf2 | Yes | Target hardware platform | +| `--lnc` | 1, 2 | Recommended | Logical NeuronCore count | +| `--verbose` | info, warning, error, debug | No | Logging verbosity level | ### --target Specifies the target Neuron hardware platform. Must match the `platform_target` in your `@nki.jit` decorator. -| Target | Hardware | Generation | FP8 Support | -|--------|----------|------------|-------------| -| `trn1` | Trainium 1 | gen2 | No | -| `trn1n` | Trainium 1n | gen2 | No | -| `inf2` | Inferentia 2 | gen2 | No | -| `trn2` | Trainium 2 | gen3 | Yes | -| `trn3` | Trainium 3 | gen4 | Yes | +| Target | Hardware | Generation | FP8 Support | +| ------- | ------------ | ---------- | ----------- | +| `trn1` | Trainium 1 | gen2 | No | +| `trn1n` | Trainium 1n | gen2 | No | +| `inf2` | Inferentia 2 | gen2 | No | +| `trn2` | Trainium 2 | gen3 | Yes | +| `trn3` | Trainium 3 | gen4 | Yes | ### --lnc (Logical NeuronCore) Controls how many NeuronCores the kernel is sharded across. -| Value | Use Case | -|-------|----------| -| `1` | Single-core debugging (recommended for development) | -| `2` | Multi-core execution (default on trn2/trn3) | +| Value | Use Case | +| ----- | --------------------------------------------------- | +| `1` | Single-core debugging (recommended for development) | +| `2` | Multi-core execution (default on trn2/trn3) | **Recommendation:** Use `--lnc 1` during debugging for simpler error messages and faster compilation. @@ -46,12 +46,12 @@ Controls how many NeuronCores the kernel is sharded across. Controls compiler output verbosity. -| Level | Output | Use When | -|-------|--------|----------| -| `info` | Progress messages | Standard debugging | -| `warning` | Diagnostic warnings | Default behavior | -| `error` | Compilation errors only | Minimal output | -| `debug` | Extensive internal info | Deep debugging | +| Level | Output | Use When | +| --------- | ----------------------- | ------------------ | +| `info` | Progress messages | Standard debugging | +| `warning` | Diagnostic warnings | Default behavior | +| `error` | Compilation errors only | Minimal output | +| `debug` | Extensive internal info | Deep debugging | ## Standard Debugging Configuration @@ -61,6 +61,7 @@ os.environ["NEURON_CC_FLAGS"] = "--target trn2 --lnc 1" ``` This configuration: + - Targets Trainium 2 hardware (gen3) - Uses single NeuronCore for simpler debugging - Uses default verbosity (warning) @@ -121,12 +122,12 @@ os.environ["NEURON_CC_FLAGS"] = "--target trn2 --lnc 1 --verbose=info" ## Common Flag Combinations -| Scenario | Flags | -|----------|-------| -| Basic debugging | `--target trn2 --lnc 1` | -| Verbose debugging | `--target trn2 --lnc 1 --verbose=info` | -| Multi-core test | `--target trn2 --lnc 2` | -| Production build | `--target trn2` (uses platform defaults) | +| Scenario | Flags | +| ----------------- | ---------------------------------------- | +| Basic debugging | `--target trn2 --lnc 1` | +| Verbose debugging | `--target trn2 --lnc 1 --verbose=info` | +| Multi-core test | `--target trn2 --lnc 2` | +| Production build | `--target trn2` (uses platform defaults) | ## Notes diff --git a/skills/neuron-nki-debugging/references/ncc-memory-resource-errors.md b/skills/neuron-nki-debugging/references/ncc-memory-resource-errors.md index a325fd0..8cf2404 100644 --- a/skills/neuron-nki-debugging/references/ncc-memory-resource-errors.md +++ b/skills/neuron-nki-debugging/references/ncc-memory-resource-errors.md @@ -4,12 +4,12 @@ Detailed reference for Neuron Compiler memory and resource limit errors. These e ## Hardware Memory Limits -| Hardware | HBM per Device | Notes | -|----------|----------------|-------| -| Trn1 (gen2) | 32 GB | 2 NeuronCores per device | -| Trn2 (gen3) | 32 GB | Enhanced compute capabilities | -| Trn3 (gen4) | 32 GB | Latest generation | -| Inf2 (gen2) | 32 GB | Inference optimized | +| Hardware | HBM per Device | Notes | +| ----------- | -------------- | ----------------------------- | +| Trn1 (gen2) | 32 GB | 2 NeuronCores per device | +| Trn2 (gen3) | 32 GB | Enhanced compute capabilities | +| Trn3 (gen4) | 32 GB | Latest generation | +| Inf2 (gen2) | 32 GB | Inference optimized | ## NCC_EOOM001 - Model Tensor Memory Exceeded @@ -20,6 +20,7 @@ Detailed reference for Neuron Compiler memory and resource limit errors. These e **Cause**: Total memory usage from I/O tensors, internal allocations, and SBUF spills exceeds available HBM. **Memory Components**: + - **I/O tensors**: Input and output activation tensors - **Internal allocations**: Scratchpad memory for intermediate computations - **SBUF spills**: Data that cannot fit in on-chip SBUF memory and must spill to HBM @@ -89,11 +90,13 @@ model.encoder = checkpoint_wrapper(model.encoder) **Error message**: The combined memory needed for the model tensors exceeds the high-bandwidth memory limit. **Cause**: Same as NCC_EOOM001. Memory usage components: + - I/O tensors - Internal allocations - SBUF spills **Resolution**: Same strategies as NCC_EOOM001: + 1. Reduce batch/tensor size 2. Use pipeline parallelism 3. Use tensor parallelism @@ -162,6 +165,7 @@ self.fc2 = RowParallelLinear(intermediate_size, hidden_size) **Cause**: Same as NCC_EBVF030, but occurs during tensor expansion phase. The expanded kernel exceeds instruction limits. **Resolution**: Same strategies as NCC_EBVF030: + 1. Apply pipeline parallelism 2. Apply tensor parallelism 3. Simplify kernel logic @@ -172,16 +176,16 @@ self.fc2 = RowParallelLinear(intermediate_size, hidden_size) ## Quick Reference -| Error Code | Phase | Summary | Quick Fix | -|------------|-------|---------|-----------| -| EOOM001 | Memory allocation | Model tensors exceed HBM | Reduce batch size, use parallelism | -| EOOM002 | Memory allocation | Memory limit exceeded | Reduce batch size, use parallelism | -| EBVF030 | Buffer verification | Instruction count exceeded | Model parallelism | -| EXTP004 | Tensor expansion | Instruction count exceeded | Model parallelism | +| Error Code | Phase | Summary | Quick Fix | +| ---------- | ------------------- | -------------------------- | ---------------------------------- | +| EOOM001 | Memory allocation | Model tensors exceed HBM | Reduce batch size, use parallelism | +| EOOM002 | Memory allocation | Memory limit exceeded | Reduce batch size, use parallelism | +| EBVF030 | Buffer verification | Instruction count exceeded | Model parallelism | +| EXTP004 | Tensor expansion | Instruction count exceeded | Model parallelism | ## Common Patterns -### Memory Errors (EOOM*) +### Memory Errors (EOOM\*) All memory errors share the same resolution strategies: @@ -227,6 +231,6 @@ if estimated_gb > 32: # HBM limit ## Related References -- `compiler-error-codes.md` - Quick reference index for all NCC_* errors +- `compiler-error-codes.md` - Quick reference index for all NCC\_\* errors - `ncc-verification-errors.md` - Verification errors (including EVRF007, EVRF009, EVRF024) - `ncc-type-operation-errors.md` - Type and operation errors diff --git a/skills/neuron-nki-debugging/references/ncc-type-operation-errors.md b/skills/neuron-nki-debugging/references/ncc-type-operation-errors.md index c43d501..41db35c 100644 --- a/skills/neuron-nki-debugging/references/ncc-type-operation-errors.md +++ b/skills/neuron-nki-debugging/references/ncc-type-operation-errors.md @@ -5,17 +5,18 @@ Detailed reference for Neuron Compiler type, argument, and operation errors. The ## Supported Data Types by Hardware | Data Type | gen2 (Trn1/Inf2) | gen3 (Trn2) | gen4 (Trn3) | -|-----------|------------------|-------------|-------------| -| float32 | Yes | Yes | Yes | -| float16 | Yes | Yes | Yes | -| bfloat16 | Yes | Yes | Yes | -| int32 | Yes | Yes | Yes | -| int16 | Yes | Yes | Yes | -| int8 | Yes | Yes | Yes | -| fp8_e4m3 | No | Yes | Yes | -| fp8_e5m2 | No | Yes | Yes | +| --------- | ---------------- | ----------- | ----------- | +| float32 | Yes | Yes | Yes | +| float16 | Yes | Yes | Yes | +| bfloat16 | Yes | Yes | Yes | +| int32 | Yes | Yes | Yes | +| int16 | Yes | Yes | Yes | +| int8 | Yes | Yes | Yes | +| fp8_e4m3 | No | Yes | Yes | +| fp8_e5m2 | No | Yes | Yes | **Unsupported types (all hardware)**: + - complex64, complex128 - float8_e4m3fnuz, float8_e4m3b11fnuz, float8_e5m2fnuz - float4_e2m1fn @@ -32,12 +33,12 @@ Detailed reference for Neuron Compiler type, argument, and operation errors. The ### Supported LNC Configurations -| Hardware | Supported LNC Values | -|----------|---------------------| -| Trn1 (gen2) | 1 | -| Inf2 (gen2) | 1, 2 | -| Trn2 (gen3) | 1, 2, 4 | -| Trn3 (gen4) | 1, 2, 4, 8 | +| Hardware | Supported LNC Values | +| ----------- | -------------------- | +| Trn1 (gen2) | 1 | +| Inf2 (gen2) | 1, 2 | +| Trn2 (gen3) | 1, 2, 4 | +| Trn3 (gen4) | 1, 2, 4, 8 | ### Understanding LNC @@ -173,6 +174,7 @@ input_tensor = input_tensor.to(torch.float16) ### Recognized Custom Call Targets (28 total) **Activation Functions**: + - `AwsNeuronErf` - `AwsNeuronGelu` - `AwsNeuronGeluApprxTanh` @@ -181,11 +183,13 @@ input_tensor = input_tensor.to(torch.float16) - `AwsNeuronSiluBackward` **Normalization**: + - `AwsNeuronRmsNorm` - `AwsNeuronSoftmax` - `AwsNeuronSoftmaxBackward` **Compute Operations**: + - `AwsNeuronCollectiveMatmul` - `AwsNeuronIntMatmult` - `AwsNeuronArgMax` @@ -193,21 +197,25 @@ input_tensor = input_tensor.to(torch.float16) - `AwsNeuronTopK` **Utility Operations**: + - `AwsNeuronDropoutMaskV1` - `AwsNeuronCustomNativeKernel` - `AwsNeuronCustomOp` - `AwsNeuronDevicePrint` **Resize Operations**: + - `ResizeNearest` - `ResizeBilinear` - `ResizeNearestGrad` **Sharding and Communication**: + - `AwsNeuronLNCShardingConstraint` - `AwsNeuronTransferWithStaticRing` **Module Markers**: + - `AwsNeuronModuleMarkerStart-Forward` - `AwsNeuronModuleMarkerStart-Backward` - `AwsNeuronModuleMarkerEnd-Forward` @@ -261,10 +269,10 @@ def lowering(ctx, x_val): ### 32-Bit Integer Limits -| Type | Min | Max | -|------|-----|-----| -| int32 | -2,147,483,648 | 2,147,483,647 | -| uint32 | 0 | 4,294,967,295 | +| Type | Min | Max | +| ------ | -------------- | ------------- | +| int32 | -2,147,483,648 | 2,147,483,647 | +| uint32 | 0 | 4,294,967,295 | ### Before (error) @@ -289,6 +297,7 @@ def test(): ``` **Note**: If you need to work with values > 4.29 billion, consider: + - Using multiple 32-bit operations - Representing values in a different scale - Offloading to CPU for 64-bit arithmetic @@ -327,11 +336,11 @@ class Model(torch.nn.Module): ### Common Unsupported Operators and Alternatives -| Unsupported | Alternative | -|-------------|-------------| -| `triangular_solve` | `inverse` + matrix multiply | -| Complex FFT | Split into real/imaginary parts | -| Some custom CUDA kernels | Rewrite using supported ops | +| Unsupported | Alternative | +| ------------------------ | ------------------------------- | +| `triangular_solve` | `inverse` + matrix multiply | +| Complex FFT | Split into real/imaginary parts | +| Some custom CUDA kernels | Rewrite using supported ops | **See also**: NCC_EVRF001 (same error message and resolution) @@ -346,6 +355,7 @@ class Model(torch.nn.Module): **Cause**: During tensor expansion phase, memory requirements exceed HBM limits. **Resolution**: Same strategies as memory errors: + 1. Reduce batch/tensor size 2. Use pipeline/tensor parallelism via neuronx-distributed @@ -386,18 +396,18 @@ class ParallelSelfAttention(transformers.models.bert.modeling_bert.BertSelfAtten ## Quick Reference -| Error Code | Category | Summary | Quick Fix | -|------------|----------|---------|-----------| -| EARG001 | Configuration | Unsupported LNC config | Use supported LNC for target hardware | -| ESPP004 | Data Type | Unsupported dtype for codegen | Use fp32/fp16/bf16 | -| ESPP047 | Data Type | Unsupported FP8 type | Convert to float16 or check gen3+ | -| EHCA005 | Custom Call | Unrecognized target | Use supported custom call target | -| ESFH002 | Constants | 64-bit constant overflow | Use uint32 constants | -| EUOC002 | Operator | Unsupported operator | Use alternative operator | -| EXSP001 | Memory | Expansion memory exceeded | Reduce size, use parallelism | +| Error Code | Category | Summary | Quick Fix | +| ---------- | ------------- | ----------------------------- | ------------------------------------- | +| EARG001 | Configuration | Unsupported LNC config | Use supported LNC for target hardware | +| ESPP004 | Data Type | Unsupported dtype for codegen | Use fp32/fp16/bf16 | +| ESPP047 | Data Type | Unsupported FP8 type | Convert to float16 or check gen3+ | +| EHCA005 | Custom Call | Unrecognized target | Use supported custom call target | +| ESFH002 | Constants | 64-bit constant overflow | Use uint32 constants | +| EUOC002 | Operator | Unsupported operator | Use alternative operator | +| EXSP001 | Memory | Expansion memory exceeded | Reduce size, use parallelism | ## Related References -- `compiler-error-codes.md` - Quick reference index for all NCC_* errors -- `ncc-verification-errors.md` - Verification errors (EVRF*) +- `compiler-error-codes.md` - Quick reference index for all NCC\_\* errors +- `ncc-verification-errors.md` - Verification errors (EVRF\*) - `ncc-memory-resource-errors.md` - Memory and resource limit errors diff --git a/skills/neuron-nki-debugging/references/ncc-verification-errors.md b/skills/neuron-nki-debugging/references/ncc-verification-errors.md index 35114fe..dbfcd03 100644 --- a/skills/neuron-nki-debugging/references/ncc-verification-errors.md +++ b/skills/neuron-nki-debugging/references/ncc-verification-errors.md @@ -1,4 +1,4 @@ -# NCC Verification Errors (NCC_EVRF*) +# NCC Verification Errors (NCC_EVRF\*) Detailed reference for Neuron Compiler verification errors. These errors occur when the compiler detects unsupported operations, data types, or configurations during verification. @@ -90,6 +90,7 @@ input_tensor = input_tensor.to(torch.float16) ``` **Supported dtypes by hardware**: + - gen2 (Trn1/Inf2): fp32, fp16, bf16 (no FP8) - gen3/gen4 (Trn2/Trn3): fp32, fp16, bf16, fp8_e4m3, fp8_e5m2 @@ -118,6 +119,7 @@ input_tensor = input_tensor.to(torch.float16) **Resolution**: Apply model parallelism to break large computational graphs into smaller subgraphs. **Strategies**: + - Use pipeline parallelism via neuronx-distributed - Use tensor parallelism to shard across devices - Simplify kernel logic or split into multiple kernels @@ -296,15 +298,15 @@ def forward(self, x): **Recognized Custom Call Targets**: -| Category | Targets | -|----------|---------| -| Activation | `AwsNeuronErf`, `AwsNeuronGelu`, `AwsNeuronGeluApprxTanh`, `AwsNeuronGeluBackward`, `AwsNeuronSilu`, `AwsNeuronSiluBackward` | -| Normalization | `AwsNeuronRmsNorm`, `AwsNeuronSoftmax`, `AwsNeuronSoftmaxBackward` | -| Compute | `AwsNeuronCollectiveMatmul`, `AwsNeuronIntMatmult`, `AwsNeuronArgMax`, `AwsNeuronArgMin`, `AwsNeuronTopK` | -| Utility | `AwsNeuronDropoutMaskV1`, `AwsNeuronCustomNativeKernel`, `AwsNeuronCustomOp`, `AwsNeuronDevicePrint` | -| Resize | `ResizeNearest`, `ResizeBilinear`, `ResizeNearestGrad` | -| Sharding | `AwsNeuronLNCShardingConstraint`, `AwsNeuronTransferWithStaticRing` | -| Markers | `AwsNeuronModuleMarkerStart-Forward`, `AwsNeuronModuleMarkerStart-Backward`, `AwsNeuronModuleMarkerEnd-Forward`, `AwsNeuronModuleMarkerEnd-Backward`, `NeuronBoundaryMarker-Start`, `NeuronBoundaryMarker-End` | +| Category | Targets | +| ------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| Activation | `AwsNeuronErf`, `AwsNeuronGelu`, `AwsNeuronGeluApprxTanh`, `AwsNeuronGeluBackward`, `AwsNeuronSilu`, `AwsNeuronSiluBackward` | +| Normalization | `AwsNeuronRmsNorm`, `AwsNeuronSoftmax`, `AwsNeuronSoftmaxBackward` | +| Compute | `AwsNeuronCollectiveMatmul`, `AwsNeuronIntMatmult`, `AwsNeuronArgMax`, `AwsNeuronArgMin`, `AwsNeuronTopK` | +| Utility | `AwsNeuronDropoutMaskV1`, `AwsNeuronCustomNativeKernel`, `AwsNeuronCustomOp`, `AwsNeuronDevicePrint` | +| Resize | `ResizeNearest`, `ResizeBilinear`, `ResizeNearestGrad` | +| Sharding | `AwsNeuronLNCShardingConstraint`, `AwsNeuronTransferWithStaticRing` | +| Markers | `AwsNeuronModuleMarkerStart-Forward`, `AwsNeuronModuleMarkerStart-Backward`, `AwsNeuronModuleMarkerEnd-Forward`, `AwsNeuronModuleMarkerEnd-Backward`, `NeuronBoundaryMarker-Start`, `NeuronBoundaryMarker-End` | ### Before (error) @@ -647,28 +649,28 @@ result = lax.scatter( ## Quick Reference -| Error Code | Summary | Quick Fix | -|------------|---------|-----------| -| EVRF001 | Unsupported operator | Use alternative operator, check `neuronx-cc list-operators` | -| EVRF004 | Complex data types | Offload to CPU or emulate with real/imag parts | -| EVRF005 | Unsupported FP8 types | Convert to float16/bfloat16 | -| EVRF006 | Unsupported RNG algorithm | Use default RNG | -| EVRF007 | Instruction limit exceeded | Apply model parallelism | -| EVRF009 | Activation memory exceeded | Reduce batch size or use parallelism | -| EVRF010 | Simultaneous dilation | Use input OR kernel dilation, not both | -| EVRF011 | Strided + dilated input | Remove stride or input dilation | -| EVRF013 | TopK integer inputs | Cast to float before TopK | -| EVRF015 | Unrecognized custom call | Use supported custom call target | -| EVRF016 | Scatter-reduce int/bool | Cast to float types | -| EVRF017 | Reduce-window base dilation | Set base_dilation to (1,1,1,1) | -| EVRF018 | Reduce-window window dilation | Set window_dilation to (1,1,1,1) | -| EVRF019 | Reduce-window wrong operands | Split into single-operand operations | -| EVRF022 | Shift-right non-32-bit | Cast first argument to 32-bit | -| EVRF024 | Output tensor > 4GB | Reduce tensor size or use parallelism | -| EVRF031 | Scatter out-of-bounds | Match iota size to operand dimension | +| Error Code | Summary | Quick Fix | +| ---------- | ----------------------------- | ----------------------------------------------------------- | +| EVRF001 | Unsupported operator | Use alternative operator, check `neuronx-cc list-operators` | +| EVRF004 | Complex data types | Offload to CPU or emulate with real/imag parts | +| EVRF005 | Unsupported FP8 types | Convert to float16/bfloat16 | +| EVRF006 | Unsupported RNG algorithm | Use default RNG | +| EVRF007 | Instruction limit exceeded | Apply model parallelism | +| EVRF009 | Activation memory exceeded | Reduce batch size or use parallelism | +| EVRF010 | Simultaneous dilation | Use input OR kernel dilation, not both | +| EVRF011 | Strided + dilated input | Remove stride or input dilation | +| EVRF013 | TopK integer inputs | Cast to float before TopK | +| EVRF015 | Unrecognized custom call | Use supported custom call target | +| EVRF016 | Scatter-reduce int/bool | Cast to float types | +| EVRF017 | Reduce-window base dilation | Set base_dilation to (1,1,1,1) | +| EVRF018 | Reduce-window window dilation | Set window_dilation to (1,1,1,1) | +| EVRF019 | Reduce-window wrong operands | Split into single-operand operations | +| EVRF022 | Shift-right non-32-bit | Cast first argument to 32-bit | +| EVRF024 | Output tensor > 4GB | Reduce tensor size or use parallelism | +| EVRF031 | Scatter out-of-bounds | Match iota size to operand dimension | ## Related References -- `compiler-error-codes.md` - Quick reference index for all NCC_* errors +- `compiler-error-codes.md` - Quick reference index for all NCC\_\* errors - `ncc-memory-resource-errors.md` - Memory and resource limit errors - `ncc-type-operation-errors.md` - Type and operation errors diff --git a/skills/neuron-nki-debugging/references/neuron-core-isolation.md b/skills/neuron-nki-debugging/references/neuron-core-isolation.md index 38204f5..eaf4705 100644 --- a/skills/neuron-nki-debugging/references/neuron-core-isolation.md +++ b/skills/neuron-nki-debugging/references/neuron-core-isolation.md @@ -14,8 +14,8 @@ TOTAL_CORES=$(neuron-ls 2>/dev/null | grep -c "NeuronCore" || echo "0") export NEURON_RT_VISIBLE_CORES="0" ``` -| Environment Variable | Purpose | -|---------------------|---------| +| Environment Variable | Purpose | +| ------------------------- | ------------------------------------------------------------ | | `NEURON_RT_VISIBLE_CORES` | Comma-separated list of core indices visible to this process | ## Allocation Strategy diff --git a/skills/neuron-nki-docs/SKILL.md b/skills/neuron-nki-docs/SKILL.md index 30c0560..aaa9e73 100644 --- a/skills/neuron-nki-docs/SKILL.md +++ b/skills/neuron-nki-docs/SKILL.md @@ -23,14 +23,14 @@ This skill provides comprehensive access to NKI (Neuron Kernel Interface) docume Route queries to the appropriate documentation based on the query pattern: -| Query Pattern | Start With | Then Read | -|---------------|------------|-----------| -| `nl.*` or `nisa.*` or API name | `references/indices/symbol-lookup.md` | Linked API doc | -| "how to [task]" | `references/indices/task-routing.md` | Linked tutorial | -| "NCC_*" or error code | `references/debugging/error-codes-index.md` | Specific error doc | -| gen2/gen3/gen4/trn1/trn2/trn3 | `references/architecture/` | Specific arch doc | -| optimize/profile/performance | `references/optimization/` | Relevant guide | -| browse/list all docs | `references/indices/hierarchical-toc.md` | Navigate tree | +| Query Pattern | Start With | Then Read | +| ------------------------------ | ------------------------------------------- | ------------------ | +| `nl.*` or `nisa.*` or API name | `references/indices/symbol-lookup.md` | Linked API doc | +| "how to [task]" | `references/indices/task-routing.md` | Linked tutorial | +| "NCC\_\*" or error code | `references/debugging/error-codes-index.md` | Specific error doc | +| gen2/gen3/gen4/trn1/trn2/trn3 | `references/architecture/` | Specific arch doc | +| optimize/profile/performance | `references/optimization/` | Relevant guide | +| browse/list all docs | `references/indices/hierarchical-toc.md` | Navigate tree | ## Search Strategy @@ -46,47 +46,53 @@ Route queries to the appropriate documentation based on the query pattern: ## Directory Guide -| Directory | Contents | -|-----------|----------| -| `references/architecture/` | Hardware architecture guides for Trainium/Inferentia generations | -| `references/debugging/` | Error codes index and individual error documentation | -| `references/downloads/` | Valid Python kernel examples (deprecated patterns excluded) | -| `references/indices/` | Navigation aids and lookup tables | -| `references/optimization/` | Performance tuning, profiling, migration guides | -| `references/programming/` | Core NKI concepts, tutorials, and API reference | -| `references/programming/api/` | Detailed API documentation by category | -| `references/programming/tutorials/` | Step-by-step kernel implementation tutorials | -| `references/reference/` | FAQ and release notes | +| Directory | Contents | +| ----------------------------------- | ---------------------------------------------------------------- | +| `references/architecture/` | Hardware architecture guides for Trainium/Inferentia generations | +| `references/debugging/` | Error codes index and individual error documentation | +| `references/downloads/` | Valid Python kernel examples (deprecated patterns excluded) | +| `references/indices/` | Navigation aids and lookup tables | +| `references/optimization/` | Performance tuning, profiling, migration guides | +| `references/programming/` | Core NKI concepts, tutorials, and API reference | +| `references/programming/api/` | Detailed API documentation by category | +| `references/programming/tutorials/` | Step-by-step kernel implementation tutorials | +| `references/reference/` | FAQ and release notes | ## Quick API Reference ### Most Common APIs **Data Movement:** + - `nki.isa.dma_copy` - Load/store data between HBM and SBUF - `nki.isa.dma_transpose` - Transpose during DMA transfer **Tensor Operations:** + - `nki.isa.nc_matmul` - Matrix multiplication on Tensor Engine - `nki.isa.tensor_tensor` - Element-wise operations - `nki.isa.tensor_scalar` - Broadcast scalar operations - `nki.isa.tensor_reduce` - Reduction along axes **Memory Allocation:** + - `nl.ndarray` - Create tensor in SBUF - `nl.zeros` - Create zero-initialized tensor **Loop Constructs:** + - `range` - Standard loop iterator (recommended) - `nl.affine_range` / `nl.sequential_range` / `nl.static_range` - Legacy aliases for `range` (all have identical effect in NKI 0.3.0+) **SPMD:** + - `nl.program_id` - Get current program index - `nl.num_programs` - Get total program count ### Module Aliases The documentation uses these standard aliases: + ```python import nki import nki.language as nl @@ -109,12 +115,12 @@ result = nisa.activation(op=nl.exp, data=tensor, scale=scale_tensor, bias=bias_t Critical limits to remember when answering questions: -| Constraint | Limit | Notes | -|------------|-------|-------| -| Partition dimension (P) | ≤ 128 | First axis of on-chip tensors | -| PSUM free dimension (F) | 512 (gen2/3); gen4: 4096 fp32 / 8192 bf16 | Second axis in PSUM buffer | -| SBUF free dimension (F) | ≤ 32767 | Second axis in SBUF buffer | -| MatMul K dimension | ≤ 2048 | Contraction dimension per tile | +| Constraint | Limit | Notes | +| ----------------------- | ----------------------------------------- | ------------------------------ | +| Partition dimension (P) | ≤ 128 | First axis of on-chip tensors | +| PSUM free dimension (F) | 512 (gen2/3); gen4: 4096 fp32 / 8192 bf16 | Second axis in PSUM buffer | +| SBUF free dimension (F) | ≤ 32767 | Second axis in SBUF buffer | +| MatMul K dimension | ≤ 2048 | Contraction dimension per tile | > The PSUM free-dimension limit is generation- and dtype-gated: 512 on gen2/gen3 > (one PSUM bank), and on gen4 up to 4096 for a `float32` `dst` or 8192 for a @@ -124,11 +130,11 @@ Critical limits to remember when answering questions: ### Hardware Generations -| Generation | Devices | Key Features | -|------------|---------|--------------| -| gen2 (v2) | Trn1, Inf2 | Baseline NKI support | -| gen3 (v3) | Trn2 | FP8 support, Double FP8 mode | -| gen4 (v4) | Trn3 | MXFP8/MXFP4, Quad-MX mode | +| Generation | Devices | Key Features | +| ---------- | ---------- | ---------------------------- | +| gen2 (v2) | Trn1, Inf2 | Baseline NKI support | +| gen3 (v3) | Trn2 | FP8 support, Double FP8 mode | +| gen4 (v4) | Trn3 | MXFP8/MXFP4, Quad-MX mode | ## Source NKI Kernel @@ -145,22 +151,27 @@ The Python example files in `references/downloads/` use the latest NKI API patte ## Common Query Examples ### API Lookup + Query: "What is nisa.nc_matmul?" → Read `indices/symbol-lookup.md` → Find link → Read `programming/api/api-nki-isa-tensor.md#nki-isa-nc_matmul` ### Error Code + Query: "What does NCC_EVRF001 mean?" → Read `debugging/error-codes-index.md` → Read `debugging/error-codes/EVRF001.md` ### Tutorial + Query: "How do I implement matrix multiplication?" → Read `indices/task-routing.md` → Read `programming/tutorials/matrix_multiplication.md` ### Architecture + Query: "What's different in Trainium3?" → Read `architecture/trainium3_arch.md` ### Optimization + Query: "How do I profile my kernel?" → Read `optimization/use-neuron-profile.md` @@ -177,6 +188,7 @@ When answering NKI documentation queries: ## Error Code Format NKI compiler errors follow the pattern `NCC_`: + - `EOOM*` - Out of memory errors - `EVRF*` - Verification/validation errors - `EUOC*` - Unsupported operation errors @@ -187,9 +199,9 @@ The full index is in `debugging/error-codes-index.md` with individual files in ` ## Related Skills -| Skill | Purpose | -|-------|---------| -| `/neuron-nki-writing` | Write NKI kernels from specifications | -| `/neuron-nki-debugging` | Debug compiler errors on device | -| `/neuron-nki-profiling` | Profile kernel performance | +| Skill | Purpose | +| ------------------------------ | ------------------------------------- | +| `/neuron-nki-writing` | Write NKI kernels from specifications | +| `/neuron-nki-debugging` | Debug compiler errors on device | +| `/neuron-nki-profiling` | Profile kernel performance | | `/neuron-nki-profile-querying` | Query and analyze kernel profile data | diff --git a/skills/neuron-nki-docs/references/architecture/nki_arch_guides.md b/skills/neuron-nki-docs/references/architecture/nki_arch_guides.md index f03a3d2..dc72fa2 100644 --- a/skills/neuron-nki-docs/references/architecture/nki_arch_guides.md +++ b/skills/neuron-nki-docs/references/architecture/nki_arch_guides.md @@ -3,20 +3,20 @@ NKI and Neuron Architecture NKI currently supports the following NeuronDevice generations: -* Trainium/Inferentia2, available on AWS `trn1`, `trn1n` and `inf2` instances +- Trainium/Inferentia2, available on AWS `trn1`, `trn1n` and `inf2` instances -* Trainium2, available on AWS `trn2` instances and UltraServers +- Trainium2, available on AWS `trn2` instances and UltraServers -* Trainium3, available on AWS `trn3` instances and UltraServers +- Trainium3, available on AWS `trn3` instances and UltraServers The documents below provide an architecture deep dive of each NeuronDevice generation, with a focus on areas that NKI developers can directly control through kernel implementation. -* [Trainium/Inferentia2 Architecture Guide](trainium_inferentia2_arch.md) serves as a foundational architecture guide for understanding the basics of any NeuronDevice generation. +- [Trainium/Inferentia2 Architecture Guide](trainium_inferentia2_arch.md) serves as a foundational architecture guide for understanding the basics of any NeuronDevice generation. -* [Trainium2 Architecture Guide](trainium2_arch.md) walks through the architecture enhancements when compared to the previous generation. +- [Trainium2 Architecture Guide](trainium2_arch.md) walks through the architecture enhancements when compared to the previous generation. -* [Trainium3 Architecture Guide](trainium3_arch.md) covers the enhancements for the next-generation Trainium ML accelerators. +- [Trainium3 Architecture Guide](trainium3_arch.md) covers the enhancements for the next-generation Trainium ML accelerators. Neuron recommends new NKI developers start with [Trainium/Inferentia2 Architecture Guide](trainium_inferentia2_arch.md) before exploring newer NeuronDevice architecture. @@ -27,4 +27,4 @@ Foundational architecture guide for understanding NeuronDevice basics. Architecture enhancements and improvements in the Trainium2 generation. [Trainium3 Architecture Guide](trainium3_arch.md#trainium3-arch) -Latest architecture features and capabilities in Trainium3 devices. \ No newline at end of file +Latest architecture features and capabilities in Trainium3 devices. diff --git a/skills/neuron-nki-docs/references/architecture/trainium2_arch.md b/skills/neuron-nki-docs/references/architecture/trainium2_arch.md index a085d6d..725c96a 100644 --- a/skills/neuron-nki-docs/references/architecture/trainium2_arch.md +++ b/skills/neuron-nki-docs/references/architecture/trainium2_arch.md @@ -5,17 +5,18 @@ In this guide, we will dive into hardware architecture of third-generation Neuro The diagram below shows a block diagram of a Trainium2 device, which consists of: -* 8 NeuronCores (v3). +- 8 NeuronCores (v3). -* 4 HBM stacks with a total device memory capacity of 96GiB and bandwidth of 3TB/s. +- 4 HBM stacks with a total device memory capacity of 96GiB and bandwidth of 3TB/s. -* 128 DMA (Direct Memory Access) engines to move data within and across devices. +- 128 DMA (Direct Memory Access) engines to move data within and across devices. -* 20 CC-Cores for collective communication. +- 20 CC-Cores for collective communication. -* 4 NeuronLink-v3 for device-to-device collective communication. +- 4 NeuronLink-v3 for device-to-device collective communication. ! + > **Figure: neuron device3** > > An architecture diagram of Trainium2 showing 8 NeuronCore-v3 units arranged in a 4x2 grid, with 128 DMA engines, 20 CC-Cores, HBM memory on all sides, and 4 NeuronLink-v3 interconnects. @@ -23,31 +24,37 @@ The diagram below shows a block diagram of a Trainium2 device, which consists of > This diagram illustrates the architecture of the Trainium2 chip, AWS's second-generation training accelerator. > > **Title and layout**: +> > - "Trainium2" label in large gray text at top left > - The chip is shown as a large rounded rectangle > > **NeuronCore arrangement**: +> > - 8 "NeuronCore-v3" units arranged in a 4-row by 2-column grid > - Each NeuronCore-v3 is shown as a white rounded rectangle > - The cores occupy the central portion of the chip > > **Memory (HBM)**: +> > - Four "HBM" blocks (blue vertical bars) positioned on all sides: > - Two on the left side (serving rows 1-2 and rows 3-4) > - Two on the right side (same arrangement) > - This provides high bandwidth memory access to all cores > > **Support components** (bottom area): +> > - "DMA (x128)": 128 DMA engines shown as stacked gray blocks > - "CC-Core (x20)": 20 Collective Communication cores shown as stacked gray blocks > - "Host PCIe": Host interface block on the right > > **Interconnects** (bottom): +> > - Four "NeuronLink-v3" blocks spanning the bottom, providing inter-chip communication > > The design shows significant scaling from Trainium (2 cores) to Trainium2 (8 cores), with proportionally increased DMA engines (32 to 128), CC-Cores (6 to 20), and NeuronLinks (4 NeuronLink-v2 to 4 NeuronLink-v3). > > **Key Elements:** +> > - **Trainium2**: Second-generation training chip > - **NeuronCore-v3**: 8 next-generation compute cores (4x2 grid) > - **HBM**: 4 High Bandwidth Memory blocks (2 left, 2 right) @@ -66,6 +73,7 @@ For a high-level architecture specification comparison from Trainium1 to Trainiu The figure below is a simplified NeuronCore-v3 diagram of the compute engines and their connectivity to the two on-chip SRAMs, SBUF and PSUM. This is similar to NeuronCore-v2. ! + > **Figure: nki trn2 arch 1** > > A NeuronCore architecture diagram showing the internal components including SBUF, four compute engines (Tensor, Vector, Scalar, GPSIMD), PSUM, and Sync Engine, with HBM external memory. @@ -75,10 +83,12 @@ The figure below is a simplified NeuronCore-v3 diagram of the compute engines an > **Main container** (labeled "NeuronCore" at top): > > **SBUF (State Buffer)** - left side: +> > - Large blue block representing the main on-chip SRAM > - Bidirectional arrows connect to HBM below and all compute engines > > **Compute Engines** - center column (from top to bottom): +> > - **Tensor Engine** (green): Matrix multiplication unit with "SEQ" block to its left > - **Vector Engine** (green): Vector operations unit with "SEQ" block to its left > - **Scalar Engine** (green): Scalar operations unit with "SEQ" block to its left @@ -87,23 +97,28 @@ The figure below is a simplified NeuronCore-v3 diagram of the compute engines an > Each engine has bidirectional arrows connecting to SBUF and receives sequencing from SEQ blocks. > > **PSUM (Partial Sum)** - top right: +> > - Blue block for accumulating matrix multiplication results > - Connected to Tensor Engine with bidirectional arrows > - Also connected to Vector Engine > > **Sync Engine** - bottom right: +> > - Green block for synchronization operations > - Positioned near GPSIMD Engine > > **HBM** - bottom: +> > - Large blue block representing off-chip High Bandwidth Memory > - Bidirectional arrow connects to SBUF > > **Data flow**: +> > - SBUF serves as the central hub connecting HBM to all compute engines > - PSUM provides fast accumulation for Tensor Engine operations > > **Key Elements:** +> > - **NeuronCore**: Main compute unit container > - **SBUF**: State Buffer - main on-chip SRAM (blue) > - **PSUM**: Partial Sum accumulator (blue) @@ -117,13 +132,12 @@ The figure below is a simplified NeuronCore-v3 diagram of the compute engines an NeuronCore-v3 SBUF capacity is **28MiB** (or, 128 partitions of 224KiB), up from 24 MiB in NeuronCore-v2. PSUM capacity remains the same at 2MiB. Engine data-path width and frequency are updated to the following: - -| Device Architecture | Compute Engine | Data-path Width (elements/cycle) | Frequency (GHz) | -| --- | --- | --- | --- | -| Trainium2 | Tensor | 4x128 (dense FP8_E4/FP8_E5 input), 2x128 (dense BF16/FP16 input) or 5x128 (sparse input); 1x128 (output) | 2.4 | -| | Vector | 512 BF16/FP16 input/output; 256 input/output for other data types | 0.96 | -| | Scalar | 128 input/output | 1.2 | -| | GpSimd | | 1.2 | +| Device Architecture | Compute Engine | Data-path Width (elements/cycle) | Frequency (GHz) | +| ------------------- | -------------- | -------------------------------------------------------------------------------------------------------- | --------------- | +| Trainium2 | Tensor | 4x128 (dense FP8_E4/FP8_E5 input), 2x128 (dense BF16/FP16 input) or 5x128 (sparse input); 1x128 (output) | 2.4 | +| | Vector | 512 BF16/FP16 input/output; 256 input/output for other data types | 0.96 | +| | Scalar | 128 input/output | 1.2 | +| | GpSimd | | 1.2 | Next, we will go over major updates to each compute engine. @@ -143,6 +157,7 @@ Logically, TensorE doubles the FP8 matmul performance by doubling the maximum co A double-FP8 matmul can perform a multiplication of a 128x256 matrix and a 256x512 matrix (that is, MxKxN matmul, M=128, K=256, N=512). The figure below shows a visualization of the two input matrices (x and y) and the matmul output matrix (output). The figure also highlights two elements (red and yellow) in the first row of the x matrix and in the first column of the y matrix. These two elements are 128 (K//2) elements apart within the rows and columns. We will use these elements to illustrate the SBUF layout requirements for these matrices next. ! + > **Figure: nki trn2 arch 2** > > A mathematical view of matrix multiplication showing three matrices (x, y, and output) with specific dimensions (M=128, K=256, N=512), with highlighted elements showing the computation pattern. @@ -152,22 +167,26 @@ A double-FP8 matmul can perform a multiplication of a 128x256 matrix and a 256x5 > **Matrix layout**: > > **Top matrix - "y" (blue)**: +> > - Dimensions: N=512 (width) by K=256 (height) > - Label "y" in the center > - A small yellow/orange square marker on the left edge indicates a specific element > - A horizontal dashed line passes through the matrix > > **Bottom left matrix - "x" (green)**: +> > - Dimensions: K=256 (width) by M=128 (height) > - Label "x" in the center > - A small red/pink square marker on the top edge > - A vertical dashed line passes through the matrix > > **Bottom right matrix - "output" (purple)**: +> > - Dimensions: N=512 (width) by M=128 (height) > - Label "output" in the center > > **Dimension annotations**: +> > - "N=512" at the top (width of y and output) > - "K=256" on the right side of y (height of y, width of x) > - "M=128" on the left (height of x and output) @@ -179,6 +198,7 @@ A double-FP8 matmul can perform a multiplication of a 128x256 matrix and a 256x5 > The highlighted markers and dashed lines illustrate how a single element of the output matrix is computed by taking the dot product of a row from x and a column from y. > > **Key Elements:** +> > - **x matrix**: Green input matrix [M=128 x K=256] > - **y matrix**: Blue input matrix [K=256 x N=512] > - **output matrix**: Purple result matrix [M=128 x N=512] @@ -192,6 +212,7 @@ A double-FP8 matmul can perform a multiplication of a 128x256 matrix and a 256x5 These tensors must still fit in the 128-partition SBUF, with each partition feeding data into each row of processing elements inside the TensorE. The contraction of size 256 is therefore split into two dimensions: (1) the partition dimension of size 128 and (2) the most major (slowest) free dimension of size 2. This is illustrated in the figure below. Both the stationary matrix (x in above figure) and the moving matrix (y in above figure) are sliced in two tiles, where the first and second tile correspond to first and second halves of the contraction dimension, respectively. ! + > **Figure: nki trn2 arch 3** > > A diagram showing the tensor layout in SBUF for matrix multiplication, with stationary matrix (green) and moving matrix (blue) both tiled with K/2=128 partition dimension. @@ -199,6 +220,7 @@ These tensors must still fit in the 128-partition SBUF, with each partition feed > This diagram illustrates how the input matrices for matrix multiplication are laid out in the State Buffer (SBUF), showing the tiling along the partition dimension. > > **Left tensor - "stationary (SBUF)" (green)**: +> > - Green rectangular block > - Dimensions: M=128+M=128 (width, showing two M tiles) by K/2=128 (height) > - Small colored markers (red and yellow squares) at the top corners @@ -206,13 +228,15 @@ These tensors must still fit in the 128-partition SBUF, with each partition feed > - Label "stationary (SBUF)" in center > > **Right tensor - "moving (SBUF)" (blue)**: -> - Blue rectangular block +> +> - Blue rectangular block > - Dimensions: N=512+N=512 (width, showing two N tiles) by K/2=128 (height) > - Small colored markers at the top > - A vertical dashed line divides the tensor into two N=512 tiles > - Label "moving (SBUF)" in center > > **Dimension annotations**: +> > - "K/2 =128" on the left side of both tensors (partition dimension) > - "M=128" twice at bottom of stationary tensor > - "N=512" twice at bottom of moving tensor @@ -221,11 +245,13 @@ These tensors must still fit in the 128-partition SBUF, with each partition feed > **Caption**: "Tensor Layout in SBUF" centered below > > The diagram shows that: +> > - Both matrices have their K dimension (256 total) tiled into K/2=128 chunks along the SBUF partition dimension > - The stationary matrix has its M dimension in the free dimension > - The moving matrix has its N dimension in the free dimension > > **Key Elements:** +> > - **stationary (SBUF)**: Green tensor loaded for Tensor Engine [K/2=128 x M=256] > - **moving (SBUF)**: Blue tensor that streams through [K/2=128 x N=1024] > - **K/2=128**: Half of K dimension per tile (partition dimension) @@ -238,6 +264,7 @@ These tensors must still fit in the 128-partition SBUF, with each partition feed Next, we invoke the LoadStationary and MultiplyMoving instructions to perform the matrix multiplications using the above tensors in SBUF. This is illustrated in figure below. The LoadStationary instruction loads the stationary tensor (K/2=128, 2, M=128) into TensorE, which stores two data elements into a single processing element (for example, the red and yellow elements land in the first processing element of TensorE as shown in ❶). Next, the MultiplyMoving instruction streams the moving tensor horizontally across the loaded stationary tensor. Similar to LoadStationary, two elements of moving tensor are sent to the same processing element simultaneously as shown in ❷, such that they can get multiplied with the corresponding pair of loaded stationary elements. ! + > **Figure: nki trn2 arch 4** > > A two-part diagram showing Tensor Engine operations: (a) LoadStationary instruction loading data from SBUF to TensorE, and (b) MultiplyMoving instruction showing the full matmul operation with output to PSUM. @@ -245,6 +272,7 @@ Next, we invoke the LoadStationary and MultiplyMoving instructions to perform th > This diagram illustrates the two key Tensor Engine instructions for matrix multiplication. > > **Part (a) - LoadStationary Instruction** (top): +> > - Shows data movement from SBUF to Tensor Engine > - Left side: "stationary (TensorE)" green block with dimensions K/2=128 (height) > - Right side: "stationary (SBUF)" green block with dimensions K/2=128 (height) and M=128 twice (width, divided by dashed line) @@ -253,6 +281,7 @@ Next, we invoke the LoadStationary and MultiplyMoving instructions to perform th > - "SBUF P-dim" label on right > > **Part (b) - MultiplyMoving Instruction** (bottom): +> > - Shows the multiplication operation with moving matrix > - Left side: "stationary (TensorE)" green block, K/2=128 height > - Center-right: "moving (SBUF)" blue block with K/2=128 height and N=512 twice width @@ -263,12 +292,14 @@ Next, we invoke the LoadStationary and MultiplyMoving instructions to perform th > - "PSUM P-dim" label below output > > **Dimension annotations**: +> > - K/2=128: Partition dimension for both operations > - M=128: Free dimension tiles for stationary > - N=512: Free dimension tiles for moving > - 512, 128: Output dimensions in PSUM > > **Key Elements:** +> > - **LoadStationary Instruction (a)**: Load stationary matrix from SBUF to Tensor Engine > - **MultiplyMoving Instruction (b)**: Multiply with moving matrix, store in PSUM > - **stationary (TensorE)**: Matrix held in Tensor Engine (green) @@ -282,7 +313,6 @@ Note that the above double FP8 `LoadStationary`/`MultiplyMoving` instruction seq NKI programmers can invoke double FP8 matmul using the `nisa.nc_matmul()` API on NeuronCore-v3: - ```python import nki.isa as nisa @@ -293,18 +323,17 @@ nisa.nc_matmul(dst, stationary, moving, perf_mode=nisa.matmul_perf_mode.double_row, ...) ``` - The `nt.tensor[128, 2, 128]` stationary and `nt.tensor[128, 2, 512]` moving tensor shapes reflect the maximum tile sizes for the double FP8 matmul instruction. Smaller tile sizes are supported, though the second dimension (the most major free dimension) of both input tensors must be two. In other words, if the contraction dimension of the matmul is not a multiple of two, programmers are required to explicitly pad the input tensors with zeros to enable the performance mode. A full NKI kernel example performing double FP8 matmul is available on [Github](https://github.com/aws-neuron/nki-samples/blob/main/src/nki_samples/reference/double_row_matmul.py). Note that Double FP8 matmul performance mode cannot be combined with the following TensorE features: -* Column tiling mode +- Column tiling mode -* Sparse matmul (new in NeuronCore-v3, discussion below) +- Sparse matmul (new in NeuronCore-v3, discussion below) -* Transpose mode (new in NeuronCore-v3, more discussion below) +- Transpose mode (new in NeuronCore-v3, more discussion below) ### Built-in Transpose Support @@ -312,27 +341,24 @@ As discussed in [Trainium/Inferentia2 Architecture Guide](trainium_inferentia2_a Starting with NeuronCore-v3, TensorE supports an explicit transpose mode, which can correctly transpose input tensors with NaN/Inf. In addition, the transpose mode provides the following benefits: -* 2x speedup in FP32 transpose, vs. no transpose mode enabled. +- 2x speedup in FP32 transpose, vs. no transpose mode enabled. -* FP16/BF16 PSUM output for FP16/BF16 transpose, vs. FP32 (default matmul output data type) PSUM output when no transpose mode enabled. This allows faster PSUM data eviction back to SBUF. +- FP16/BF16 PSUM output for FP16/BF16 transpose, vs. FP32 (default matmul output data type) PSUM output when no transpose mode enabled. This allows faster PSUM data eviction back to SBUF. > **Note** > > Note -> -> +> > NeuronCore-v3 TensorE transpose mode for FP8 input data produces 16-bit output elements in PSUM, with the upper 8 bits filled with zeros. NKI programmers can enable TensorE transpose mode on NeuronCore-v3 through the following APIs: - ```python nisa.nc_matmul(..., is_transpose=True) # OR nisa.nc_transpose(..., engine=nisa.constants.engine.tensor) ``` - ## Vector Engine Vector Engine (VectorE) is specially designed to accelerate vector operations where every element in the output tensor typically depends on multiple elements from input tensor(s), such as vector reduction and element-wise operators between two tensors. NeuronCore-v3 Vector Engine delivers a total of 1.0 TFLOPS of FP32 computations and can handle various input/output data-types, including FP8, FP16, BF16, TF32, FP32, INT8, INT16, and INT32. @@ -343,26 +369,23 @@ NeuronCore-v3 Vector Engine provides a new performance mode BF16/FP16 data types In particular, the following instructions could see a 4x throughput lift compared to NeuronCore-v2: -* -`nisa.tensor_copy` and `nisa.tensor_scalar` when both input/output tensors: +- `nisa.tensor_copy` and `nisa.tensor_scalar` when both input/output tensors: are in SBUF -* are in BF16/FP16 (input and output data types do not need to match) +- are in BF16/FP16 (input and output data types do not need to match) -* have physically contiguous elements in the inner-most (most minor) free dimension +- have physically contiguous elements in the inner-most (most minor) free dimension The following instructions could see a 2x throughput lift compared to NeuronCore-v2: -* -`nisa.tensor_copy` and `nisa.tensor_scalar`: +- `nisa.tensor_copy` and `nisa.tensor_scalar`: when both input/output tensors satisfy 1a and 1b, but not 1c conditions above, or -* when both input/output tensors satisfy 1b and 1c, but one of input and output tensors is in PSUM +- when both input/output tensors satisfy 1b and 1c, but one of input and output tensors is in PSUM -* -`nisa.tensor_tensor`: +- `nisa.tensor_tensor`: when both input tensors are SBUF and all of input/output tensors are in BF16/FP16 @@ -385,6 +408,7 @@ In NeuronCore-v3, each processor in GpsimdE also comes with an integrated DMA en Trainium2 consists of a three-tiered memory hierarchy: HBM, SBUF and PSUM, from highest to lowest memory capacity. Figures below show the specifications of these memories and their connectivity for one NeuronCore-v3. ! + > **Figure: nki trn2 arch 5 1** > > A memory hierarchy pyramid diagram showing four levels from Host CPU memory at the bottom to PSUM at the top, with capacity and bandwidth specifications for each level and data flow operations labeled. @@ -394,37 +418,44 @@ Trainium2 consists of a three-tiered memory hierarchy: HBM, SBUF and PSUM, from > **Pyramid levels (top to bottom)**: > > **Level 1 - PSUM (top, yellow)**: +> > - Capacity: ~2 MB > - Bandwidth: ~10 TB/sec > - Blue arrow up: "MatMult" (data flows up from SBUF for matrix multiplication) > - Red arrow down: "Use MatMult result" (results flow back to SBUF) > > **Level 2 - SBUF (yellow/green)**: +> > - Capacity: ~25 MB > - Bandwidth: ~10 TB/sec > - Central position in the on-chip hierarchy > > **Level 3 - Device memory (HBM) (green)**: +> > - Capacity: ~50 GB > - Bandwidth: ~0.5 TB/sec per NC (NeuronCore) > - Blue arrow up: "Refill, or Start NKI kernel" > - Red arrow down: "Spill, or End NKI kernel" > > **Level 4 - Host (CPU) memory (DRAM) (blue/gray)**: +> > - Capacity: ~1 TB > - Bandwidth: ~16 GB/sec > - Blue arrow up: "Start compute graph" > - Red arrow down: "End compute graph" > > **Right side annotations**: +> > - Bracket labeled "Memory within NeuronCore (on-chip)" encompasses PSUM and SBUF > - Bracket labeled "Memory within Neuron Device" encompasses PSUM, SBUF, and HBM > > **Color coding for arrows**: +> > - Blue arrows: Data moving up the hierarchy (toward compute) > - Red arrows: Data moving down the hierarchy (results/spill) > > **Key Elements:** +> > - **PSUM**: ~2 MB, ~10 TB/sec - fastest, smallest (top) > - **SBUF**: ~25 MB, ~10 TB/sec - main on-chip buffer > - **Device memory (HBM)**: ~50 GB, ~0.5 TB/sec per NC @@ -436,6 +467,7 @@ Trainium2 consists of a three-tiered memory hierarchy: HBM, SBUF and PSUM, from > - **On-chip vs Device memory**: Hierarchy classification ! + > **Figure: nki trn2 arch 6** > > A NeuronCore memory hierarchy diagram showing the relationship between on-chip components (PSUM, compute engines, SBUF) and off-chip HBM, with DMA engines facilitating data transfer. @@ -445,11 +477,13 @@ Trainium2 consists of a three-tiered memory hierarchy: HBM, SBUF and PSUM, from > **On-chip section** (enclosed in dashed rectangle on right): > > **PSUM (top)** - peach/orange colored block: +> > - Spans full width > - Partial Sum accumulator for Tensor Engine outputs > - Bidirectional arrows connect to compute engines below > > **Compute Engines** - four blocks in a row: +> > - **TensorE**: Tensor Engine (leftmost) > - **VectorE**: Vector Engine > - **ScalarE**: Scalar Engine @@ -457,11 +491,13 @@ Trainium2 consists of a three-tiered memory hierarchy: HBM, SBUF and PSUM, from > - Each has bidirectional arrows to PSUM above and SBUF below > > **SBUF (middle)** - green colored block: +> > - State Buffer - main on-chip SRAM > - Spans full width > - Central hub connecting compute engines to external memory > > **DMA Engines** - multiple blocks below SBUF: +> > - Four "DMA" blocks shown with "..." indicating more > - Bidirectional arrows connect to SBUF above and HBM below > - Facilitate data movement between on-chip and off-chip memory @@ -469,16 +505,19 @@ Trainium2 consists of a three-tiered memory hierarchy: HBM, SBUF and PSUM, from > **Off-chip section**: > > **HBM (bottom)** - light blue colored block: +> > - High Bandwidth Memory > - External device memory > - Connected to DMA engines above > > **Labels**: +> > - "on-chip" bracket on right encompassing PSUM through DMA > - "off-chip" bracket on right for HBM > - Dashed line separates on-chip from off-chip > > **Key Elements:** +> > - **PSUM**: Partial Sum buffer (peach/orange) > - **TensorE, VectorE, ScalarE, GpSimdE**: Four compute engines > - **SBUF**: State Buffer - main on-chip SRAM (green) @@ -490,25 +529,24 @@ Trainium2 consists of a three-tiered memory hierarchy: HBM, SBUF and PSUM, from As shown in the above figures, data movement between HBM and SBUF is performed using on-chip DMA (Direct Memory Access) engines, which can run in parallel to computation within the NeuronCore. Data movement between PSUM and SBUF is done through ISA instructions on the compute engines. In NeuronCore-v3, two restrictions in engine parallel accesses to SBUF/PSUM are lifted to improve programming flexibility compared to NeuronCore-v2: -* -VectorE and GpSimdE can access SBUF in parallel. +- VectorE and GpSimdE can access SBUF in parallel. This was disallowed in NeuronCore-v2. -* VectorE’s performance mode leverages a shared memory bus between the VectorE and GpsimdE engines to deliver 2-4x performance improvement for select VectorE instructions. The hardware automatically coordinates access between engines to optimize bus utilization, including arbitrating between GpsimdE and relevant VectorE instructions. +- VectorE’s performance mode leverages a shared memory bus between the VectorE and GpsimdE engines to deliver 2-4x performance improvement for select VectorE instructions. The hardware automatically coordinates access between engines to optimize bus utilization, including arbitrating between GpsimdE and relevant VectorE instructions. -* -VectorE and ScalarE can access PSUM in parallel. +- VectorE and ScalarE can access PSUM in parallel. This was disallowed in NeuronCore-v2. -* Both VectorE and ScalarE can access PSUM at full bandwidth in parallel, as long as their accesses do not collide on the same PSUM bank. +- Both VectorE and ScalarE can access PSUM at full bandwidth in parallel, as long as their accesses do not collide on the same PSUM bank. ### DMA Transpose Trainium2 DMA engines can perform a tensor transpose while moving data from HBM into SBUF, or from SBUF to SBUF itself. The figure below illustrates these two supported DMA transpose data flows. Trainium2 DMA transpose supports bit-accurate transposition for both 2-byte and 4-byte data types. ! + > **Figure: nki trn2 arch 7** > > A diagram showing the DMA transpose mechanism where data flows from HBM through multiple DMA engines and an Xpose Block to produce both transposed (data.T) and non-transposed (data) outputs in SBUF. @@ -516,17 +554,20 @@ Trainium2 DMA engines can perform a tensor transpose while moving data from HBM > This diagram illustrates the hardware transpose capability during DMA transfers from HBM to SBUF. > > **Left side - HBM (gray)**: +> > - Large gray block representing High Bandwidth Memory > - Contains "data" label indicating source data > - Arrow points right showing data flow > > **Center - DMA and Xpose Block**: +> > - Multiple "DMA" blocks stacked vertically (4 shown with "..." indicating more) > - Arrows flow from HBM through DMA blocks > - All DMA outputs feed into a central "Xpose Block" (purple) > - Numbers "(1)" and "(2)" in purple indicate two output paths from the Xpose Block > > **Right side - SBUF (green)**: +> > - Large green block representing State Buffer > - Two outputs from Xpose Block: > - "data.T" (transposed data) - upper path @@ -536,6 +577,7 @@ Trainium2 DMA engines can perform a tensor transpose while moving data from HBM > The Xpose Block is the hardware unit that performs on-the-fly transpose during DMA operations, allowing data to be written to SBUF in either transposed or non-transposed format without additional compute operations. > > **Key Elements:** +> > - **HBM**: Source High Bandwidth Memory (gray) > - **data**: Source data in HBM > - **DMA**: Multiple DMA engines (stacked blocks) @@ -546,16 +588,16 @@ Trainium2 DMA engines can perform a tensor transpose while moving data from HBM > - **SBUF**: Destination State Buffer (green) > - **Ellipsis (...)**: Indicates additional DMA engines - #### HBM2SBUF DMA transpose Before diving into how HBM2SBUF transpose works, let’s revisit a simple DMA copy from a packed HBM tensor `[128, 512]` to an SBUF tensor `[nl.par_dim(128), 512]`. Following Numpy convention, these tensor shapes follow a major to minor ordering. The figure below visualizes these HBM and SBUF tensors. A packed `[128, 512]` HBM tensor consists of 128 chunks of 512 elements, laid out back to back in the HBM linear memory. The most minor (that is, inner-most) dimension consists of 512 contiguous elements in memory. Once loaded into the SBUF, the most minor HBM tensor dimension (512) is mapped to the free dimension of the SBUF, while the most major dimension is mapped to the SBUF partition dimension. In Trainium2, each NeuronCore-v3 is typically paired with 16x DMA engines to drive its corresponding SBUF bandwidth. In the above DMA copy, each DMA engine would be responsible for moving 128/16 = 8 chunks of 512 elements. -* HBM tensor [128, 512]: 512 is the inner-most (minor) dimension +- HBM tensor [128, 512]: 512 is the inner-most (minor) dimension ! + > **Figure: nki trn2 arch 8** > > A diagram showing DMA copy operation from HBM tensor [128, 512] to SBUF tensor with partition dimension, where contiguous 512-element rows become columns in the SBUF 2D layout. @@ -563,6 +605,7 @@ In Trainium2, each NeuronCore-v3 is typically paired with 16x DMA engines to dri > This diagram illustrates how a DMA copy operation maps a linear HBM tensor to a 2D SBUF tensor layout. > > **Left side - Source HBM tensor**: +> > - Horizontal strip showing "src: HBM tensor [128, 512]" > - Three colored blocks (blue, purple, green) each representing 512 elements > - Dimensions labeled as "512", "512", "512" above each block @@ -570,10 +613,12 @@ In Trainium2, each NeuronCore-v3 is typically paired with 16x DMA engines to dri > - Ellipsis (...) indicates additional rows > > **Center - Arrow**: +> > - "DMA copy" label with arrow pointing right > - Indicates the data movement operation > > **Right side - Destination SBUF tensor**: +> > - Large rectangular block showing "dst: SBUF tensor [nl.par_dim(128), 512]" > - Dimensions: "512 F" (free dimension) width, "128 P" (partition dimension) height > - Same colored blocks (blue, purple, green) now arranged vertically as rows @@ -582,11 +627,13 @@ In Trainium2, each NeuronCore-v3 is typically paired with 16x DMA engines to dri > - Ellipsis (...) indicates additional partitions > > The key transformation: +> > - HBM tensor is logically [128, 512] (128 rows of 512 elements each) > - In SBUF, these become 128 partitions (P-dim) with 512 elements (F-dim) each > - The nl.par_dim(128) indicates the partition dimension designation > > **Key Elements:** +> > - **src: HBM tensor [128, 512]**: Source tensor with 128 rows of 512 elements > - **DMA copy**: Data movement operation > - **dst: SBUF tensor [nl.par_dim(128), 512]**: Destination with partition dimension @@ -598,6 +645,7 @@ In Trainium2, each NeuronCore-v3 is typically paired with 16x DMA engines to dri In contrast, in a DMA transpose operation, we take an HBM tensor of opposite layout [512, 128]: ! + > **Figure: nki trn2 arch 9** > > A diagram showing DMA transpose operation from HBM tensor [512, 128] to SBUF tensor, where the data is transposed during transfer so rows become distributed across the free dimension. @@ -605,6 +653,7 @@ In contrast, in a DMA transpose operation, we take an HBM tensor of opposite lay > This diagram illustrates how a DMA transpose operation maps and transposes a linear HBM tensor to a 2D SBUF tensor layout. > > **Left side - Source HBM tensor**: +> > - Horizontal strip showing "src: HBM tensor [512, 128]" > - Three colored blocks (blue, purple, green) each representing 128 elements > - Dimensions labeled as "128", "128", "128" above each block @@ -612,10 +661,12 @@ In contrast, in a DMA transpose operation, we take an HBM tensor of opposite lay > - Ellipsis (...) indicates additional rows > > **Center - Arrow**: +> > - "DMA transpose" label with arrow pointing right > - Indicates the data movement with transpose operation > > **Right side - Destination SBUF tensor**: +> > - Large rectangular block showing "dst: SBUF tensor [nl.par_dim(128), 512]" > - Dimensions: "512 F" (free dimension) width, "128 P" (partition dimension) height > - Same colored blocks (blue, purple, green) now arranged as vertical columns within the tensor @@ -624,12 +675,14 @@ In contrast, in a DMA transpose operation, we take an HBM tensor of opposite lay > - Ellipsis (...) indicates additional columns > > The key transformation: +> > - HBM tensor is [512, 128] (512 rows of 128 elements) > - After transpose, becomes [128, 512] in SBUF > - Original 128-element rows become columns in the free dimension > - The 512 original rows become distributed across the free dimension > > **Key Elements:** +> > - **src: HBM tensor [512, 128]**: Source tensor (512 x 128) > - **DMA transpose**: Data movement with transpose operation > - **dst: SBUF tensor [nl.par_dim(128), 512]**: Transposed destination @@ -644,7 +697,6 @@ HBM2SBUF DMA transpose is commonly seen in ML workloads where the data layout in In NKI, programmers can invoke an HBM2SBUF DMA transpose using the `nisa.dma_transpose` API. - ```python import nki import nki.language as nl @@ -655,12 +707,10 @@ import nki.isa as nisa sbuf_dst = nisa.dma_transpose(src=hbm_src) ``` - > **Note** > > Performance Consideration -> -> +> > DMA transpose on Trainium2 can achieve up to 90% DMA throughput utilization given hardware-friendly tensor access patterns, compared to up to 100% throughput utilization for a DMA copy. #### SBUF2SBUF DMA transpose @@ -669,7 +719,6 @@ SBUF2SBUF DMA transpose works in a similar fashion as HBM2SBUF transpose, where The same `nisa.dma_transpose` API can be used to perform an SBUF2SBUF DMA transpose: - ```python import nki import nki.language as nl @@ -680,7 +729,6 @@ import nki.isa as nisa sbuf_dst = nisa.dma_transpose(src=hbm_src) ``` - Performance Consideration. SBUF2SBUF transpose can achieve up to 50% of DMA throughput on Trainium2. Compared to TensorE transpose that is more performant but requires ScalarE/VectorE to evict the transposed output from PSUM back to SBUF, DMA transpose can read from and write to SBUF directly. Therefore, DMA transpose is particularly useful in operators that are ScalarE/VectorE bound, such as self attention. ### Descriptor Generation Engine (DGE) @@ -688,6 +736,7 @@ Performance Consideration. SBUF2SBUF transpose can achieve up to 50% of DMA thro The Descriptor Generation Engine (DGE) is a new hardware block in NeuronCore-v3 that accelerates DMA descriptor generation to perform either DMA copy or transpose on the DMA engines. Each NeuronCore-v3 comes with two instances of DGE, which can be commanded through either SyncE or ScalarE sequencer. The figure below shows the connectivity of the DGE instances. ! + > **Figure: nki trn2 arch 10** > > A diagram showing DMA engines and DGE (DMA Gather Engine) components interfacing with NeuronCore's Scalar Engine and Sync Engine, illustrating the descriptor-based DMA command architecture. @@ -695,33 +744,39 @@ The Descriptor Generation Engine (DGE) is a new hardware block in NeuronCore-v3 > This diagram illustrates the DMA subsystem architecture and how it interfaces with the NeuronCore for data movement operations. > > **Left side - DMA engines**: +> > - Four "DMA" blocks shown as stacked gray rectangles (with "..." indicating more) > - These represent the pool of DMA engines available for data transfers > > **Center - DGE (DMA Gather Engines)**: +> > - "DGE[0]" (purple block) at top > - "DGE[1]" (blue block) at bottom > - Each DGE connects to multiple DMA engines via "desc" (descriptor) arrows > - The DGEs gather and dispatch DMA operations > > **Arrows and connections**: +> > - "desc" arrows: From DMA engines to both DGE[0] and DGE[1], showing descriptor-based control > - Lines cross between DMA engines and DGEs, indicating flexible mapping > - "cmd" arrows: From DGE[0] and DGE[1] to NeuronCore components > > **Right side - NeuronCore**: +> > - Rounded rectangle labeled "NeuronCore" > - Contains two components: > - "SEQ" block (gray) with "Scalar Engine" (green) - receives "cmd" from DGE[0] > - "Sync Engine" (green) - receives "cmd" from DGE[1] > > The diagram shows that: +> > 1. DMA engines are controlled via descriptors > 2. DGEs aggregate DMA commands > 3. Scalar Engine controls DGE[0] for general data movement > 4. Sync Engine controls DGE[1] for synchronized transfers > > **Key Elements:** +> > - **DMA**: Multiple DMA engines for data movement (gray blocks) > - **DGE[0]**: DMA Gather Engine 0 (purple) - controlled by Scalar Engine > - **DGE[1]**: DMA Gather Engine 1 (blue) - controlled by Sync Engine @@ -741,13 +796,11 @@ NKI programmers can invoke hardware-based DGE on NeuronCore-v3 using `nisa.dma_c > **Note** > > Note -> -> +> > NeuronCore-v3 hardware DGE currently does not support indirect DMA operations (gather/scatter). Refer to nisa API documentation for detailed implementation guidelines. > **Note** > > Performance Consideration -> -> -> When triggered from ScalarE, execution of the DGE-based DMA instruction could be hidden behind earlier compute instructions (such as `nisa.activate()`) in program order, since DGE and the compute pipeline of ScalarE are independent hardware resources. Each DGE-based DMA instruction takes about 600 ns to execute on NeuronCore-v3. \ No newline at end of file +> +> When triggered from ScalarE, execution of the DGE-based DMA instruction could be hidden behind earlier compute instructions (such as `nisa.activate()`) in program order, since DGE and the compute pipeline of ScalarE are independent hardware resources. Each DGE-based DMA instruction takes about 600 ns to execute on NeuronCore-v3. diff --git a/skills/neuron-nki-docs/references/architecture/trainium3_arch.md b/skills/neuron-nki-docs/references/architecture/trainium3_arch.md index 08c1841..1f2f2f9 100644 --- a/skills/neuron-nki-docs/references/architecture/trainium3_arch.md +++ b/skills/neuron-nki-docs/references/architecture/trainium3_arch.md @@ -5,25 +5,25 @@ Trainium3 Architecture Guide for NKI > **Note** > > Note -> -> +> > If nisa API is mentioned for a given architectural feature, that means NKI support is ready yet. In this guide, we will dive into hardware architecture of fourth-generation NeuronDevices: Trainium3. This guide will highlight major architectural updates compared to the previous generation (Trainium2). Therefore, we assume readers are familiar with [Trainium/Inferentia2 Architecture Guide](trainium_inferentia2_arch.md) and [Trainium2 Architecture Guide for NKI](trainium2_arch.md) to understand the basics of NeuronDevice Architecture. The diagram below shows a block diagram of a Trainium3 device, which consists of: -* 8 NeuronCores (v4). +- 8 NeuronCores (v4). -* 4 HBM stacks with a total device memory capacity of 144 GiB and bandwidth of 4.7 TB/s. +- 4 HBM stacks with a total device memory capacity of 144 GiB and bandwidth of 4.7 TB/s. -* 128 DMA (Direct Memory Access) engines to move data within and across devices. +- 128 DMA (Direct Memory Access) engines to move data within and across devices. -* 20 CC-Cores for collective communication. +- 20 CC-Cores for collective communication. -* 4 NeuronLink-v4 for device-to-device collective communication. +- 4 NeuronLink-v4 for device-to-device collective communication. ! + > **Figure: nki trn3 arch 1** > > An architecture diagram of AWS Trainium3 showing 8 NeuronCore-v4 units arranged in a 2x4 grid, each containing On-chip SRAM, Tensor Engine, Vector Engine, Scalar Engine, and GPSIMD Engine, with HBM, DMA, CC-Core, and NeuronLink-v4 interconnects. @@ -33,6 +33,7 @@ The diagram below shows a block diagram of a Trainium3 device, which consists of > **Title**: "Trainium3" in blue text at top left > > **NeuronCore arrangement**: +> > - 8 "NeuronCore-v4" units arranged in a 2-row by 4-column grid > - Each NeuronCore-v4 contains: > - **On-chip SRAM memory**: Database/cylinder icon representing local memory @@ -42,21 +43,25 @@ The diagram below shows a block diagram of a Trainium3 device, which consists of > - **GPSIMD Engine**: Multiple small grid icons for general-purpose SIMD > > **Memory (HBM)**: +> > - Two "HBM" blocks on the left side (serving top and bottom rows) > - Two "HBM" blocks on the right side (serving top and bottom rows) > - High bandwidth memory provides external storage > > **Support components** (bottom area): +> > - "DMA": DMA engines block (stacked appearance indicating multiple) > - "CC-Core": Collective Communication cores (stacked) > - "Host PCIe": Host interface on the right > > **Interconnects**: +> > - Four "NeuronLink-v4" blocks at the bottom for inter-chip communication > > The Trainium3 represents a significant evolution with NeuronCore-v4 units and NeuronLink-v4 interconnects, maintaining the proven architecture pattern while scaling compute capabilities. > > **Key Elements:** +> > - **Trainium3**: Third-generation training chip > - **NeuronCore-v4**: 8 next-generation compute cores (2x4 grid) > - **On-chip SRAM memory**: Local storage in each core @@ -79,13 +84,12 @@ The figure below is a simplified NeuronCore-v4 diagram of the compute engines an ![../../../_images/nki-trn3-arch-2.png](../../../_images/nki-trn3-arch-2.png) The NeuronCore-v4 SBUF capacity is 32MiB (up from 28 MiB in NeuronCore-v3), while the PSUM capacity remains the same at 2MiB. The engine data-path widths and frequencies are updated to the following: - -| Device Architecture | Compute Engine | Data-path Width (elements/cycle) | Frequency (GHz) | -| --- | --- | --- | --- | -| Trainium3 | Tensor | 8x128 (MXFP8 dense input) or 2x128 (non-MXFP8 dense input) or 5x128 (sparse input); 1x128 (output) | 2.4 | -| | Vector | 512 BF16/FP16/FP8 input/output; 256 input/output for other data types | 1.2 | -| | Scalar | 256 BF16/FP16/FP8 input/output; 128 input/output for other data types | 1.2 | -| | GpSimd | 128 input/output for all data types | 1.2 | +| Device Architecture | Compute Engine | Data-path Width (elements/cycle) | Frequency (GHz) | +| ------------------- | -------------- | -------------------------------------------------------------------------------------------------- | --------------- | +| Trainium3 | Tensor | 8x128 (MXFP8 dense input) or 2x128 (non-MXFP8 dense input) or 5x128 (sparse input); 1x128 (output) | 2.4 | +| | Vector | 512 BF16/FP16/FP8 input/output; 256 input/output for other data types | 1.2 | +| | Scalar | 256 BF16/FP16/FP8 input/output; 128 input/output for other data types | 1.2 | +| | GpSimd | 128 input/output for all data types | 1.2 | Sync Engine has not changed since [previous Trainium architectures](trainium_inferentia2_arch.md). Next, we will go over major architectural updates to each compute engine. @@ -121,29 +125,26 @@ A single scaling group corresponds to one 8-bit integer scale. Therefore, for ev ![../../../_images/nki-trn3-arch-7.png](../../../_images/nki-trn3-arch-7.png) The moving data and scale tensor layout follows the same rules. Therefore, an MX matmul on TensorE requires four input tensors: -* stationary data +- stationary data -* stationary scale +- stationary scale -* moving data +- moving data -* moving scale +- moving scale In NKI, programmers can define MX data tensors using the special x4 data types. The maximum tile size for stationary MX data tensor is [128, 128] in x4 data types ([128, 512] of actual values), while the maximum tile size for moving MX data tensor is [128, 512] in x4 data types ([128, 2048] of actual values). One convenience of the x4 datatypes is that the output matrix dimensions map directly to the sizes of the free dimensions of the input matrices. Similarly, the maximum tile size for stationary and moving MX scale tensors are [128, 128] and [128, 512] in nl.uint8, respectively. The API to invoke an MX matmul is: - ```python nisa.nc_matmul_mx(moving, stationary, moving_scale, stationary_scale) ``` - ### BF16 Matmul Results in PSUM Prior to the NeuronCore-v4, the Tensor Engine always passes FP32 matrix multiplication results to PSUM unless transpose mode is turned on. Similarly, the PSUM buffer was restricted to FP32 near-memory accumulation (fp32_psum_tensor += fp32_matmul_output). Starting with the NeuronCore-v4, the Tensor Engine allows the matrix multiplication instruction (nisa.nc_matmul) to store BF16 data into the PSUM buffer directly and also to perform addition to a BF16 tensor stored in PSUM. NKI programmers can use this feature through the existing nisa.nc_matmul API: - ```python psum_tensor = nl.ndarray((128, 512), dtype=nl.bfloat16, buffer=nl.psum) @@ -151,17 +152,16 @@ nisa.nc_matmul(..., dst=psum_tensor, psum_accumulate_flags=1) nisa.nc_matmul(..., dst=psum_tensor, psum_accumulate_flags=0) ``` - Note that the accumulation performed during a matmul operation within the systolic array is still performed using FP32 data. When writing the matmul results into a BF16 PSUM tensor location, the downcast from FP32 to BF16 is performed immediately before the write. The downcast can use the RNE (round nearest even) or SR (stochastic rounding) mode. The figure below illustrates this data flow. ![../../../_images/nki-trn3-arch-8.png](../../../_images/nki-trn3-arch-8.png) When adding the matmul results to an existing BF16 tensor stored in PSUM the following operations are performed: -* The existing PSUM tensor (red) is upcast to FP32. +- The existing PSUM tensor (red) is upcast to FP32. -* The PSUM tensor (now in FP32) and the TensorE output (yellow) are added together at FP32 precision. +- The PSUM tensor (now in FP32) and the TensorE output (yellow) are added together at FP32 precision. -* The result of the addition (green) is converted to BF16 using the given rounding mode, and written back to PSUM. +- The result of the addition (green) is converted to BF16 using the given rounding mode, and written back to PSUM. ![../../../_images/nki-trn3-arch-9.png](../../../_images/nki-trn3-arch-9.png) @@ -180,6 +180,7 @@ The Vector Engine is optimized for vector computations, in which every element o The NeuronCore-v4 VectorE supports quantizing FP16/BF16 data to MXFP8 tensors (both data and scales) in a layout that TensorE can directly consume for MX matmul, as described in the Quad-MXFP8/MXFP4 Matmul Performance section above. As a reminder, an MxK MXFP8 matrix, where K is the contraction dimension, requires the following data and scale layout in SBUF: ! + > **Figure: nki trn3 arch 10** > > A diagram showing the SBUF layout for MX (microscaling) format data, with the main data tensor having 8P block groups and a separate scale tensor with 4P blocks, both occupying 32P total partition height. @@ -187,8 +188,9 @@ The NeuronCore-v4 VectorE supports quantizing FP16/BF16 data to MXFP8 tensors (b > This diagram illustrates how MX (microscaling) format data is laid out in the State Buffer (SBUF), showing the relationship between data and scale tensors. > > **Left tensor - "data" (green)**: +> > - Large green rectangular block -> - Width: "M * 4" (free dimension) +> - Width: "M \* 4" (free dimension) > - Height: "K/4 = 128" (annotation on left), total "32P" partitions > - Contains a highlighted block group in upper left with colored markers (red, blue, yellow, green) > - The block group is "8P" tall (8 partitions) @@ -196,6 +198,7 @@ The NeuronCore-v4 VectorE supports quantizing FP16/BF16 data to MXFP8 tensors (b > - Label "data" (italic) above > > **Right tensor - "scale" (green with gray stripes)**: +> > - Narrow vertical tensor > - Width: "M" (free dimension) > - Height: "32P" partitions total @@ -207,18 +210,20 @@ The NeuronCore-v4 VectorE supports quantizing FP16/BF16 data to MXFP8 tensors (b > **Caption**: "Data and Scale Layout in SBUF" centered below > > The MX format stores: +> > - Main data in larger blocks (8P partition groups) > - Corresponding scale factors in smaller blocks (4P partition groups) > - Scale factor count is half of data partition count due to microscaling sharing > > **Key Elements:** +> > - **data tensor**: Main MX data [K/4=128 partitions x M*4 free elements] > - **scale tensor**: Scale factors [32P x M] > - **8P**: Data block group size in partitions > - **4P**: Scale block size in partitions > - **32P**: Total partition dimension height > - **K/4 = 128**: Partition dimension size -> - **M * 4, M**: Free dimension sizes +> - **M \* 4, M**: Free dimension sizes > - **Colored markers**: Block group boundary indicators > - **Dashed lines**: Block group boundaries > - **Green stripes**: Scale factor locations @@ -226,6 +231,7 @@ The NeuronCore-v4 VectorE supports quantizing FP16/BF16 data to MXFP8 tensors (b The VectorE can natively quantize BF16/FP16 data to produce this layout using the QuantizeMX instruction. QuantizeMX calculates the required scales for each group of 32 values, divides them by the calculated scale, and casts to the target MXFP8 datatype (as per the OCP specification): ! + > **Figure: nki trn3 arch 11** > > A diagram showing the QuantizeMX() operation on VectorE, converting BF16/FP16 data to MXFP8 format, producing both quantized data and scale tensors. @@ -233,8 +239,9 @@ The VectorE can natively quantize BF16/FP16 data to produce this layout using th > This diagram illustrates the MX quantization operation that converts higher-precision floating-point data to MXFP8 format using the Vector Engine. > > **Left side - Input "BF16/FP16 data" (blue)**: +> > - Large blue rectangular tensor -> - Width: "M * 4" (free dimension) +> - Width: "M \* 4" (free dimension) > - Height: "K/4 = 128", with "32P" total partitions > - Contains block group indicator with colored markers (red, blue, yellow, green) > - Block group height: "8P" (8 partitions) @@ -242,6 +249,7 @@ The VectorE can natively quantize BF16/FP16 data to produce this layout using th > - Label "BF16/FP16 data" (italic) above > > **Center - Operation**: +> > - Arrow pointing right > - "QuantizeMX()" label above > - "VectorE" label in a box below the arrow @@ -250,13 +258,15 @@ The VectorE can natively quantize BF16/FP16 data to produce this layout using th > **Right side - Outputs**: > > **"MXFP8 data" (green)**: +> > - Green rectangular tensor with same dimensions as input -> - Width: "M * 4" +> - Width: "M \* 4" > - Height: "K/4 = 128", "32P" partitions > - Same block structure with "8P" groups > - Label "MXFP8 data" (italic) above > > **"MXFP8 scale" (green with gray)**: +> > - Narrow tensor to the right > - Width: "M" > - Height: "32P" @@ -265,6 +275,7 @@ The VectorE can natively quantize BF16/FP16 data to produce this layout using th > - Label "MXFP8 scale" (italic) above > > **Key Elements:** +> > - **BF16/FP16 data**: Input tensor in 16-bit format (blue) > - **QuantizeMX()**: Quantization operation > - **VectorE**: Vector Engine performs the conversion @@ -273,7 +284,7 @@ The VectorE can natively quantize BF16/FP16 data to produce this layout using th > - **8P, 4P**: Block group sizes > - **32P**: Total partition height > - **K/4=128**: Partition dimension -> - **M * 4, M**: Free dimension sizes +> - **M \* 4, M**: Free dimension sizes The source FP16/BF16 data must be in SBUF, and has to be in a layout that exactly matches the target MXFP8 data layout (QuantizeMX preserves the data layout). The target MXFP8 data and scales also have to be in SBUF. The quantization instruction can quantize four input elements per partition, per cycle (i.e., 4x Vector performance mode). @@ -283,7 +294,6 @@ In NKI, programmers can perform such an MX data type quantization using the nisa The NeuronCore-v4 Vector Engine introduces a new instruction to perform fast exponential evaluation (nisa.exponential(dst=out_tile, src=in_tile, …)), at 4x the throughput compared to the nisa.activation(op=nl.exp) instruction on the Scalar Engine. In addition to the exponential function, the instruction on Vector Engine can also apply a subtraction before the exponential function and an accumulation after: - ```python # Inputs: # src tile [M, N] @@ -298,7 +308,6 @@ for i in range(M): # parallel (partition) dimension row_max[i, 0] += dst[i, j] ``` - This particular pattern is useful to speed up the Softmax operator, which is commonly on the critical path of long context length self-attention in large language models (LLMs): \[Softmax(X)=\frac{e^{X_i-max(X)}}{\sum_i e^{X_i-max(X)}}\] @@ -329,6 +338,7 @@ Trainium3 introduces the Activation2 instruction, which provides more flexibilit The NeuronCore-v4 SBUF/PSUM introduce a new indirect addressing mode for all compute engines (TensorE/VectorE/ScalarE/GpsimdE), which allows gathering or scattering SBUF and PSUM tensors along the free (F) dimension. Consider a tensor of shape [128, 512] located in SBUF, which occupies 128 partitions with 512 elements per partition. Suppose a user is interested in only accessing the elements 0, 128 and 384 along the free dimension across all 128 partitions for a single computation operation, such as nisa.nc_matmul: ! + > **Figure: nki trn3 arch 12** > > A diagram showing an SBUF tensor layout with dimensions 128 P (partition) by 512 F (free), with colored column stripes at positions 0, 128, and 384 indicating data placement. @@ -336,11 +346,13 @@ The NeuronCore-v4 SBUF/PSUM introduce a new indirect addressing mode for all com > This diagram illustrates a destination SBUF tensor layout showing how data is organized with specific free dimension offsets. > > **Tensor structure**: +> > - Large rectangular block representing an SBUF tensor > - Dimensions: "512 F" (free dimension, horizontal) by "128 P" (partition dimension, vertical) > - Dark gray fill for the main tensor body > > **Colored column stripes**: +> > - Three groups of colored vertical stripes positioned at different free dimension offsets: > - Position 0: Blue and lighter blue stripes on the left edge > - Position 128: Green stripes @@ -348,6 +360,7 @@ The NeuronCore-v4 SBUF/PSUM introduce a new indirect addressing mode for all com > - Each stripe group shows data placement within the free dimension > > **Dimension annotations**: +> > - "512 F" at top indicating free dimension width > - "128 P" on right indicating partition dimension height > - Position markers at bottom: "0", "128", "384" showing free dimension offsets @@ -358,6 +371,7 @@ The NeuronCore-v4 SBUF/PSUM introduce a new indirect addressing mode for all com > The diagram shows how different data chunks (colored stripes) are placed at specific offsets within the free dimension of the SBUF tensor, useful for understanding memory layout and data placement in NKI programming. > > **Key Elements:** +> > - **SBUF tensor [128, 512]**: Destination tensor with 128 partitions, 512 free elements > - **512 F**: Free dimension (horizontal extent) > - **128 P**: Partition dimension (vertical extent) @@ -372,6 +386,7 @@ Since these three vectors do not have a uniform stride along the free dimension; In NeuronCore-v4, all compute engines can perform a gather access pattern to directly access those three vectors in a single instruction: ! + > **Figure: nki trn3 arch 13** > > A diagram showing data flow from SBUF tensor to compute engine, with colored column stripes indicating the data being read and processed along the partition dimension. @@ -379,6 +394,7 @@ In NeuronCore-v4, all compute engines can perform a gather access pattern to dir > This diagram illustrates how data flows from the SBUF tensor to a compute engine for processing. > > **Left side - SBUF tensor**: +> > - Large rectangular block labeled "dst: SBUF tensor [128, 512]" (128 underlined) > - Dimensions: "512 F" (free dimension) width, "128 P" (partition dimension) height > - Dark gray fill @@ -389,10 +405,12 @@ In NeuronCore-v4, all compute engines can perform a gather access pattern to dir > - Position markers at bottom: "0", "128", "384" > > **Center - Data flow**: +> > - Large black arrow pointing right labeled "128 P" > - Indicates data flows along the partition dimension from SBUF to compute > > **Right side - Compute engine**: +> > - Rectangular block labeled "compute engine" > - Same colored vertical stripes (blue, green, purple) showing the data being processed > - The stripes appear in the same relative positions as in SBUF @@ -400,6 +418,7 @@ In NeuronCore-v4, all compute engines can perform a gather access pattern to dir > The diagram shows that compute engines read data from SBUF along the partition dimension, maintaining the same data layout/structure. The 128 P annotation on the arrow indicates all 128 partitions are involved in the data transfer to the compute engine. > > **Key Elements:** +> > - **dst: SBUF tensor [128, 512]**: Source tensor with 128 partitions, 512 free elements > - **512 F**: Free dimension in SBUF > - **128 P**: Partition dimension (both as dimension label and on arrow) @@ -411,6 +430,7 @@ In NeuronCore-v4, all compute engines can perform a gather access pattern to dir Similarly, an indirect scatter operation allows any engine to scatter a set of vectors into a target tensor: ! + > **Figure: nki trn3 arch 14** > > A diagram showing data flow from compute engine back to SBUF tensor, illustrating how computed results are written back to the State Buffer. @@ -418,6 +438,7 @@ Similarly, an indirect scatter operation allows any engine to scatter a set of v > This diagram illustrates the reverse data flow from a compute engine back to the SBUF tensor for storing results. > > **Left side - Compute engine**: +> > - Rectangular block labeled "compute engine" > - Contains colored vertical stripes showing data layout: > - Blue stripes on the left @@ -426,10 +447,12 @@ Similarly, an indirect scatter operation allows any engine to scatter a set of v > - Represents the compute engine holding processed data > > **Center - Data flow**: +> > - Large black arrow pointing right > - Indicates data flows from compute engine back to SBUF > > **Right side - SBUF tensor**: +> > - Large rectangular block labeled "dst: SBUF tensor [128, 512]" (128 underlined) > - Dimensions: "512 F" (free dimension) width, "128 P" (partition dimension) height > - Dark gray fill @@ -443,6 +466,7 @@ Similarly, an indirect scatter operation allows any engine to scatter a set of v > This diagram complements nki-trn3-arch-13.png by showing the write-back path. Together they illustrate the bidirectional data flow between SBUF and compute engines in NeuronCore operations. > > **Key Elements:** +> > - **compute engine**: Source of computed results > - **dst: SBUF tensor [128, 512]**: Destination tensor for results > - **512 F**: Free dimension in SBUF @@ -461,6 +485,7 @@ NeuronCore-v4 introduces an enhanced SBUF capability that enables on-the-fly ten The figure below illustrates the data flow that is used to enable this SBUF accumulation feature. As the first, a DMA unit transfers tensor A to the ReadAddWrite unit adjacent to the SBUF. The ReadAddWrite unit then retrieves tensor B from SBUF, performs the addition of A and B, and writes the result back to tensor B’s original location in SBUF. ! + > **Figure: nki trn3 arch 15** > > A diagram showing the ReadAddWrite DMA operation where multiple DMA engines read new data, add it to existing SBUF data, and write the accumulated result back to SBUF. @@ -468,25 +493,30 @@ The figure below illustrates the data flow that is used to enable this SBUF accu > This diagram illustrates the atomic read-add-write capability of the DMA subsystem for in-place accumulation operations. > > **Left side - DMA engines**: +> > - Four "DMA" blocks shown vertically (gray rectangles) > - Ellipsis (...) indicates additional DMA engines > - Black arrows flow right from each DMA block > > **Center - ReadAddWrite blocks**: +> > - Four "ReadAddWrite" blocks (green) aligned with DMA blocks > - Each receives input from its corresponding DMA engine > - These blocks perform the atomic read-add-write operation > > **Right side - SBUF**: +> > - Large light blue block labeled "SBUF" > - Receives outputs from all ReadAddWrite blocks > > **Arrow legend** (bottom): +> > - **Black solid arrow**: "New data to add (A)" - incoming data from DMA > - **Blue dashed arrow**: "Existing SBUF data (B)" - data read from SBUF > - **Green solid arrow**: "Accumulated data to write (A+B)" - result written to SBUF > > **Data flow**: +> > 1. DMA brings new data (A) from HBM > 2. ReadAddWrite reads existing data (B) from SBUF > 3. ReadAddWrite computes A + B @@ -495,6 +525,7 @@ The figure below illustrates the data flow that is used to enable this SBUF accu > This operation is essential for gradient accumulation and other reduction operations where partial results need to be accumulated in-place without separate read and write operations. > > **Key Elements:** +> > - **DMA**: Multiple DMA engines providing new data (gray blocks) > - **ReadAddWrite**: Atomic read-add-write units (green blocks) > - **SBUF**: State Buffer for accumulated storage (light blue) @@ -507,4 +538,4 @@ Trainium3 DMA engines support Traffic Shaping, which enables configurable bandwi ### DMA QoS -The Trainium3 DMA engines support QoS (quality-of-service), configured per DMA queue through user registers. Note that this implementation of QoS uses a “strict priority”: the transfers in a DMA queue with the highest priority are always scheduled first, before any other DMA queues are serviced. This DMA queue-based QoS feature is particularly useful in the context of parallelizing computation and communication (collectives operation) among multiple NeuronCores. \ No newline at end of file +The Trainium3 DMA engines support QoS (quality-of-service), configured per DMA queue through user registers. Note that this implementation of QoS uses a “strict priority”: the transfers in a DMA queue with the highest priority are always scheduled first, before any other DMA queues are serviced. This DMA queue-based QoS feature is particularly useful in the context of parallelizing computation and communication (collectives operation) among multiple NeuronCores. diff --git a/skills/neuron-nki-docs/references/architecture/trainium_inferentia2_arch.md b/skills/neuron-nki-docs/references/architecture/trainium_inferentia2_arch.md index 6f61e36..ed351c4 100644 --- a/skills/neuron-nki-docs/references/architecture/trainium_inferentia2_arch.md +++ b/skills/neuron-nki-docs/references/architecture/trainium_inferentia2_arch.md @@ -10,16 +10,15 @@ through [NKI Language Guide](../programming/nki-language-guide.md) and familiari [Fig. 47](#fig-arch-neuron-device-v2) shows a block diagram of a Trainium and Inferentia2 device. At a high level, both Trainium and Inferentia2 devices consist of: -* 2 NeuronCores (v2). +- 2 NeuronCores (v2). -* 2 HBM stacks with a total device memory capacity of 32GiB and bandwidth of 820 GB/s. +- 2 HBM stacks with a total device memory capacity of 32GiB and bandwidth of 820 GB/s. -* 32 DMA (Direct Memory Access) engines to move data within and across devices. +- 32 DMA (Direct Memory Access) engines to move data within and across devices. -* 6 CC-Cores for collective communication. - -* 2 (Inferentia2) or 4 (Trainium) NeuronLink-v2 for device-to-device collective communication. +- 6 CC-Cores for collective communication. +- 2 (Inferentia2) or 4 (Trainium) NeuronLink-v2 for device-to-device collective communication. > **Figure: neuron device2** > @@ -28,6 +27,7 @@ At a high level, both Trainium and Inferentia2 devices consist of: > This diagram compares the architecture of two AWS Neuron devices side by side. > > **Left side - Trainium**: +> > - Title "Trainium" in blue at top left > - Host PCIe interface at top right > - "32x DMA" and "6x CC-Core" blocks below Host PCIe @@ -42,6 +42,7 @@ At a high level, both Trainium and Inferentia2 devices consist of: > - Four "NeuronLink-v2" blocks at bottom for inter-device communication > > **Right side - Inferentia2**: +> > - Title "Inferentia2" in blue at top > - Same basic layout as Trainium > - Host PCIe at top @@ -58,6 +59,7 @@ At a high level, both Trainium and Inferentia2 devices consist of: > Both devices share the NeuronCore-v2 architecture but Trainium has more NeuronLink-v2 connections (4 vs 1), reflecting its focus on training workloads requiring more inter-device communication. > > **Key Elements:** +> > - **Trainium**: Training-focused device (left) > - **Inferentia2**: Inference-focused device (right) > - **NeuronCore-v2**: Two compute cores per device @@ -69,7 +71,6 @@ At a high level, both Trainium and Inferentia2 devices consist of: > - **NeuronLink-v2**: Inter-device links (4 for Trainium, 1 for Inferentia2) > - **Host PCIe**: Host interface - Fig. 47 Trainium/Inferentia2 Device Diagrams. The rest of this guide will go into details of each compute engine in NeuronCore-v2 and supported data movement @@ -93,14 +94,12 @@ Within each NeuronCore, there is also a Sync Engine, which functions as an engin In addition, it is often useful to take engine data-path width and frequency into account when optimizing performance for a multi-engine operator: - -| Device Architecture | Compute Engine | Data-path Width (elements/cycle) | Frequency (GHz) | -| --- | --- | --- | --- | -| Trainium/Inferentia2 | Tensor | 2x128 (input); 1x128 (output) | 2.8 | -| Vector | 128 input/output | 1.12 | -| Scalar | 1.4 | -| GpSimd | 1.4 | - +| Device Architecture | Compute Engine | Data-path Width (elements/cycle) | Frequency (GHz) | +| -------------------- | ---------------- | -------------------------------- | --------------- | +| Trainium/Inferentia2 | Tensor | 2x128 (input); 1x128 (output) | 2.8 | +| Vector | 128 input/output | 1.12 | +| Scalar | 1.4 | +| GpSimd | 1.4 | Memory-wise, a NeuronCore-v2 consists of two software-managed on-chip SRAMs, a 24MiB SBUF as the main data storage and a 2MiB PSUM as a dedicated accumulation buffer for Tensor Engine. Both SBUF and PSUM are considered two-dimensional memories @@ -109,11 +108,11 @@ more details on data movements with SBUF/PSUM later [here](#arch-sec-data-moveme The rest of this section will cover the following topics for each compute engine: -* Key functionalities. +- Key functionalities. -* Layout and tile size requirement for input and output tensors. +- Layout and tile size requirement for input and output tensors. -* Best practices to achieve good performance on the engine. +- Best practices to achieve good performance on the engine. ### Tensor Engine @@ -133,7 +132,6 @@ is always in FP32. and PSUM as below. Note, PSUM partition dimension is purposely rotated 90 degrees compared to SBUF partition dimension due to systolic array data flow. - > **Figure: tensor engine** > > A block diagram showing the Tensor Engine's 128x128 systolic array architecture with its connections to SBUF (input) and PSUM (output), illustrating the data flow for matrix multiplication operations. @@ -141,12 +139,14 @@ due to systolic array data flow. > This diagram illustrates the core architecture of the NeuronCore Tensor Engine and its relationship with the on-chip memory buffers used for matrix multiplication operations. > > **Tensor Engine (Left, Green Square):** +> > - Represented as a large green square > - Dimensions labeled: "128 rows" (vertical, left side) x "128 columns" (horizontal, top) > - Represents a 128x128 systolic array for matrix multiplication > - Label: "Tensor Engine" centered in the block > > **SBUF - State Buffer (Right, Blue Rectangle):** +> > - Represented as a tall blue rectangle to the right of the Tensor Engine > - Height labeled: "128 partitions" (right side, vertical) > - Provides input operands to the Tensor Engine @@ -154,6 +154,7 @@ due to systolic array data flow. > - Label: "SBUF" centered in the block > > **PSUM - Partial Sum Buffer (Bottom, Blue Rectangle):** +> > - Represented as a wide blue rectangle below the Tensor Engine > - Width labeled: "128 partitions" (bottom, horizontal) > - Receives output/accumulation results from the Tensor Engine @@ -161,28 +162,31 @@ due to systolic array data flow. > - Label: "PSUM" centered in the block > > **Data Flow Pattern:** +> > 1. Input data streams from SBUF (right) into the Tensor Engine > 2. Matrix multiplication is performed in the 128x128 systolic array > 3. Results accumulate into PSUM (bottom) > > **Dimensional Alignment:** +> > - SBUF's 128 partitions align with the Tensor Engine's 128 columns (for "moving" matrix) > - PSUM's 128 partitions align with the Tensor Engine's 128 rows (for output) > - This creates a natural flow for matrix operations where one operand is stationary and one "moves" through the array > > **Key Architecture Insights:** +> > - 128x128 = 16,384 multiply-accumulate units operating in parallel > - SBUF provides streaming input at 10 TB/s bandwidth > - PSUM accumulates partial results for large matrix products > > **Key Elements:** +> > - **Tensor Engine (128x128)**: Systolic array for matrix multiplication > - **SBUF**: Input buffer with 128 partitions feeding the engine > - **PSUM**: Output accumulator with 128 partitions receiving results > - **Arrows**: Data flow from SBUF to Engine to PSUM > - **Partition alignment**: 128 partitions match engine dimensions - Fig. 49 Tensor Engine and SRAM Connectivity. As shown in the diagram above, TensorE must **read** input matrices from **SBUF** and **write** output matrices to **PSUM**. @@ -201,11 +205,11 @@ in an output with dimensions `[M,N]`. For every `nki.isa.nc_matmul(stationary, moving)` call, TensorE executes two distinct Neuron ISA instructions: -* LoadStationary (short for LS): This instruction loads the `stationary` from SBUF and caches it in internal storage of TensorE +- LoadStationary (short for LS): This instruction loads the `stationary` from SBUF and caches it in internal storage of TensorE -* MultiplyMoving (short for MM): This instruction loads the `moving` from SBUF and multiplies `moving` across the pre-loaded -`stationary` matrix from the previous LoadStationary instruction. The output of this instruction is the -output of the `nki.isa.nc_matmul` call written to PSUM. +- MultiplyMoving (short for MM): This instruction loads the `moving` from SBUF and multiplies `moving` across the pre-loaded + `stationary` matrix from the previous LoadStationary instruction. The output of this instruction is the + output of the `nki.isa.nc_matmul` call written to PSUM. With the above instruction sequence, we as NKI programmers effectively map input tile `stationary` as the stationary tensor and input tile `moving` as the moving tensor for TensorE. As a rule-of-thumb for layout analysis, the **free** axis of the @@ -213,7 +217,6 @@ and input tile `moving` as the moving tensor for TensorE. As a rule-of-thumb for **moving** tensor becomes the free axis of the output. [Fig 50](#fig-arch-matmul) below visualizes this concept by showing a matrix multiplication in both mathematical and TensorE views. - > **Figure: matmul** > > A comprehensive diagram comparing the mathematical view of matrix multiplication with the NeuronCore Tensor Engine implementation view, showing how matrices map to stationary (Tensor Engine), moving (SBUF), and output (PSUM) components. @@ -221,12 +224,14 @@ by showing a matrix multiplication in both mathematical and TensorE views. > This diagram is divided into two parts by a vertical dashed line, illustrating how mathematical matrix multiplication maps to NeuronCore hardware. > > Part (a) "Mathematical View" (left side) shows standard matrix multiplication: +> > - A blue matrix "y" at the top with dimensions N (width) by K (height) > - A green matrix "x" at the bottom left with dimensions K (width) by M (height) > - A purple matrix "output" at the bottom right with dimensions N (width) by M (height) -> - The matrices are arranged to show x * y = output multiplication +> - The matrices are arranged to show x \* y = output multiplication > > Part (b) "Tensor Engine View" (right side) shows the hardware mapping: +> > - A green matrix labeled "stationary (Tensor Engine)" with dimensions M (stationary_fsize) width by K height - this matrix is loaded into the Tensor Engine and held stationary > - A blue matrix labeled "moving (SBUF)" with dimensions N (moving_fsize) width by K (rhs_psize) height - this matrix streams from the State Buffer > - A purple matrix labeled "output (PSUM)" with dimensions N (moving_fsize) width by M (stationary_fsize) height - partial sums accumulate here @@ -234,13 +239,15 @@ by showing a matrix multiplication in both mathematical and TensorE views. > - A "Copy" arrow shows the PSUM output being copied to a final "output (SBUF)" tensor with dimensions N width by M height, stored in State Buffer > > Dimension annotations include: +> > - M (stationary_fsize): Free dimension size of stationary matrix -> - N (moving_fsize): Free dimension size of moving matrix +> - N (moving_fsize): Free dimension size of moving matrix > - K (lhs_psize, rhs_psize): Contraction dimension > - PSUM P-dim and SBUF P-dim labels indicate partition dimension orientations > > **Key Elements:** -> - **Mathematical View (a)**: Standard matrix multiplication x * y = output +> +> - **Mathematical View (a)**: Standard matrix multiplication x \* y = output > - **Tensor Engine View (b)**: Hardware-mapped implementation > - **stationary (Tensor Engine)**: Green matrix held in Tensor Engine > - **moving (SBUF)**: Blue matrix streamed from State Buffer @@ -249,7 +256,6 @@ by showing a matrix multiplication in both mathematical and TensorE views. > - **M, N, K dimensions**: Matrix dimension labels > - **Copy arrow**: Data movement from PSUM to SBUF - Fig. 50 MxKxN Matrix Multiplication Visualization. However, programmers are also free to map `stationary` tile to the moving tensor instead, which would lead to the same output tile @@ -261,19 +267,19 @@ for more discussion. **Tile Size.** The `nki.isa.nc_matmul` API enforces the following constraints on the input/output tile sizes: -* `stationary` tensor free axis size (`stationary_fsize`) must never exceed 128, due to the number of PE columns in TensorE. +- `stationary` tensor free axis size (`stationary_fsize`) must never exceed 128, due to the number of PE columns in TensorE. -* `stationary/moving` tensor partition axis size (`stationary_psize/moving_psize`) must never exceed 128, due to the number of PE rows and -also the number of SBUF partitions. +- `stationary/moving` tensor partition axis size (`stationary_psize/moving_psize`) must never exceed 128, due to the number of PE rows and + also the number of SBUF partitions. -* `moving` tensor free axis size (`moving_fsize`) must never exceed 512, due to the fact that each `nc_matmul` can only write -to a single PSUM bank, which can only hold 512 FP32 elements per PSUM partition. +- `moving` tensor free axis size (`moving_fsize`) must never exceed 512, due to the fact that each `nc_matmul` can only write + to a single PSUM bank, which can only hold 512 FP32 elements per PSUM partition. When the shapes of the input matrices defined in the user-level operator exceed any of the above tile size limitation, we must tile the input matrices and invoke multiple `nki.isa.nc_matmul` calls to perform the matrix multiplication. Exceeding the `stationary_fsize` (#1) or `moving_fsize` (#3) tile limitations for M or N should lead to fully independent `nki.isa.nc_matmul` with disjoint output tiles. However, when `K` exceeds the `stationary_psize/moving_psize` limit, we need to tile the input matrices -in the contraction dimension and invoke multiple `nki.isa.nc_matmul` to accumulate into the *same* output buffer in PSUM. +in the contraction dimension and invoke multiple `nki.isa.nc_matmul` to accumulate into the _same_ output buffer in PSUM. Refer to the [Tiling Matrix Multiplications](../programming/tutorials/matrix_multiplication.md#tutorial-matmul-tiling) tutorial for a NKI code example. @@ -287,7 +293,6 @@ As an example, we can perform a 128x128 matrix transposition (i.e., swap the fre `identity` is a 128x128 identity matrix. In fact, this is exactly what nki.isa.nc_transpose() does, when TensorE is chosen as the compute engine. - > **Figure: mm transpose** > > A diagram showing how to implement matrix transpose using the Tensor Engine by multiplying with an identity matrix, producing the transposed result in PSUM and then copying to SBUF. @@ -305,13 +310,15 @@ as the compute engine. > **Top right** - Blue matrix "x^T (SBUF)" with dimensions N (width) by M (height), the final transposed result in State Buffer. > > Arrows show the flow: +> > - Arrow from Identity to x (indicating multiplication setup) > - Arrow from x down to x^T (PSUM) > - Curved arrow labeled "Copy" from x^T (PSUM) to x^T (SBUF) > -> The mathematical insight is: x * I = x^T when x is loaded as the moving matrix and I as stationary, effectively transposing the result due to the Tensor Engine's output layout. +> The mathematical insight is: x \* I = x^T when x is loaded as the moving matrix and I as stationary, effectively transposing the result due to the Tensor Engine's output layout. > > **Key Elements:** +> > - **x**: Green input matrix [M x N] > - **Identity**: Blue identity matrix [N x N] with diagonal 1s > - **x^T (PSUM)**: Purple transposed result in Partial Sum @@ -320,7 +327,6 @@ as the compute engine. > - **M, N dimensions**: Matrix dimension labels > - **1s and 0s**: Identity matrix pattern - Fig. 51 Transposition. Similarly, we can broadcast a vector occupying a single partition to M (M <= 128) partitions using `nki.isa.nc_matmul(ones, @@ -328,7 +334,6 @@ broadcast_input, is_stationary_onezero=True)`, where `ones` is a 1xM vector fill the vector to be broadcast. In fact, NKI invokes such matmul under the hood when `broadcast_input.broadcast_to((M, broadcast_input.shape[1]))` is called. - > **Figure: mm broadcast** > > A diagram showing how to implement broadcast operations using matrix multiplication, where a vector y is broadcast to a full matrix y_bcast by multiplying with a ones vector, then copying from PSUM to SBUF. @@ -336,6 +341,7 @@ is called. > This diagram illustrates a technique for implementing tensor broadcast using the Tensor Engine's matrix multiplication capability on NeuronCore. > > At the top of the diagram, two input tensors are shown: +> > - A green horizontal tensor of all ones (labeled "1, 1, ..., 1") with dimension M (width) by 1 (height), representing a ones vector > - A blue horizontal tensor "y" with dimension N (width) by 1 (height), representing the input vector to be broadcast > @@ -346,11 +352,13 @@ is called. > A curved arrow labeled "Copy" points from y_bcast (PSUM) to a blue square tensor "y_bcast (SBUF)" on the right, with dimensions N (width) by M (height). This represents copying the broadcast result from PSUM to the State Buffer. > > The dimension annotations show: +> > - M: Width of the ones vector and the broadcast output > - N: Height of the input vector y and the broadcast output > - The tensor is effectively broadcast from shape [1, N] to [M, N] > > **Key Elements:** +> > - **Ones vector**: Green tensor of all 1s with shape [1, M] > - **y**: Blue input vector to broadcast with shape [1, N] > - **y_bcast (PSUM)**: Purple broadcast result in Partial Sum buffer [N, M] @@ -358,7 +366,6 @@ is called. > - **Copy arrow**: Data movement from PSUM to SBUF > - **M, N dimensions**: Size annotations for the broadcast operation - Fig. 52 Partition Broadcast. In general, we can achieve many more complex data reshapes in TensorE, such as shuffling partitions of a SBUF tensor, by @@ -372,7 +379,6 @@ best use of TensorE. If you can do summation within each partition (F-dim summat for an alternative reduction implementation on Vector Engine. It is recommended to choose the engine based on the natural layout of your input data to avoid any transpositions. - > **Figure: mm cross partition** > > A diagram showing how to perform cross-partition reduction using matrix multiplication, where a vector y is multiplied with a ones vector to produce a scalar sum in PSUM, then copied to SBUF. @@ -380,6 +386,7 @@ layout of your input data to avoid any transpositions. > This diagram illustrates using matrix multiplication to implement cross-partition reduction (summing across partitions) on NeuronCore. > > On the left side, two input vectors are shown vertically: +> > - A green vertical tensor of all ones (labeled "1, 1, ..., 1") with dimension 1 (width) by N (height) > - A blue vertical tensor "y" with dimension 1 (width) by N (height) > @@ -390,12 +397,14 @@ layout of your input data to avoid any transpositions. > A curved arrow labeled "Copy" points from sum (PSUM) to a small blue square "sum (SBUF)" in the upper right, representing the final scalar result copied to the State Buffer. > > The key insight is that multiplying a vector by a ones vector computes the sum of all elements: -> - y^T * ones = sum(y) +> +> - y^T \* ones = sum(y) > - This effectively performs a reduction across the N partition dimension > > This technique is useful when you need to sum values across partitions but only have access to the Tensor Engine for computation. > > **Key Elements:** +> > - **Ones vector**: Green tensor of all 1s with shape [N, 1] > - **y**: Blue input vector with shape [N, 1] > - **sum (PSUM)**: Purple scalar sum result in Partial Sum buffer @@ -404,18 +413,17 @@ layout of your input data to avoid any transpositions. > - **N dimension**: Length of the vectors being reduced > - **1 dimension**: Single element width/output - Fig. 53 Cross-Partition Accumulation As TensorE is the most performant compute engine of the NeuronCore in terms of FLOPS, the goal is to have it execute meaningful -computation at high utilization as much as possible. The above “alternative use cases” stop TensorE from performing *useful* -computations at *high* throughput and therefore, should generally be avoided. However, there are situations where it is +computation at high utilization as much as possible. The above “alternative use cases” stop TensorE from performing _useful_ +computations at _high_ throughput and therefore, should generally be avoided. However, there are situations where it is advisable to use them: -* Operators that do not require heavy matmuls anyhow, e.g. normalization, softmax. +- Operators that do not require heavy matmuls anyhow, e.g. normalization, softmax. -* Layout conflicts between producer and consumer engines where broadcast/transpose are absolutely unavoidable (see example -in fused attention tutorial). +- Layout conflicts between producer and consumer engines where broadcast/transpose are absolutely unavoidable (see example + in fused attention tutorial). #### **Performance Consideration** @@ -423,18 +431,18 @@ As a rule of thumb, TensorE can achieve the best throughput when it runs many ba input matrices at the largest possible tiles sizes (`stationary` is 128x128 and `moving` is 128x512). In this ideal scenario, TensorE sees the below instruction sequence: -* `LoadStationary (LS[0])` (128x128) +- `LoadStationary (LS[0])` (128x128) -* `MultiplyMoving (MM[0])` (128x512) +- `MultiplyMoving (MM[0])` (128x512) -* `LoadStationary (LS[1])` (128x128) +- `LoadStationary (LS[1])` (128x128) -* `MultiplyMoving (MM[1])` (128x512) +- `MultiplyMoving (MM[1])` (128x512) -* … +- … **Cost Model:** TensorE is a deeply pipelined engine; therefore, the engine can have several `LS&MM` instruction pairs -in-flight at a given time. Due to this pipelining nature, it is often *not* useful to use end-to-end execution *latency* +in-flight at a given time. Due to this pipelining nature, it is often _not_ useful to use end-to-end execution _latency_ of a single instruction when estimating the instruction cost. Instead, we can focus on the **initiation interval** of such instructions, that is, the number of cycles between successive instruction launches. Therefore, we can estimate the cost of an instruction `I` by how soon TensorE can issue the next instruction after `I`. @@ -448,7 +456,6 @@ FP32 input matrix data type to one of BF16/FP16/TF32/cFP8 before performing matr Figure below visualizes two pipelined `MM` instructions: - > **Figure: mm pipeline** > > A timing diagram showing how matrix multiplication operations are pipelined across multiple pipeline stages (0, 1, through P-1), with annotations for initiation interval and full execution latency. @@ -458,18 +465,22 @@ Figure below visualizes two pipelined `MM` instructions: > The diagram shows multiple horizontal timelines representing different pipeline stages: > > **Pipeline Stage 0** (top): +> > - Shows two consecutive operations MM[0] (blue) and MM[1] (purple) > - Operations are displayed as rectangular blocks on the timeline > > **Pipeline Stage 1** (second row): +> > - Same MM[0] and MM[1] operations, but shifted right by one cycle > - The offset demonstrates the pipeline initiation interval > > **Pipeline Stage P-1** (bottom row): +> > - Shows the same operations at the end of the pipeline > - MM[0] and MM[1] blocks appear much later in time > > Key timing annotations: +> > - **Initiation Interval**: Marked with a green double-headed arrow at the top, showing the time between starting successive operations (approximately one cycle) > - **1 cycle**: Small annotation showing the cycle boundary > - **Full execution Latency of MM[0] on TensorE**: A long green double-headed arrow at the bottom spanning from when MM[0] enters Pipeline Stage 0 to when it exits Pipeline Stage P-1 @@ -477,6 +488,7 @@ Figure below visualizes two pipelined `MM` instructions: > Vertical dashed lines help align the timing across pipeline stages. Ellipsis (...) between Stage 1 and Stage P-1 indicates intermediate pipeline stages not shown. > > **Key Elements:** +> > - **Pipeline Stage 0, 1, P-1**: Multiple pipeline stages in Tensor Engine > - **MM[0], MM[1]**: Consecutive matrix multiplication operations > - **Initiation Interval**: Time between starting new operations @@ -485,7 +497,6 @@ Figure below visualizes two pipelined `MM` instructions: > - **Blue/purple blocks**: Color-coded operations showing overlap > - **Dashed vertical lines**: Timing alignment markers - Fig. 54 Pipelined multiplyMoving instructions. **Background LoadStationary:** In typical workloads, TensorE would be alternating between LS and MM instructions with different @@ -496,7 +507,6 @@ As a result, depending on the relative sizes of the `stationary` and `moving` ma TensorE performance can be bounded by either `LS` or `MM` instructions. Figure below visualizes these two cases. In the ideal scenario where `stationary` and `moving` use the largest tile sizes, TensorE should operate in case (a). - > **Figure: mm bottleneck** > > A timing diagram comparing two execution scenarios for matrix multiplication: MultiplyMoving Bounded (where compute is the bottleneck) and LoadStationary Bounded (where memory loading is the bottleneck). @@ -504,12 +514,14 @@ the ideal scenario where `stationary` and `moving` use the largest tile sizes, T > This diagram shows two execution timeline scenarios illustrating different bottleneck conditions in matrix multiplication on NeuronCore, helping developers understand performance limiting factors. > > Part (a) "MultiplyMoving Bounded" (top section) shows two parallel timelines: +> > - **LoadStationary row**: Shows sequential loading operations LS[0], LS[1], LS[2], LS[3], ... with blocks colored in shades of blue/green. These complete relatively quickly with gaps between them. > - **MultiplyMoving row**: Shows sequential computation operations MM[0], MM[1], MM[2], MM[3], ... with blocks colored in shades of blue, green, and purple. These operations are longer and continuous, forming the critical path. > > In this scenario, LoadStationary completes before MultiplyMoving needs the data, indicating compute is the bottleneck. The computation (MultiplyMoving) takes longer than data loading (LoadStationary). > > Part (b) "LoadStationary Bounded" (bottom section) shows two parallel timelines: +> > - **LoadStationary row**: Shows the same LS[0] through LS[3] operations, but now they are longer and form a continuous sequence. > - **MultiplyMoving row**: Shows MM[0] through MM[3] operations with gaps between them, waiting for data to be loaded. > @@ -518,6 +530,7 @@ the ideal scenario where `stationary` and `moving` use the largest tile sizes, T > Both timelines have arrows extending to the right with ellipsis (...) indicating the pattern continues. > > **Key Elements:** +> > - **LoadStationary (LS)**: Operations loading the stationary matrix into Tensor Engine > - **MultiplyMoving (MM)**: Matrix multiplication operations with moving matrix > - **LS[0]-LS[3]**: Individual load operations (blue/teal colors) @@ -527,7 +540,6 @@ the ideal scenario where `stationary` and `moving` use the largest tile sizes, T > - **Timeline arrows**: Show execution sequence over time > - **Gaps vs continuous**: Visual indication of which operation is bottleneck - Possible execution timeline execution with background LoadStationary **Fast LoadStationary:** Since `LoadStationary` is a pure data movement with no computation, TensorE can perform `LoadStationary` @@ -572,7 +584,6 @@ below shows connectivity between SBUF and VectorE banks. VectorE consists of fou Bank connects to 32 SBUF/PSUM partitions and outputs 32 parallel streams of data, while each Compute Bank can process 32 parallel data streams using 32 vector lanes. The Compute Bank can write back to 32 SBUF/PSUM partitions. - > **Figure: vector engine cross partition** > > A diagram showing the Vector Engine architecture with 128 SBUF/PSUM partitions mapped to 4 banks (32 partitions each), illustrating how the Reshape Banks and Compute Banks enable cross-partition operations. @@ -581,6 +592,7 @@ parallel data streams using 32 vector lanes. The Compute Bank can write back to > > **Left Column - SBUF/PSUM (Input):** > A vertical stack of 128 partitions labeled "SBUF/PSUM" at the top: +> > - **Bank 0 partitions (Purple)**: p[0], p[1], ..., p[31] > - **Bank 1 partitions (Blue)**: p[32], p[33], ..., p[63] > - **Bank 2 partitions (Green)**: p[64], p[65], ..., p[95] @@ -591,6 +603,7 @@ parallel data streams using 32 vector lanes. The Compute Bank can write back to > > **Middle Column - Reshape Banks:** > Four "Reshape Bank" blocks corresponding to the partition groupings: +> > - **Reshape Bank[0]** (Purple): Handles partitions p[0]-p[31] > - **Reshape Bank[1]** (Blue): Handles partitions p[32]-p[63] > - **Reshape Bank[2]** (Green): Handles partitions p[64]-p[95] @@ -601,6 +614,7 @@ parallel data streams using 32 vector lanes. The Compute Bank can write back to > > **Right Column - Compute Banks:** > Four "Compute Bank" blocks matching the Reshape Banks: +> > - **Compute Bank[0]** (Purple) > - **Compute Bank[1]** (Blue) > - **Compute Bank[2]** (Green) @@ -613,6 +627,7 @@ parallel data streams using 32 vector lanes. The Compute Bank can write back to > The right two columns (Reshape Banks and Compute Banks) are grouped under the "Vector Engine" label. > > **Data Flow:** +> > 1. Data enters from SBUF/PSUM partitions > 2. Partitions are grouped into 4 banks of 32 partitions each > 3. Reshape Banks reorganize data for cross-partition operations @@ -620,6 +635,7 @@ parallel data streams using 32 vector lanes. The Compute Bank can write back to > 5. Results flow out for further processing or storage > > **Key Elements:** +> > - **128 partitions**: Total SBUF/PSUM partition count > - **4 banks**: Partitions grouped into banks of 32 > - **Reshape Bank[0-3]**: Data reorganization for cross-partition ops @@ -627,19 +643,18 @@ parallel data streams using 32 vector lanes. The Compute Bank can write back to > - **Color coding**: Purple (0-31), Blue (32-63), Green (64-95), Orange (96-127) > - **Cross-partition capability**: Enables operations across partition boundaries within a bank - Fig. 56 Vector Engine reshape and compute banks. The Reshape Bank supports the following data movement: -* *32x32 transpose*: Each Reshape Bank can read in 32 elements per SBUF/PSUM partitions and transpose the partition and -free dimension of the incoming 32x32 matrix. This can be invoked by [nki.isa.nc_transpose](../programming/api/api-nki-isa-tensor.md#nki-isa-nc_transpose) -API by selecting VectorE as the execution engine. +- _32x32 transpose_: Each Reshape Bank can read in 32 elements per SBUF/PSUM partitions and transpose the partition and + free dimension of the incoming 32x32 matrix. This can be invoked by [nki.isa.nc_transpose](../programming/api/api-nki-isa-tensor.md#nki-isa-nc_transpose) + API by selecting VectorE as the execution engine. -* *32 partition shuffle*: Each Reshape Bank can take an arbitrary *shuffle mask* -`SM`* of length 32. The integer value of `SM[i]` indicates the source partition ID (modulo 32) that the Reshape Bank -output stream `i` will get. For example, we can broadcast partition[0] to partition[0-31] using a SM of 32 zeros. -This can be invoked by [nki.isa.nc_stream_shuffle](../programming/api/api-nki-isa-tensor.md#nki-isa-nc_stream_shuffle) API. +- _32 partition shuffle_: Each Reshape Bank can take an arbitrary _shuffle mask_ + `SM`\* of length 32. The integer value of `SM[i]` indicates the source partition ID (modulo 32) that the Reshape Bank + output stream `i` will get. For example, we can broadcast partition[0] to partition[0-31] using a SM of 32 zeros. + This can be invoked by [nki.isa.nc_stream_shuffle](../programming/api/api-nki-isa-tensor.md#nki-isa-nc_stream_shuffle) API. Refer [here](#arch-sec-cross-partition-connect) later in this doc for cross-bank data movement. @@ -659,10 +674,10 @@ Vector instruction. Refer to NKI Performance Guide for more detailed discussion **Cost Model:** In the most common cases where the free axis size (`N`) of the input tile(s) is sufficiently large (`N > 128`), the execution cost of an instruction on VectorE is correlated to `N`: -* If there is only one input tile, most VectorE instructions can execute in roughly `N` cycles (example: -[nki.isa.tensor_scalar](../programming/api/api-nki-isa-tensor.md#nki-isa-tensor_scalar)) +- If there is only one input tile, most VectorE instructions can execute in roughly `N` cycles (example: + [nki.isa.tensor_scalar](../programming/api/api-nki-isa-tensor.md#nki-isa-tensor_scalar)) -* If there are two input tiles, the instruction can execute in roughly `2N` cycles (example: nki.isa.tensor_tensor) +- If there are two input tiles, the instruction can execute in roughly `2N` cycles (example: nki.isa.tensor_tensor) There are a few exceptions to the above rule, depending on the data types and instruction type. See [NKI ISA API doc](../programming/api/nki.isa.md) @@ -701,7 +716,6 @@ must not exceed 128, while the free dimension size can be up to 64K elements for Each ScalarE compute lane also supports an additional multiply-add **before** the non-linear function (`func`) is applied in a pipeline fashion. Mathematically, ScalarE implements: - ```python # Case 1: scale is SBUF/PSUM vector # Input: 2D in_tile, 1D scale, 1D bias @@ -718,7 +732,6 @@ for lane_id in range(in_tile.shape[0]): + bias[lane_id]) ``` - This functionality can be invoked using the [nki.isa.activation](../programming/api/api-nki-isa-scalar.md#nki-isa-activation) API by specifying a `scale` for multiplication and `bias` for addition. The scale can either be a tile from SBUF/PSUM with one element/partition or a compile-time constant. On the other hand, the bias can only be a tile from SBUF/PSUM with @@ -732,7 +745,6 @@ in a pipeline fashion. On NeuronCore-v2, the reduction operator can only be addi Mathematically, ScalarE with accumulation enabled implements: - ```python # Input: 2D in_tile, 1D scale (similarly for scalar scale), 1D bias # Output: 2D out_tile, 1D reduce_res @@ -743,7 +755,6 @@ for lane_id in range(in_tile.shape[0]): reduce_res[lane_id] += out_tile[lane_id][k] ``` - This functionality can be invoked using the [nki.isa.activation_reduce](../programming/api/api-nki-isa-scalar.md#nki-isa-activation_reduce) API by specifying `reduce_op` as `nki.language.add` and `reduce_res` as the output reduction tile, passed by reference. @@ -783,11 +794,11 @@ uses a vectorized kernel implementation that Neuron engineers hand-tune for the **Data Types.** Each processor in GpSimd supports vectorized computation for -* 16x FP32/INT32/UINT32, or +- 16x FP32/INT32/UINT32, or -* 32x FP16/INT16/UINT16, or +- 32x FP16/INT16/UINT16, or -* 64x INT8/UINT8 +- 64x INT8/UINT8 This is in contrast to ScalarE/VectorE which can only perform arithmetic operations in FP32. However, if the GpSimdE program chooses to, it can also access SBUF data of any [supported data types in NKI](../programming/api/nki.api.shared.md#nki-dtype) @@ -811,7 +822,6 @@ from all 16 connected partitions collectively (up to 32-bit per partition) to th 512-bit SIMD width. Similarly, on the write side, the tensor write interface can accept 512-bit of data for writing back to the connected SBUF partitions per cycle. - > **Figure: gpsimd sbuf connectivity** > > A connectivity diagram showing how SBUF (State Buffer) partitions map to GpSimd Engine cores, with 8 partition groups each connecting to one of 8 GpSimd cores via bidirectional arrows. @@ -819,6 +829,7 @@ to the connected SBUF partitions per cycle. > This diagram illustrates the memory-to-compute connectivity between the State Buffer (SBUF) and the GpSimd (General Purpose SIMD) Engine in NeuronCore architecture. > > On the left side, a column labeled "SBUF" (in italics) contains 8 stacked rectangular blocks in light blue/cyan color, each representing a group of 16 partitions: +> > - Partition [0-15] at the top > - Partition [16-31] > - Partition [32-47] @@ -829,6 +840,7 @@ to the connected SBUF partitions per cycle. > - Partition [112-127] at the bottom > > On the right side, a column labeled "GpSimd Engine" (in italics) contains 8 stacked rectangular blocks in light purple/lavender color, each representing a compute core: +> > - Core[0] at the top > - Core[1] > - Core[2] @@ -841,6 +853,7 @@ to the connected SBUF partitions per cycle. > Between each SBUF partition group and its corresponding GpSimd core, bidirectional arrows (pointing both left and right) indicate two-way data flow. Each partition group connects exclusively to one core: Partition [0-15] connects to Core[0], Partition [16-31] connects to Core[1], and so on through Partition [112-127] connecting to Core[7]. > > **Key Elements:** +> > - **SBUF**: State Buffer memory divided into 8 partition groups > - **Partition [0-15] through [112-127]**: Eight groups of 16 partitions each > - **GpSimd Engine**: General Purpose SIMD compute engine with 8 cores @@ -848,7 +861,6 @@ to the connected SBUF partitions per cycle. > - **Bidirectional arrows**: Two-way data flow between partition groups and cores > - **One-to-one mapping**: Each partition group maps to exactly one GpSimd core - Fig. 57 Connectivity between GpSimdE and SBUF. #### **Performance Consideration** @@ -870,7 +882,6 @@ and also how to do it efficiently. As a reminder, there are three main types of PSUM, from highest to lowest capacity. Figure below shows the specifications of these memories and their connectivity for one NeuronCore-v2: - > **Figure: memory hierarchy** > > A hierarchical architecture diagram showing the NeuronCore memory system with on-chip components (PSUM, compute engines, SBUF) and off-chip HBM, connected through DMA engines. @@ -880,6 +891,7 @@ for one NeuronCore-v2: > At the top of the on-chip section (enclosed in a dashed rectangle), the "PSUM" block (peach/orange color) spans the full width, representing the Partial Sum accumulator memory. > > Below PSUM, four compute engine blocks are arranged horizontally: +> > - "TensorE" (Tensor Engine) - leftmost > - "VectorE" (Vector Engine) - second from left > - "ScalarE" (Scalar Engine) - third from left @@ -898,6 +910,7 @@ for one NeuronCore-v2: > The bidirectional arrows throughout indicate data can flow in both directions between all connected components. > > **Key Elements:** +> > - **PSUM**: Partial Sum accumulator at top (peach/orange) > - **TensorE**: Tensor Engine compute unit > - **VectorE**: Vector Engine compute unit @@ -909,7 +922,6 @@ for one NeuronCore-v2: > - **on-chip / off-chip**: Labels indicating memory hierarchy levels > - **Bidirectional arrows**: Data flow between all components - Fig. 58 Memory hierarchy. As shown in the above figure, data movement between HBM and SBUF is performed using on-chip DMA @@ -918,9 +930,9 @@ parallel to computation within the NeuronCore. Data movement between PSUM and SB compute engines. However, different compute engines have different connectivity to SBUF/PSUM as indicated by the arrows in the figure. In addition, NeuronCore-v2 has the following restrictions: -* VectorE and GpSimdE cannot access SBUF in parallel. +- VectorE and GpSimdE cannot access SBUF in parallel. -* VectorE and ScalarE cannot access PSUM in parallel. +- VectorE and ScalarE cannot access PSUM in parallel. Therefore, VectorE and GpSimdE instructions that access SBUF must be serialized, similarly for VectorE and ScalarE instructions that access PSUM. This is enforced by Neuron Compiler during NKI kernel compilation, so NKI developers are not required @@ -928,11 +940,11 @@ to program such serializations. The rest of this section will discuss the following topics in detail: -* Data movement between HBM and SBUF using DMAs. +- Data movement between HBM and SBUF using DMAs. -* Accessing SBUF/PSUM tensors using compute engines. +- Accessing SBUF/PSUM tensors using compute engines. -* In-memory accumulation using TensorE and PSUM. +- In-memory accumulation using TensorE and PSUM. ### Data movement between HBM and SBUF using DMAs @@ -957,7 +969,6 @@ address map. `sbuf_base_addr` is a 64-bit address dependent on which NeuronCore-v2 on the device the SBUF is located in. The SBUF addresses start from the first byte of partition 0, increment along the free dimension first and then advance onto the next partition. - > **Figure: sbuf addr space** > > A diagram illustrating the State Buffer (SBUF) address space organization showing the two-dimensional addressing scheme with Partition dimension (128 partitions) and Free dimension, along with base address offsets. @@ -969,15 +980,18 @@ increment along the free dimension first and then advance onto the next partitio > > **Address Markers (Top, Green Text):** > Three address markers are shown along the top of the diagram indicating memory positions in the Free dimension: +> > - **sbuf_base_addr**: Starting address (leftmost position) > - **sbuf_base_addr+1**: Second position > - **sbuf_base_addr+192KiB**: Position at 192 KiB offset (right side) > > **Additional Address (Left Side, Green Text):** +> > - **sbuf_base_addr+256KiB**: Shows the address after traversing through partition addresses > > **Partition Dimension (Right Side Labels):** > The vertical axis shows partition indices with labels: +> > - Partition 0 (top) > - Partition 1 > - Partition 2 @@ -987,22 +1001,26 @@ increment along the free dimension first and then advance onto the next partitio > - Partition 127 (bottom) > > **Dimension Labels:** +> > - **Partition (P) Dimension**: Labeled on the right side with bidirectional arrow > - **Free (F) Dimension**: Labeled at the bottom with bidirectional arrow > > **Visual Elements:** +> > - Green highlighted cells at the top-left corner showing the first few elements > - Dashed lines indicating address boundaries > - Arrow from sbuf_base_addr pointing to the top-left cell > - The grid structure shows 128 rows (partitions) and variable columns (free dimension) > > **Memory Layout Interpretation:** +> > - The Free dimension extends horizontally with 192 KiB of addressable space per partition > - The Partition dimension has 128 partitions (0-127) > - Total SBUF size: 128 partitions x 192 KiB = 24 MiB (approximate) > - The 256 KiB offset represents wrapping through partition addresses > > **Key Elements:** +> > - **sbuf_base_addr**: Starting address of SBUF allocation > - **128 Partitions**: Partition indices 0-127 along vertical axis > - **192 KiB per partition**: Free dimension size per partition @@ -1010,7 +1028,6 @@ increment along the free dimension first and then advance onto the next partitio > - **Green cells**: Highlighted memory elements being accessed > - **Partition stride**: Moving down increments partition, not contiguous in memory - Fig. 59 SBUF memory address space. As discussed in [NKI Language Guide](../programming/nki-language-guide.md), @@ -1030,7 +1047,6 @@ In NKI, moving data from HBM to SBUF and from SBUF to HBM are done with calls to assigning these transfers to different DMA engines. As an example, loading a 128x512 FP32 HBM tensor to SBUF is best done through 16 DMA transfers (one per DMA engine), each moving a scatter-gather list of 8 DMA buffers: - ```python import nki.language as nl import nki.isa as nisa @@ -1038,14 +1054,13 @@ tile = nl.ndarray((128, 512), dtype=in_tensor.dtype, buffer=nl.sbuf) nisa.dma_copy(dst=tile, src=in_tensor[0:128, 0:512]) ``` - To achieve good performance out of the DMAs, we generally aim to: -* Move a large amount of contiguous data in each DMA buffer to amortize DMA buffer overhead +- Move a large amount of contiguous data in each DMA buffer to amortize DMA buffer overhead -* Move a large amount of data in each DMA transfer to amortize DMA transfer overhead. +- Move a large amount of data in each DMA transfer to amortize DMA transfer overhead. -* Invoke as many parallel DMA transfers on the available DMA engines as possible. +- Invoke as many parallel DMA transfers on the available DMA engines as possible. These goals ultimately boil down to a quick optimization rule: maximize **both free (4KiB or above) and partition (ideally 128) dimension sizes** when moving tensors between SBUF and HBM using `nki.language.load` @@ -1064,10 +1079,9 @@ In every cycle, each engine can read 128 elements across 128 SBUF/PSUM partition perform a computation on previously read 128 elements, and write 128 previously computed results to SBUF/PSUM. In other words, the P axis of a tensor -is the *parallel* dimension for SBUF/PSUM data accessing, while the F axis of the tensor is the *time* dimension for data +is the _parallel_ dimension for SBUF/PSUM data accessing, while the F axis of the tensor is the _time_ dimension for data accessing. - > **Figure: data streaming** > > A time-sequenced diagram showing how data streams through a compute engine pipeline, illustrating the progression from source tensor (SBUF/PSUM) through the compute engine to destination tensor across multiple time steps (Time = 0, 1, through N). @@ -1081,6 +1095,7 @@ accessing. > At Time = N (bottom section), the full pipeline is active. The source tensor shows ellipsis (...) indicating ongoing data reads, the Compute Engine shows internal "Engine pipelines" with multiple processing stages indicated by ellipsis, and arrows flow to the destination tensor which also shows ellipsis indicating ongoing data writes. > > **Key Elements:** +> > - **SBUF/PSUM: src_tensor**: Source tensor providing input data (light blue) > - **Compute Engine**: Central processing unit (peach/orange color) with internal pipelines > - **SBUF/PSUM: dst_tensor**: Destination tensor receiving output data (dashed outlines initially) @@ -1090,7 +1105,6 @@ accessing. > - **Arrows**: Data flow direction from source through compute to destination > - **Engine pipelines**: Internal pipeline stages within the compute engine - Fig. 61 Data streaming between SBUF and compute engine. When accessing SBUF/PSUM tensors in an instruction, we need to follow different rules in the P and F dimensions. First, @@ -1098,11 +1112,11 @@ hardware does not allow P dimension striding when accessing data from a single S tensor of an instruction must occupy a continuous number of partitions. In addition, the hardware further enforces which partition a tensor can start from (`start_partition`) based on the number of partitions the tensor occupies (`num_partition`). This is currently handled by the tensor allocator in Neuron Compiler during NKI kernel compilation process: -* If `64 < num_partition <= 128`, `start_partition` must be 0 +- If `64 < num_partition <= 128`, `start_partition` must be 0 -* If `32 < num_partition <= 64`, `start_partition` must be 0 or 64 +- If `32 < num_partition <= 64`, `start_partition` must be 0 or 64 -* If `0 < num_partition <= 32`, `start_partition` must be one of 0/32/64/96 +- If `0 < num_partition <= 32`, `start_partition` must be one of 0/32/64/96 On the other hand, data accessing along the free dimension is a lot more flexible: the src/dst tensor of an engine instruction can support up to four-dimensional tensorized access pattern with a stride in each dimension @@ -1114,7 +1128,7 @@ to learn about how to index SBUF/PSUM tensors to achieve F dimension striding in Lastly, as implied in [Figure 61](#fig-arch-data-streaming), when accessing a SBUF/PSUM tensor, all active partitions must follow the same F dimension access pattern. In other words, -at every time step, the engine read/write interface will access data elements at the same *offset* within each active partition. +at every time step, the engine read/write interface will access data elements at the same _offset_ within each active partition. #### Cross-Partition Connectivity @@ -1130,7 +1144,6 @@ for `0 < num_partition <= 32`. Figure below illustrates these two patterns for ` The shaded portion of the `Engine` block indicates the active lanes for the given instruction. With these movement patterns, each partition in `src_tensor` still has a one-to-one mapping to each partition in `dst_tensor`. - > **Figure: cross quadrant** > > A technical diagram illustrating two types of tensor data movement patterns: Cross-Half Movement and Cross-Quadrant Movement, showing how data flows between SBUF/PSUM source tensors and destination tensors through the VectorE/ScalarE/GpSimdE compute engines. @@ -1142,6 +1155,7 @@ each partition in `src_tensor` still has a one-to-one mapping to each partition > In part (b) "Cross-Quadrant Movement", the layout is similar but with four distinct color-coded groups of partitions: partitions 0-1 (light blue), partitions 31-33 (light green), partitions 63-65 (orange), partitions 95-97 (light purple), and partition 127. This shows a more complex four-way redistribution pattern where data is exchanged across four quadrants of the partition space. > > **Key Elements:** +> > - **SBUF/PSUM: src_tensor**: Source tensor on the left side with partition dimension (P) vertical and free dimension (F) horizontal > - **VectorE/ScalarE/GpSimdE**: Central gray compute engine block processing the data movement > - **SBUF/PSUM: dst_tensor**: Destination tensor on the right side receiving redistributed data @@ -1150,7 +1164,6 @@ each partition in `src_tensor` still has a one-to-one mapping to each partition > - **Cross-Quadrant Movement (b)**: Four-way data exchange across quadrants with color-coded partitions (blue, green, orange, purple) > - **Ellipsis (...)**: Indicates additional partitions not explicitly shown in the diagram - Fig. 62 Cross-partition connectivity. #### Performance Consideration @@ -1166,13 +1179,13 @@ be half of the peak bandwidth, which translates to roughly 50% performance hit o **Concurrent SBUF/PSUM accesses by engines.** As mentioned earlier, NeuronCore-v2 has the following on-chip RAM access restrictions: -* Vector Engine and GpSimd Engine cannot access SBUF in parallel +- Vector Engine and GpSimd Engine cannot access SBUF in parallel -* Vector Engine and Scalar Engine cannot access PSUM in parallel +- Vector Engine and Scalar Engine cannot access PSUM in parallel Despite these restrictions, SBUF is capable of driving peak bandwidth in each tensor read/write interface connected to VectorE/ScalarE/TensorE -or GpSimdE/ScalarE/TensorE *simultaneously* without bandwidth interference. Similarly, PSUM can drive peak bandwidth for -VectorE/TensorE or ScalarE/TensorE *simultaneously*. +or GpSimdE/ScalarE/TensorE _simultaneously_ without bandwidth interference. Similarly, PSUM can drive peak bandwidth for +VectorE/TensorE or ScalarE/TensorE _simultaneously_. **Tensor access overhead.** Initiating a tensor access request from an engine to its SBUF/PSUM read/write interface incurs a static overhead approximately 60 cycles on NeuronCore-v2. Compute engines can typically hide some of this latency through @@ -1184,35 +1197,34 @@ whenever possible to amortize this overhead. As shown in [Figure 48](#fig-arch-neuron-core-v2), both VectorE and ScalarE have read and write access to PSUM, while TensorE only has write access. In fact, PSUM is designed to be a landing buffer for TensorE with near-memory accumulation capabilities that allows read-accumulate-write to every -4B element in memory. Note, this accumulation mechanism can *only* be controlled by TensorE. VectorE and ScalarE can only +4B element in memory. Note, this accumulation mechanism can _only_ be controlled by TensorE. VectorE and ScalarE can only access PSUM like a regular SRAM similar to SBUF. -Next, let’s discuss how TensorE can write outputs to PSUM. As previously discussed, PSUM is organized into 128 *partitions,* +Next, let’s discuss how TensorE can write outputs to PSUM. As previously discussed, PSUM is organized into 128 _partitions,_ each consisting of 16KB of memory. Each partition is further divided into 8 PSUM banks, with each bank holding up to 512 32-bit values. The output tile of a TensorE matrix multiplication instruction (`nki.isa.nc_matmul`) must **fit** into one PSUM bank per partition, which is the fundamental reason for the [free dimension size limitation](#arch-matmul-tile-size) for the `moving` tensor. -Every `nc_matmul` instruction can choose whether to *override* existing bank data with instruction output or *accumulate* +Every `nc_matmul` instruction can choose whether to _override_ existing bank data with instruction output or _accumulate_ instruction output into existing bank data element-wise. The accumulation mode of PSUM is particularly useful when the high-level matmul operator has a contraction dimension (i.e., `stationary/moving` partition dimension of `nki.isa.nc_matmul`) greater than 128. As an example, let’s assume the following matmul dimensions: -* `x.shape = [128, 256]` +- `x.shape = [128, 256]` -* `y.shape = [256, 512]` +- `y.shape = [256, 512]` Figure below shows this matmul mathematically and also how we would tile the contraction dimension. With tiling, we slice both `x` and `y` in the contraction dimension to get `[x0, x1]` and `[y0, y1]` input tiles. To get the final output result, we need to perform: -* output0 = matmul(x0, y0) - -* output1 = matmul(x1, y1) +- output0 = matmul(x0, y0) -* output = output0 + output1 +- output1 = matmul(x1, y1) +- output = output0 + output1 > **Figure: mm tiling** > @@ -1222,6 +1234,7 @@ final output result, we need to perform: > > **Part (a) Mathematical View** (top section): > Shows the full matrix multiplication setup: +> > - A blue matrix "y" at top, divided into y_0 and y_1 by a red dashed line, with dimensions 512 (width) by 256 (height) > - A green matrix "x" divided into x_0 and x_1, with dimensions 256 (width) by 128 (height) > - A purple "output" matrix with dimensions 512 (width) by 128 (height) @@ -1231,22 +1244,26 @@ final output result, we need to perform: > Shows the sequential computation process: > > **Step 1**: +> > - Blue tile y_0 (512 x 128) with red dashed border indicating current tile > - Green tile x_0 (128 x 128) > - Purple output_0 (512 x 128) - first partial result > > **Step 2**: +> > - Blue tile y_1 (512 x 128) with red dashed border > - Green tile x_1 (128 x 128) > - Purple output_1 (512 x 128) - second partial result > > **Step 3**: +> > - Shows output_0 + output_1 = output > - Purple tiles being summed to produce final result > > The red dashed borders indicate which tiles are currently being processed in each step. > > **Key Elements:** +> > - **y_0, y_1**: Tiles of the y matrix (blue) > - **x_0, x_1**: Tiles of the x matrix (green) > - **output_0, output_1**: Partial output tiles (purple) @@ -1255,14 +1272,12 @@ final output result, we need to perform: > - **Step 1, 2, 3**: Sequential computation phases > - **Plus sign (+)**: Accumulation of partial results - Fig. 63 Matmul tiling (mathematical view). PSUM accumulation effectively combines Step 2 and 3 above into a single TensorE `nki.isa.nc_matmul` instruction. Assuming we have `x` in the transposed layout in SBUF, visually the above tiled matmul example will have two back-to-back `nki.isa.nc_matmul` instructions on TensorE: - > **Figure: mm tiling hw** > > A hardware-level view of matrix multiplication tiling showing two iterations with Tensor Engine, SBUF, and PSUM components, demonstrating the Overwrite (first iteration) and Accumulate (second iteration) operations. @@ -1270,6 +1285,7 @@ instructions on TensorE: > This diagram shows the hardware-level execution of tiled matrix multiplication across two iterations, illustrating how partial results are accumulated in PSUM. > > **Left side (First iteration)**: +> > - **Tensor Engine** contains xT_0 (green tile) with dimensions 128(F) x 128(P), representing the transposed x_0 tile > - **SBUF** contains y_0 (blue tile) with dimensions 512(F) x 128(P), the moving matrix from State Buffer > - Arrow labeled "Overwrite" points down from these inputs @@ -1277,12 +1293,14 @@ instructions on TensorE: > - This is the first partial result, written fresh to PSUM > > **Right side (Second iteration)**: +> > - **Tensor Engine** contains xT_1 (green tile) with dimensions 128(F) x 128(P), the transposed x_1 tile > - **SBUF** contains y_1 (blue tile) with dimensions 512(F) x 128(P), the next moving matrix > - Arrow labeled "Accumulate" points down > - **PSUM** at bottom shows "output_0 + output_1" (purple tile), indicating accumulation of partial results into the same PSUM location > > Dimension annotations throughout: +> > - 128 (F): Free dimension size (128 elements) > - 128 (P): Partition dimension size (128 partitions) > - 512 (F): Larger free dimension for the y/output tiles @@ -1290,6 +1308,7 @@ instructions on TensorE: > Red dashed borders on tiles indicate the current data being processed. > > **Key Elements:** +> > - **Tensor Engine**: Holds stationary matrix (xT_0, xT_1) > - **SBUF**: State Buffer holding moving matrix (y_0, y_1) > - **PSUM**: Partial Sum buffer for output accumulation @@ -1298,19 +1317,17 @@ instructions on TensorE: > - **128(F), 128(P), 512(F)**: Dimension annotations > - **output_0 + output_1**: Shows partial result accumulation - Fig. 64 Matmul tiling (hardware view). Effectively, the first `nki.isa.nc_matmul` instruction overwrites the destination PSUM bank with the instruction output. The second instruction accumulates instruction output onto the previous instruction’s result in the same PSUM. The PSUM accumulation is always done in FP32. A series of TensorE matmul instructions with the first one writing to a PSUM bank and -more subsequent instructions accumulating into the same PSUM bank data is called a *matmul accumulation group*. +more subsequent instructions accumulating into the same PSUM bank data is called a _matmul accumulation group_. In current release of NKI, the `nki.isa.nc_matmul` does not have an explicit control field to indicate `overwrite` or `accumulate` for the PSUM. Instead, NeuronCompiler relies on the following NKI code pattern to trigger PSUM accumulation: - ```python # condition 1: a psum buffer with zeros psum_buf = nl.zeros(..., buffer=nl.psum) @@ -1321,7 +1338,6 @@ for i in range(N): psum_buf += nl.matmul(stationary_tile, moving_tile) # or nisa.nc_matmul ``` - Refer to the [Tiling Matrix Multiplications](../programming/tutorials/matrix_multiplication.md#tutorial-matmul-tiling) tutorial for a detailed implementation. @@ -1329,8 +1345,7 @@ tutorial for a detailed implementation. > **Note** > > Note -> -> +> > Due to current limitations in NKI, `psum_buf[...] = psum_buf + nisa.nc_matmul(stationary_tile, moving_tile)` > will not reliably trigger the PSUM accumulation architecture feature. Therefore, even though this alternative > syntax is functionally equivalent to the use of `+=`, it may get lowered to nisa.tensor_tensor on VectorEngine for @@ -1339,4 +1354,4 @@ tutorial for a detailed implementation. Finally, with 8 PSUM banks per partition, TensorE can have up to eight outstanding matmul accumulation groups, which allows flexible scheduling of matmul instructions on TensorE. Also, the extra buffering from multiple PSUM banks allows us to pipeline TensorE computation with other compute engines: TensorE can move onto the next accumulation group without waiting for VectorE/ScalarE -to evict previous accumulation group results. \ No newline at end of file +to evict previous accumulation group results. diff --git a/skills/neuron-nki-docs/references/debugging/error-codes-index.md b/skills/neuron-nki-docs/references/debugging/error-codes-index.md index 36a74ea..c49fcc8 100644 --- a/skills/neuron-nki-docs/references/debugging/error-codes-index.md +++ b/skills/neuron-nki-docs/references/debugging/error-codes-index.md @@ -5,37 +5,35 @@ Neuron Compiler Error Codes This page lists the error codes you can encounter while developing with the Neuron Compiler. For more details on any individual error, click the link for that error code in the table below. +| Error Code | Error Message | Recommendation | +| ---------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------- | +| [ NCC_EARG001 ](error-codes/EARG001.md#error-code-earg001) | Unsupported Logical Neuron Core (LNC) configuration. | You attempted to use a Logical Neuron Core configuration that is not supported by the target Neuron architecture. | +| [ NCC_EBVF030 ](error-codes/EBVF030.md#error-code-ebvf030) | The number of instructions generated exceeds the limit. | Consider applying model parallelism as partitioning the model will help break large computational graphs into smaller subgraphs. | +| [ NCC_EHCA005 ](error-codes/EHCA005.md#error-code-ehca005) | The compiler encountered a custom call instruction with a target name that is not recognized. | Use a supported custom call target from the list of recognized targets. | +| [ NCC_EOOM001 ](error-codes/EOOM001.md#error-code-eoom001) | The combined memory needed for the model’s activation tensors exceeds the high-bandwidth memory limit. | You may need to reduce batch/tensor size or utilize pipeline/tensor parallelism via neuronx-distributed. | +| [ NCC_EOOM002 ](error-codes/EOOM002.md#error-code-eoom002) | The combined memory needed for the model’s activation tensors exceeds the high-bandwidth memory limit. | You may need to reduce batch/tensor size or utilize pipeline/tensor parallelism via neuronx-distributed. | +| [ NCC_ESFH002 ](error-codes/ESFH002.md#error-code-esfh002) | The compiler encountered a unsigned 64-bit integer constant with a value that cannot be safely converted to 32-bit representation. | Try to use uint32 for constants when possible and restructure code to avoid large constants. | +| [ NCC_ESPP004 ](error-codes/ESPP004.md#error-code-espp004) | The compiler encountered a data type that is not supported for code generation. | Use a supported data type as listed in the Neuron documentation. | +| [ NCC_ESPP047 ](error-codes/ESPP047.md#error-code-espp047) | Unsupported 8-bit floating-point data type. | The compiler found usage of an unsupported 8-bit floating-point data type. Convert to a supported type like torch.float16. | +| [ NCC_EUOC002 ](error-codes/EUOC002.md#error-code-euoc002) | An unsupported operator was used. | Try using alternative operators from the full list of supported operators via neuronx-cc list-operators –framework XLA to workaround the limitation. | +| [ NCC_EVRF001 ](error-codes/EVRF001.md#error-code-evrf001) | An unsupported operator was used. | Try using alternative operators from the full list of supported operators to workaround the limitation. | +| [ NCC_EVRF004 ](error-codes/EVRF004.md#error-code-evrf004) | Complex data types are not supported on the Neuron device. | You cannot use complex data types (such as complex64 , complex128 , and others) on the Neuron device directly. | +| [ NCC_EVRF005 ](error-codes/EVRF005.md#error-code-evrf005) | Unsupported F8E4M3FNUZ, F8E4M3B11FNUZ, or F8E5M2FNUZ data type. | The compiler found usage of unsupported 8-bit floating-point data types. Convert to a supported type like torch.float16. | +| [ NCC_EVRF006 ](error-codes/EVRF006.md#error-code-evrf006) | The compiler encountered a RNGBitGenerator operation using a random number generation algorithm other than RNG_DEFAULT. | Ensure that you are using standard JAX/PyTorch random APIs and not explicity specifying an RNG algorithm. | +| [ NCC_EVRF007 ](error-codes/EVRF007.md#error-code-evrf007) | The number of instructions generated exceeds the limit. | Consider applying model parallelism as partitioning the model will help break large computational graphs into smaller subgraphs. | +| [ NCC_EVRF009 ](error-codes/EVRF009.md#error-code-evrf009) | The combined memory needed for the model’s activation tensors exceeds the high-bandwidth memory limit. | You may need to reduce batch/tensor size or utilize pipeline/tensor parallelism via neuronx-distributed. | +| [ NCC_EVRF010 ](error-codes/EVRF010.md#error-code-evrf010) | The compiler encountered simultaneous use of input and kernel dilation, which is not supported. | If possible, use only input or kernel dilation, not both simultaneously. | +| [ NCC_EVRF011 ](error-codes/EVRF011.md#error-code-evrf011) | The compiler encountered strided convolution combined with dilated input, which is not supported. | If possible, remove stride or input dilation, or apply upsampling and downsampling separately. | +| [ NCC_EVRF013 ](error-codes/EVRF013.md#error-code-evrf013) | TopK does not support integer input tensors (int32, int64). | The TopK operation cannot be performed on integer data types. | +| [ NCC_EVRF015 ](error-codes/EVRF015.md#error-code-evrf015) | The compiler encountered a custom call instruction with a target name that is not recognized. | Use a supported custom call target from the list of recognized targets. | +| [ NCC_EVRF016 ](error-codes/EVRF016.md#error-code-evr016) | The scatter-reduce operation cannot perform reduction logic if the data being scattered or the destination tensor is using an integer or boolean data type. | Cast your input and source tensors to a floating-point data type (e.g., torch.float32 or torch.bfloat16). | +| [ NCC_EVRF017 ](error-codes/EVRF017.md#error-code-evrf017) | Reduce-window operation with base dilation greater than 1 is not supported. | Change base dilation to be all 1s or consider manual dilation if necessary. | +| [ NCC_EVRF018 ](error-codes/EVRF018.md#error-code-evrf018) | Reduce-window operation with window dilation greater than 1 is not supported. | Remove window_dilation or change values to be all 1s, or consider manual dilation if necessary. | +| [ NCC_EVRF019 ](error-codes/EVRF019.md#error-code-evrf019) | The compiler encountered a reduce-window operation with more or less than 2 operands. | If possible, split multi-operand reduce_window with multiple single-operand reduce_window operations. | +| [ NCC_EVRF022 ](error-codes/EVRF022.md#error-code-evrf022) | Shift-right-arithmetic operation on non 32-bit inputs is not supported. Cast the first argument’s data type to be S32, U32, or F32. | You need to use 32-bit data types for shift operations. Cast inputs to int32, uint32, or float32. | +| [ NCC_EVRF024 ](error-codes/EVRF024.md#error-code-evrf024) | The output tensor size limit of 4GB was exceeded. | Reduce batch/tensor size or utilize tensor parallelism via neuronx-distributed. | +| [ NCC_EVRF031 ](error-codes/EVRF031.md#error-code-evrf031) | The compiler encountered a scatter out-of-bounds error. | Ensure that the iota size matches the operand dimension size. | +| [ NCC_EXSP001 ](error-codes/EXSP001.md#error-code-exsp001) | The combined memory needed for the model’s activation tensors exceeds the high-bandwidth memory limit. | You may need to reduce batch/tensor size or utilize pipeline/tensor parallelism via neuronx-distributed. | +| [ NCC_EXTP004 ](error-codes/EXTP004.md#error-code-extp004) | The number of instructions generated exceeds the limit. | Consider applying model parallelism as partitioning the model will help break large computational graphs into smaller subgraphs. | -| Error Code | Error Message | Recommendation | -| --- | --- | --- | -| [ NCC_EARG001 ](error-codes/EARG001.md#error-code-earg001) | Unsupported Logical Neuron Core (LNC) configuration. | You attempted to use a Logical Neuron Core configuration that is not supported by the target Neuron architecture. | -| [ NCC_EBVF030 ](error-codes/EBVF030.md#error-code-ebvf030) | The number of instructions generated exceeds the limit. | Consider applying model parallelism as partitioning the model will help break large computational graphs into smaller subgraphs. | -| [ NCC_EHCA005 ](error-codes/EHCA005.md#error-code-ehca005) | The compiler encountered a custom call instruction with a target name that is not recognized. | Use a supported custom call target from the list of recognized targets. | -| [ NCC_EOOM001 ](error-codes/EOOM001.md#error-code-eoom001) | The combined memory needed for the model’s activation tensors exceeds the high-bandwidth memory limit. | You may need to reduce batch/tensor size or utilize pipeline/tensor parallelism via neuronx-distributed. | -| [ NCC_EOOM002 ](error-codes/EOOM002.md#error-code-eoom002) | The combined memory needed for the model’s activation tensors exceeds the high-bandwidth memory limit. | You may need to reduce batch/tensor size or utilize pipeline/tensor parallelism via neuronx-distributed. | -| [ NCC_ESFH002 ](error-codes/ESFH002.md#error-code-esfh002) | The compiler encountered a unsigned 64-bit integer constant with a value that cannot be safely converted to 32-bit representation. | Try to use uint32 for constants when possible and restructure code to avoid large constants. | -| [ NCC_ESPP004 ](error-codes/ESPP004.md#error-code-espp004) | The compiler encountered a data type that is not supported for code generation. | Use a supported data type as listed in the Neuron documentation. | -| [ NCC_ESPP047 ](error-codes/ESPP047.md#error-code-espp047) | Unsupported 8-bit floating-point data type. | The compiler found usage of an unsupported 8-bit floating-point data type. Convert to a supported type like torch.float16. | -| [ NCC_EUOC002 ](error-codes/EUOC002.md#error-code-euoc002) | An unsupported operator was used. | Try using alternative operators from the full list of supported operators via neuronx-cc list-operators –framework XLA to workaround the limitation. | -| [ NCC_EVRF001 ](error-codes/EVRF001.md#error-code-evrf001) | An unsupported operator was used. | Try using alternative operators from the full list of supported operators to workaround the limitation. | -| [ NCC_EVRF004 ](error-codes/EVRF004.md#error-code-evrf004) | Complex data types are not supported on the Neuron device. | You cannot use complex data types (such as complex64 , complex128 , and others) on the Neuron device directly. | -| [ NCC_EVRF005 ](error-codes/EVRF005.md#error-code-evrf005) | Unsupported F8E4M3FNUZ, F8E4M3B11FNUZ, or F8E5M2FNUZ data type. | The compiler found usage of unsupported 8-bit floating-point data types. Convert to a supported type like torch.float16. | -| [ NCC_EVRF006 ](error-codes/EVRF006.md#error-code-evrf006) | The compiler encountered a RNGBitGenerator operation using a random number generation algorithm other than RNG_DEFAULT. | Ensure that you are using standard JAX/PyTorch random APIs and not explicity specifying an RNG algorithm. | -| [ NCC_EVRF007 ](error-codes/EVRF007.md#error-code-evrf007) | The number of instructions generated exceeds the limit. | Consider applying model parallelism as partitioning the model will help break large computational graphs into smaller subgraphs. | -| [ NCC_EVRF009 ](error-codes/EVRF009.md#error-code-evrf009) | The combined memory needed for the model’s activation tensors exceeds the high-bandwidth memory limit. | You may need to reduce batch/tensor size or utilize pipeline/tensor parallelism via neuronx-distributed. | -| [ NCC_EVRF010 ](error-codes/EVRF010.md#error-code-evrf010) | The compiler encountered simultaneous use of input and kernel dilation, which is not supported. | If possible, use only input or kernel dilation, not both simultaneously. | -| [ NCC_EVRF011 ](error-codes/EVRF011.md#error-code-evrf011) | The compiler encountered strided convolution combined with dilated input, which is not supported. | If possible, remove stride or input dilation, or apply upsampling and downsampling separately. | -| [ NCC_EVRF013 ](error-codes/EVRF013.md#error-code-evrf013) | TopK does not support integer input tensors (int32, int64). | The TopK operation cannot be performed on integer data types. | -| [ NCC_EVRF015 ](error-codes/EVRF015.md#error-code-evrf015) | The compiler encountered a custom call instruction with a target name that is not recognized. | Use a supported custom call target from the list of recognized targets. | -| [ NCC_EVRF016 ](error-codes/EVRF016.md#error-code-evr016) | The scatter-reduce operation cannot perform reduction logic if the data being scattered or the destination tensor is using an integer or boolean data type. | Cast your input and source tensors to a floating-point data type (e.g., torch.float32 or torch.bfloat16). | -| [ NCC_EVRF017 ](error-codes/EVRF017.md#error-code-evrf017) | Reduce-window operation with base dilation greater than 1 is not supported. | Change base dilation to be all 1s or consider manual dilation if necessary. | -| [ NCC_EVRF018 ](error-codes/EVRF018.md#error-code-evrf018) | Reduce-window operation with window dilation greater than 1 is not supported. | Remove window_dilation or change values to be all 1s, or consider manual dilation if necessary. | -| [ NCC_EVRF019 ](error-codes/EVRF019.md#error-code-evrf019) | The compiler encountered a reduce-window operation with more or less than 2 operands. | If possible, split multi-operand reduce_window with multiple single-operand reduce_window operations. | -| [ NCC_EVRF022 ](error-codes/EVRF022.md#error-code-evrf022) | Shift-right-arithmetic operation on non 32-bit inputs is not supported. Cast the first argument’s data type to be S32, U32, or F32. | You need to use 32-bit data types for shift operations. Cast inputs to int32, uint32, or float32. | -| [ NCC_EVRF024 ](error-codes/EVRF024.md#error-code-evrf024) | The output tensor size limit of 4GB was exceeded. | Reduce batch/tensor size or utilize tensor parallelism via neuronx-distributed. | -| [ NCC_EVRF031 ](error-codes/EVRF031.md#error-code-evrf031) | The compiler encountered a scatter out-of-bounds error. | Ensure that the iota size matches the operand dimension size. | -| [ NCC_EXSP001 ](error-codes/EXSP001.md#error-code-exsp001) | The combined memory needed for the model’s activation tensors exceeds the high-bandwidth memory limit. | You may need to reduce batch/tensor size or utilize pipeline/tensor parallelism via neuronx-distributed. | -| [ NCC_EXTP004 ](error-codes/EXTP004.md#error-code-extp004) | The number of instructions generated exceeds the limit. | Consider applying model parallelism as partitioning the model will help break large computational graphs into smaller subgraphs. | - - -**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` \ No newline at end of file +**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` diff --git a/skills/neuron-nki-docs/references/debugging/error-codes/EARG001.md b/skills/neuron-nki-docs/references/debugging/error-codes/EARG001.md index 55fc36e..ee3f606 100644 --- a/skills/neuron-nki-docs/references/debugging/error-codes/EARG001.md +++ b/skills/neuron-nki-docs/references/debugging/error-codes/EARG001.md @@ -7,7 +7,6 @@ NCC_EARG001 For example, a trn1 instance running the following code will run into this error: - ```python traced_model = torch_neuronx.trace( model, @@ -16,21 +15,20 @@ traced_model = torch_neuronx.trace( ) ``` - On trn1, only lnc=1 is supported. Physical Neuron Core: -* Actual hardware compute unit on the chip +- Actual hardware compute unit on the chip -* Has dedicated compute resources, memory, etc. +- Has dedicated compute resources, memory, etc. Logical Neuron Core: -* Software abstraction grouping multiple physical cores +- Software abstraction grouping multiple physical cores -* Controlled via the NEURON_LOGICAL_NC_CONFIG environment variable or the –lnc flag (when using neuronx-cc directly) +- Controlled via the NEURON_LOGICAL_NC_CONFIG environment variable or the –lnc flag (when using neuronx-cc directly) For more information: [https://awsdocs-neuron.readthedocs-hosted.com/en/latest/neuron-runtime/explore/device-memory.html#logical-neuron-cores](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/neuron-runtime/explore/device-memory.html#logical-neuron-cores) -**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` \ No newline at end of file +**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` diff --git a/skills/neuron-nki-docs/references/debugging/error-codes/EBVF030.md b/skills/neuron-nki-docs/references/debugging/error-codes/EBVF030.md index f69b063..d641dd7 100644 --- a/skills/neuron-nki-docs/references/debugging/error-codes/EBVF030.md +++ b/skills/neuron-nki-docs/references/debugging/error-codes/EBVF030.md @@ -9,8 +9,8 @@ Consider applying model parallelism as partitioning the model will help break la For more information: -* [https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/neuronx-distributed/api_guide.html#api-guide](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/neuronx-distributed/api_guide.html#api-guide) +- [https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/neuronx-distributed/api_guide.html#api-guide](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/neuronx-distributed/api_guide.html#api-guide) -* [https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/nxd-training/app_notes/nxd-training-pp-appnote.html](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/nxd-training/app_notes/nxd-training-pp-appnote.html) +- [https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/nxd-training/app_notes/nxd-training-pp-appnote.html](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/nxd-training/app_notes/nxd-training-pp-appnote.html) -**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` \ No newline at end of file +**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` diff --git a/skills/neuron-nki-docs/references/debugging/error-codes/EHCA005.md b/skills/neuron-nki-docs/references/debugging/error-codes/EHCA005.md index 4318558..23a2ed3 100644 --- a/skills/neuron-nki-docs/references/debugging/error-codes/EHCA005.md +++ b/skills/neuron-nki-docs/references/debugging/error-codes/EHCA005.md @@ -7,67 +7,66 @@ NCC_EHCA005 The Neuron compiler currently recognizes the following custom call targets: -* AwsNeuronErf +- AwsNeuronErf -* AwsNeuronGelu +- AwsNeuronGelu -* AwsNeuronGeluApprxTanh +- AwsNeuronGeluApprxTanh -* AwsNeuronGeluBackward +- AwsNeuronGeluBackward -* AwsNeuronSilu +- AwsNeuronSilu -* AwsNeuronSiluBackward +- AwsNeuronSiluBackward -* AwsNeuronRmsNorm +- AwsNeuronRmsNorm -* AwsNeuronSoftmax +- AwsNeuronSoftmax -* AwsNeuronSoftmaxBackward +- AwsNeuronSoftmaxBackward -* AwsNeuronCollectiveMatmul +- AwsNeuronCollectiveMatmul -* AwsNeuronIntMatmult +- AwsNeuronIntMatmult -* AwsNeuronArgMax +- AwsNeuronArgMax -* AwsNeuronArgMin +- AwsNeuronArgMin -* AwsNeuronTopK +- AwsNeuronTopK -* AwsNeuronDropoutMaskV1 +- AwsNeuronDropoutMaskV1 -* AwsNeuronCustomNativeKernel +- AwsNeuronCustomNativeKernel -* AwsNeuronCustomOp +- AwsNeuronCustomOp -* AwsNeuronDevicePrint +- AwsNeuronDevicePrint -* ResizeNearest +- ResizeNearest -* ResizeBilinear +- ResizeBilinear -* ResizeNearestGrad +- ResizeNearestGrad -* AwsNeuronLNCShardingConstraint +- AwsNeuronLNCShardingConstraint -* AwsNeuronTransferWithStaticRing +- AwsNeuronTransferWithStaticRing -* AwsNeuronModuleMarkerStart-Forward +- AwsNeuronModuleMarkerStart-Forward -* AwsNeuronModuleMarkerStart-Backward +- AwsNeuronModuleMarkerStart-Backward -* AwsNeuronModuleMarkerEnd-Forward +- AwsNeuronModuleMarkerEnd-Forward -* AwsNeuronModuleMarkerEnd-Backward +- AwsNeuronModuleMarkerEnd-Backward -* NeuronBoundaryMarker-Start +- NeuronBoundaryMarker-Start -* NeuronBoundaryMarker-End +- NeuronBoundaryMarker-End Erroneous code example: - ```python def lowering(ctx, x_val): result_type = ir.RankedTensorType(x_val.type) @@ -80,10 +79,8 @@ def lowering(ctx, x_val): ).results ``` - Use a supported custom call target: - ```python def lowering(ctx, x_val): result_type = ir.RankedTensorType(x_val.type) @@ -97,5 +94,4 @@ def lowering(ctx, x_val): ).results ``` - -**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` \ No newline at end of file +**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` diff --git a/skills/neuron-nki-docs/references/debugging/error-codes/EOOM001.md b/skills/neuron-nki-docs/references/debugging/error-codes/EOOM001.md index 27f763b..ca32e43 100644 --- a/skills/neuron-nki-docs/references/debugging/error-codes/EOOM001.md +++ b/skills/neuron-nki-docs/references/debugging/error-codes/EOOM001.md @@ -7,21 +7,20 @@ NCC_EOOM001 The memory usage consists of: -* I/O tensors: Input and output activation tensors +- I/O tensors: Input and output activation tensors -* Internal allocations: Scratchpad memory for intermediate computations +- Internal allocations: Scratchpad memory for intermediate computations -* SBUF spills: Data that cannot fit in on-chip SBUF memory and must spill to HBM +- SBUF spills: Data that cannot fit in on-chip SBUF memory and must spill to HBM There are several ways to potentially fix this issue. -* Simply reduce the batch/tensor size if possible +- Simply reduce the batch/tensor size if possible -* Utilize pipeline/tensor parallelism via neuronx-distributed +- Utilize pipeline/tensor parallelism via neuronx-distributed Short snippet of tensor parallelism: - ```python class ParallelSelfAttention(transformers.models.bert.modeling_bert.BertSelfAttention): def __init__(self, config, position_embedding_type=None): @@ -43,11 +42,10 @@ class ParallelSelfAttention(transformers.models.bert.modeling_bert.BertSelfAtten self.all_head_size = self.all_head_size // tp_size ``` - For more information: -* [https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/neuronx-distributed/activation_memory_reduction.html](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/neuronx-distributed/activation_memory_reduction.html) +- [https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/neuronx-distributed/activation_memory_reduction.html](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/neuronx-distributed/activation_memory_reduction.html) -* [https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/nxd-training/app_notes/nxd-training-pp-appnote.html](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/nxd-training/app_notes/nxd-training-pp-appnote.html) +- [https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/nxd-training/app_notes/nxd-training-pp-appnote.html](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/nxd-training/app_notes/nxd-training-pp-appnote.html) -**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` \ No newline at end of file +**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` diff --git a/skills/neuron-nki-docs/references/debugging/error-codes/EOOM002.md b/skills/neuron-nki-docs/references/debugging/error-codes/EOOM002.md index 8f7686d..ed6b96e 100644 --- a/skills/neuron-nki-docs/references/debugging/error-codes/EOOM002.md +++ b/skills/neuron-nki-docs/references/debugging/error-codes/EOOM002.md @@ -7,21 +7,20 @@ NCC_EOOM002 The memory usage consists of: -* I/O tensors: Input and output activation tensors +- I/O tensors: Input and output activation tensors -* Internal allocations: Scratchpad memory for intermediate computations +- Internal allocations: Scratchpad memory for intermediate computations -* SBUF spills: Data that cannot fit in on-chip SBUF memory and must spill to HBM +- SBUF spills: Data that cannot fit in on-chip SBUF memory and must spill to HBM There are several ways to potentially fix this issue. -* Simply reduce the batch/tensor size if possible +- Simply reduce the batch/tensor size if possible -* Utilize pipeline/tensor parallelism via neuronx-distributed +- Utilize pipeline/tensor parallelism via neuronx-distributed Short snippet of tensor parallelism: - ```python class ParallelSelfAttention(transformers.models.bert.modeling_bert.BertSelfAttention): def __init__(self, config, position_embedding_type=None): @@ -43,11 +42,10 @@ class ParallelSelfAttention(transformers.models.bert.modeling_bert.BertSelfAtten self.all_head_size = self.all_head_size // tp_size ``` - For more information: -* [https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/neuronx-distributed/activation_memory_reduction.html](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/neuronx-distributed/activation_memory_reduction.html) +- [https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/neuronx-distributed/activation_memory_reduction.html](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/neuronx-distributed/activation_memory_reduction.html) -* [https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/nxd-training/app_notes/nxd-training-pp-appnote.html](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/nxd-training/app_notes/nxd-training-pp-appnote.html) +- [https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/nxd-training/app_notes/nxd-training-pp-appnote.html](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/nxd-training/app_notes/nxd-training-pp-appnote.html) -**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` \ No newline at end of file +**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` diff --git a/skills/neuron-nki-docs/references/debugging/error-codes/ESFH002.md b/skills/neuron-nki-docs/references/debugging/error-codes/ESFH002.md index 501d9ed..6504885 100644 --- a/skills/neuron-nki-docs/references/debugging/error-codes/ESFH002.md +++ b/skills/neuron-nki-docs/references/debugging/error-codes/ESFH002.md @@ -9,7 +9,6 @@ The Neuron hardware operates on 32-bit or narrower data types and attempts to co Erroneous code example: - ```python @jax.jit def foo(): @@ -20,10 +19,8 @@ def foo(): return x + large_constant ``` - Use uint32 for constants when possible: - ```python @jax.jit def test(): @@ -32,5 +29,4 @@ def test(): return x + large_constant ``` - -**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` \ No newline at end of file +**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` diff --git a/skills/neuron-nki-docs/references/debugging/error-codes/ESPP004.md b/skills/neuron-nki-docs/references/debugging/error-codes/ESPP004.md index f267f57..a6e7d0f 100644 --- a/skills/neuron-nki-docs/references/debugging/error-codes/ESPP004.md +++ b/skills/neuron-nki-docs/references/debugging/error-codes/ESPP004.md @@ -7,7 +7,6 @@ NCC_ESPP004 Erroneous code example: - ```python import numpy as np import jax.numpy as jnp @@ -20,10 +19,8 @@ dtype = np.dtype(dtypes.float4_e2m1fn) val = lax_internal._convert_element_type(0, dtype, weak_type=False) ``` - Use a supported data type: - ```python import numpy as np import jax.numpy as jnp @@ -36,7 +33,6 @@ dtype = jnp.bfloat16 val = lax_internal._convert_element_type(0, dtype, weak_type=False) ``` - More information on supported data types [https://awsdocs-neuron.readthedocs-hosted.com/en/latest/about-neuron/arch/neuron-features/data-types.html](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/about-neuron/arch/neuron-features/data-types.html) -**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` \ No newline at end of file +**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` diff --git a/skills/neuron-nki-docs/references/debugging/error-codes/ESPP047.md b/skills/neuron-nki-docs/references/debugging/error-codes/ESPP047.md index 67bb62b..a9ff0bb 100644 --- a/skills/neuron-nki-docs/references/debugging/error-codes/ESPP047.md +++ b/skills/neuron-nki-docs/references/debugging/error-codes/ESPP047.md @@ -7,7 +7,6 @@ NCC_ESPP047 Erroneous code example: - ```python class Model(nn.Module): def __init__(self): @@ -25,10 +24,8 @@ class Model(nn.Module): input_tensor = torch.randn(1, 10).to(torch.float8_e4m3fn) ``` - To fix this error: - ```python class Model(nn.Module): def __init__(self): @@ -47,5 +44,4 @@ input_tensor = torch.randn(1, 10).to(torch.float8_e4m3fn) input_tensor = input_tensor.to(torch.float16) ``` - -**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` \ No newline at end of file +**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` diff --git a/skills/neuron-nki-docs/references/debugging/error-codes/EUOC002.md b/skills/neuron-nki-docs/references/debugging/error-codes/EUOC002.md index e19a3c5..ca4e64f 100644 --- a/skills/neuron-nki-docs/references/debugging/error-codes/EUOC002.md +++ b/skills/neuron-nki-docs/references/debugging/error-codes/EUOC002.md @@ -9,17 +9,14 @@ Try using alternative operators from the full list of supported operators via ne Before: - ```python class Model(torch.nn.Module): def forward(self, A, b): return torch.triangular_solve(b, A) ``` - Possible workaround: - ```python class Model(torch.nn.Module): def forward(self, A, b): @@ -28,5 +25,4 @@ class Model(torch.nn.Module): return A_inv @ b ``` - -**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` \ No newline at end of file +**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` diff --git a/skills/neuron-nki-docs/references/debugging/error-codes/EVRF001.md b/skills/neuron-nki-docs/references/debugging/error-codes/EVRF001.md index 94c92e8..da6964a 100644 --- a/skills/neuron-nki-docs/references/debugging/error-codes/EVRF001.md +++ b/skills/neuron-nki-docs/references/debugging/error-codes/EVRF001.md @@ -9,17 +9,14 @@ Try using alternative operators from the full list of supported operators via ne Before: - ```python class Model(torch.nn.Module): def forward(self, A, b): return torch.triangular_solve(b, A) ``` - Possible workaround: - ```python class Model(torch.nn.Module): def forward(self, A, b): @@ -28,5 +25,4 @@ class Model(torch.nn.Module): return A_inv @ b ``` - -**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` \ No newline at end of file +**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` diff --git a/skills/neuron-nki-docs/references/debugging/error-codes/EVRF004.md b/skills/neuron-nki-docs/references/debugging/error-codes/EVRF004.md index 19bdb4b..c1278f3 100644 --- a/skills/neuron-nki-docs/references/debugging/error-codes/EVRF004.md +++ b/skills/neuron-nki-docs/references/debugging/error-codes/EVRF004.md @@ -9,22 +9,18 @@ You cannot use complex data types (such as `complex64`, `complex128`, and others One fix is to offload complex operations to CPU, like so: - ```python x = torch.tensor([1+2j, 3+4j], dtype=torch.complex64).to('cpu') ``` - > **Note** > > Note -> -> +> > Since data transfer between CPU and device is expensive, this is best used when complex operations are rare. You can also address this error by manually emulating complex tensors using real and imaginary parts: - ```python real = x.real imag = x.imag @@ -34,5 +30,4 @@ real_out = a_real * b_real - a_imag * b_imag imag_out = a_real * b_imag + a_imag * b_real ``` - -**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` \ No newline at end of file +**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` diff --git a/skills/neuron-nki-docs/references/debugging/error-codes/EVRF005.md b/skills/neuron-nki-docs/references/debugging/error-codes/EVRF005.md index 8decf6a..742f962 100644 --- a/skills/neuron-nki-docs/references/debugging/error-codes/EVRF005.md +++ b/skills/neuron-nki-docs/references/debugging/error-codes/EVRF005.md @@ -7,7 +7,6 @@ NCC_EVRF005 Erroneous code example: - ```python class Model(nn.Module): def __init__(self): @@ -22,10 +21,8 @@ class Model(nn.Module): input_tensor = torch.randn(1, 10).to(torch.float8_e4m3fnuz) ``` - To fix this error: - ```python class Model(nn.Module): def __init__(self): @@ -42,7 +39,6 @@ input_tensor = torch.randn(1, 10).to(torch.float8_e4m3fnuz) input_tensor = input_tensor.to(torch.float16) ``` +- More information on supported data types: [https://awsdocs-neuron.readthedocs-hosted.com/en/latest/about-neuron/arch/neuron-features/data-types.html](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/about-neuron/arch/neuron-features/data-types.html) -* More information on supported data types: [https://awsdocs-neuron.readthedocs-hosted.com/en/latest/about-neuron/arch/neuron-features/data-types.html](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/about-neuron/arch/neuron-features/data-types.html) - -**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` \ No newline at end of file +**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` diff --git a/skills/neuron-nki-docs/references/debugging/error-codes/EVRF006.md b/skills/neuron-nki-docs/references/debugging/error-codes/EVRF006.md index f5ce14f..9895ce7 100644 --- a/skills/neuron-nki-docs/references/debugging/error-codes/EVRF006.md +++ b/skills/neuron-nki-docs/references/debugging/error-codes/EVRF006.md @@ -8,4 +8,4 @@ NCC_EVRF006 Ensure that you are using standard JAX/PyTorch random APIs and not explicity specifying an RNG algorithm. -**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` \ No newline at end of file +**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` diff --git a/skills/neuron-nki-docs/references/debugging/error-codes/EVRF007.md b/skills/neuron-nki-docs/references/debugging/error-codes/EVRF007.md index 93f1ebf..b2d1aa8 100644 --- a/skills/neuron-nki-docs/references/debugging/error-codes/EVRF007.md +++ b/skills/neuron-nki-docs/references/debugging/error-codes/EVRF007.md @@ -9,8 +9,8 @@ Consider applying model parallelism as partitioning the model will help break la For more information: -* [https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/neuronx-distributed/api_guide.html#api-guide](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/neuronx-distributed/api_guide.html#api-guide) +- [https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/neuronx-distributed/api_guide.html#api-guide](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/neuronx-distributed/api_guide.html#api-guide) -* [https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/nxd-training/app_notes/nxd-training-pp-appnote.html](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/nxd-training/app_notes/nxd-training-pp-appnote.html) +- [https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/nxd-training/app_notes/nxd-training-pp-appnote.html](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/nxd-training/app_notes/nxd-training-pp-appnote.html) -**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` \ No newline at end of file +**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` diff --git a/skills/neuron-nki-docs/references/debugging/error-codes/EVRF009.md b/skills/neuron-nki-docs/references/debugging/error-codes/EVRF009.md index 65cb1e2..55fc7e7 100644 --- a/skills/neuron-nki-docs/references/debugging/error-codes/EVRF009.md +++ b/skills/neuron-nki-docs/references/debugging/error-codes/EVRF009.md @@ -7,13 +7,12 @@ NCC_EVRF009 There are several ways to potentially fix this issue. -* Simply reduce the batch/tensor size if possible +- Simply reduce the batch/tensor size if possible -* Utilize pipeline/tensor parallelism via neuronx-distributed +- Utilize pipeline/tensor parallelism via neuronx-distributed Short snippet of tensor parallelism: - ```python class ParallelSelfAttention(transformers.models.bert.modeling_bert.BertSelfAttention): def __init__(self, config, position_embedding_type=None): @@ -36,11 +35,10 @@ class ParallelSelfAttention(transformers.models.bert.modeling_bert.BertSelfAtten self.all_head_size = self.all_head_size // tp_size ``` - For more information: -* [https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/neuronx-distributed/activation_memory_reduction.html](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/neuronx-distributed/activation_memory_reduction.html) +- [https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/neuronx-distributed/activation_memory_reduction.html](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/neuronx-distributed/activation_memory_reduction.html) -* [https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/nxd-training/app_notes/nxd-training-pp-appnote.html](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/nxd-training/app_notes/nxd-training-pp-appnote.html) +- [https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/nxd-training/app_notes/nxd-training-pp-appnote.html](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/nxd-training/app_notes/nxd-training-pp-appnote.html) -**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` \ No newline at end of file +**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` diff --git a/skills/neuron-nki-docs/references/debugging/error-codes/EVRF010.md b/skills/neuron-nki-docs/references/debugging/error-codes/EVRF010.md index 920b027..d54ec22 100644 --- a/skills/neuron-nki-docs/references/debugging/error-codes/EVRF010.md +++ b/skills/neuron-nki-docs/references/debugging/error-codes/EVRF010.md @@ -7,7 +7,6 @@ NCC_EVRF010 Erroneous code example: - ```python x = jnp.ones((1, 4, 4, 1), dtype=jnp.float32) kernel = jnp.ones((3, 3, 1, 1), dtype=jnp.float32) @@ -23,10 +22,8 @@ result = lax.conv_general_dilated( ) ``` - If possible, use only only input or kernel dilation: - ```python x = jnp.ones((1, 4, 4, 1), dtype=jnp.float32) kernel = jnp.ones((3, 3, 1, 1), dtype=jnp.float32) @@ -42,7 +39,6 @@ result = lax.conv_general_dilated( ) ``` - Or apply dilation manually and apply convolution to the remainder. -**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` \ No newline at end of file +**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` diff --git a/skills/neuron-nki-docs/references/debugging/error-codes/EVRF011.md b/skills/neuron-nki-docs/references/debugging/error-codes/EVRF011.md index 06df921..506d702 100644 --- a/skills/neuron-nki-docs/references/debugging/error-codes/EVRF011.md +++ b/skills/neuron-nki-docs/references/debugging/error-codes/EVRF011.md @@ -7,7 +7,6 @@ NCC_EVRF011 Erroneous code example: - ```python x = jnp.ones((1, 4, 4, 1), dtype=jnp.float32) kernel = jnp.ones((3, 3, 1, 1), dtype=jnp.float32) @@ -23,10 +22,8 @@ result = lax.conv_general_dilated( ) ``` - If possible, remove stride or input dilation: - ```python x = jnp.ones((1, 4, 4, 1), dtype=jnp.float32) kernel = jnp.ones((3, 3, 1, 1), dtype=jnp.float32) @@ -41,7 +38,6 @@ result = lax.conv_general_dilated( ) ``` - Or apply upsampling and downsampling separately. -**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` \ No newline at end of file +**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` diff --git a/skills/neuron-nki-docs/references/debugging/error-codes/EVRF013.md b/skills/neuron-nki-docs/references/debugging/error-codes/EVRF013.md index c6f2bd7..ec82bfc 100644 --- a/skills/neuron-nki-docs/references/debugging/error-codes/EVRF013.md +++ b/skills/neuron-nki-docs/references/debugging/error-codes/EVRF013.md @@ -7,7 +7,6 @@ NCC_EVRF013 Erroneous code example: - ```python def forward(self, x): # assume x is an integer tensor @@ -17,10 +16,8 @@ def forward(self, x): return values, indices ``` - To fix this error, you can cast your tensor to a supported floating point dtype. - ```python def forward(self, x): x = x.float() @@ -29,5 +26,4 @@ def forward(self, x): return values, indices ``` - -**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` \ No newline at end of file +**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` diff --git a/skills/neuron-nki-docs/references/debugging/error-codes/EVRF015.md b/skills/neuron-nki-docs/references/debugging/error-codes/EVRF015.md index a6298ab..19160e4 100644 --- a/skills/neuron-nki-docs/references/debugging/error-codes/EVRF015.md +++ b/skills/neuron-nki-docs/references/debugging/error-codes/EVRF015.md @@ -7,67 +7,66 @@ NCC_EVRF015 The Neuron compiler currently recognizes the following custom call targets: -* AwsNeuronErf +- AwsNeuronErf -* AwsNeuronGelu +- AwsNeuronGelu -* AwsNeuronGeluApprxTanh +- AwsNeuronGeluApprxTanh -* AwsNeuronGeluBackward +- AwsNeuronGeluBackward -* AwsNeuronSilu +- AwsNeuronSilu -* AwsNeuronSiluBackward +- AwsNeuronSiluBackward -* AwsNeuronRmsNorm +- AwsNeuronRmsNorm -* AwsNeuronSoftmax +- AwsNeuronSoftmax -* AwsNeuronSoftmaxBackward +- AwsNeuronSoftmaxBackward -* AwsNeuronCollectiveMatmul +- AwsNeuronCollectiveMatmul -* AwsNeuronIntMatmult +- AwsNeuronIntMatmult -* AwsNeuronArgMax +- AwsNeuronArgMax -* AwsNeuronArgMin +- AwsNeuronArgMin -* AwsNeuronTopK +- AwsNeuronTopK -* AwsNeuronDropoutMaskV1 +- AwsNeuronDropoutMaskV1 -* AwsNeuronCustomNativeKernel +- AwsNeuronCustomNativeKernel -* AwsNeuronCustomOp +- AwsNeuronCustomOp -* AwsNeuronDevicePrint +- AwsNeuronDevicePrint -* ResizeNearest +- ResizeNearest -* ResizeBilinear +- ResizeBilinear -* ResizeNearestGrad +- ResizeNearestGrad -* AwsNeuronLNCShardingConstraint +- AwsNeuronLNCShardingConstraint -* AwsNeuronTransferWithStaticRing +- AwsNeuronTransferWithStaticRing -* AwsNeuronModuleMarkerStart-Forward +- AwsNeuronModuleMarkerStart-Forward -* AwsNeuronModuleMarkerStart-Backward +- AwsNeuronModuleMarkerStart-Backward -* AwsNeuronModuleMarkerEnd-Forward +- AwsNeuronModuleMarkerEnd-Forward -* AwsNeuronModuleMarkerEnd-Backward +- AwsNeuronModuleMarkerEnd-Backward -* NeuronBoundaryMarker-Start +- NeuronBoundaryMarker-Start -* NeuronBoundaryMarker-End +- NeuronBoundaryMarker-End Erroneous code example: - ```python def lowering(ctx, x_val): result_type = ir.RankedTensorType(x_val.type) @@ -80,10 +79,8 @@ def lowering(ctx, x_val): ).results ``` - Use a supported custom call target: - ```python def lowering(ctx, x_val): result_type = ir.RankedTensorType(x_val.type) @@ -97,5 +94,4 @@ def lowering(ctx, x_val): ).results ``` - -**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` \ No newline at end of file +**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` diff --git a/skills/neuron-nki-docs/references/debugging/error-codes/EVRF016.md b/skills/neuron-nki-docs/references/debugging/error-codes/EVRF016.md index fda10a1..2fb2b76 100644 --- a/skills/neuron-nki-docs/references/debugging/error-codes/EVRF016.md +++ b/skills/neuron-nki-docs/references/debugging/error-codes/EVRF016.md @@ -13,7 +13,6 @@ The hardware instructions used on the Neuron device for these specific scatter-a The following example shows the **NCC_EVRF016** error because the `input_tensor` is defined using an integer data type (`torch.int32`) while being used with a reduction function (`reduce='sum'`) in the `scatter_reduce_` operation. - ```python def forward(self, input_tensor, indices_tensor, src_tensor): output = input_tensor.clone() @@ -31,12 +30,10 @@ input_tensor = torch.zeros(BATCH_SIZE, DIM_SIZE, dtype=torch.int32) ... ``` - **How to fix** To fix this error, you must cast your input and source tensors to a floating-point data type (e.g., torch.float32 or torch.bfloat16). - ```python def forward(self, input_tensor, indices_tensor, src_tensor): output = input_tensor.clone() @@ -55,5 +52,4 @@ input_tensor = torch.zeros(BATCH_SIZE, DIM_SIZE, dtype=torch.float32) ... ``` - -**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` \ No newline at end of file +**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` diff --git a/skills/neuron-nki-docs/references/debugging/error-codes/EVRF017.md b/skills/neuron-nki-docs/references/debugging/error-codes/EVRF017.md index 73f9fa2..8742c62 100644 --- a/skills/neuron-nki-docs/references/debugging/error-codes/EVRF017.md +++ b/skills/neuron-nki-docs/references/debugging/error-codes/EVRF017.md @@ -7,7 +7,6 @@ NCC_EVRF017 Erroneous code example: - ```python result = lax.reduce_window( x, -jnp.inf, lax.max, @@ -18,10 +17,8 @@ result = lax.reduce_window( ) ``` - If possible, change base dilation to be all 1s: - ```python result = lax.reduce_window( x, -jnp.inf, lax.max, @@ -32,7 +29,6 @@ result = lax.reduce_window( ) ``` - Or consider manual dilation if necessary. -**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` \ No newline at end of file +**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` diff --git a/skills/neuron-nki-docs/references/debugging/error-codes/EVRF018.md b/skills/neuron-nki-docs/references/debugging/error-codes/EVRF018.md index 895bdf0..490fc2a 100644 --- a/skills/neuron-nki-docs/references/debugging/error-codes/EVRF018.md +++ b/skills/neuron-nki-docs/references/debugging/error-codes/EVRF018.md @@ -7,7 +7,6 @@ NCC_EVRF018 Erroneous code example: - ```python result = lax.reduce_window( jnp.ones((1, 4, 4, 1)), -jnp.inf, lax.max, @@ -18,10 +17,8 @@ result = lax.reduce_window( ) ``` - If possible, remove window_dilation or change values to be all 1s: - ```python result = lax.reduce_window( jnp.ones((1, 4, 4, 1)), -jnp.inf, lax.max, @@ -32,7 +29,6 @@ result = lax.reduce_window( ) ``` - Or consider manual dilation if necessary. -**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` \ No newline at end of file +**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` diff --git a/skills/neuron-nki-docs/references/debugging/error-codes/EVRF019.md b/skills/neuron-nki-docs/references/debugging/error-codes/EVRF019.md index 8ae5d3d..adbae39 100644 --- a/skills/neuron-nki-docs/references/debugging/error-codes/EVRF019.md +++ b/skills/neuron-nki-docs/references/debugging/error-codes/EVRF019.md @@ -7,7 +7,6 @@ NCC_EVRF019 Erroneous code example: - ```python # reduce-window operation with more or less than 2 operands is not supported # 4 operands are being provided instead of 2 @@ -21,10 +20,8 @@ lax.reduce_window( ) ``` - If possible, split multi-operand reduce_window with multiple single-operand reduce_window operations. - ```python # For max pooling # 2 operands are correctly being provided @@ -49,5 +46,4 @@ min_pool = lax.reduce_window( ) ``` - -**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` \ No newline at end of file +**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` diff --git a/skills/neuron-nki-docs/references/debugging/error-codes/EVRF022.md b/skills/neuron-nki-docs/references/debugging/error-codes/EVRF022.md index 7bfcb08..f7c202e 100644 --- a/skills/neuron-nki-docs/references/debugging/error-codes/EVRF022.md +++ b/skills/neuron-nki-docs/references/debugging/error-codes/EVRF022.md @@ -7,7 +7,6 @@ NCC_EVRF022 Erroneous code example: - ```python def forward(self, input, other): return torch.bitwise_right_shift(input, other) @@ -18,10 +17,8 @@ input = torch.tensor([16, 32, 64], dtype=torch.int16) other = torch.tensor([1, 2, 3], dtype=torch.int16) ``` - To fix this error: - ```python def forward(self, input, other): return torch.bitwise_right_shift(input, other) @@ -31,5 +28,4 @@ input = torch.tensor([16, 32, 64], dtype=torch.int32) other = torch.tensor([1, 2, 3], dtype=torch.int16) ``` - -**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` \ No newline at end of file +**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` diff --git a/skills/neuron-nki-docs/references/debugging/error-codes/EVRF024.md b/skills/neuron-nki-docs/references/debugging/error-codes/EVRF024.md index 27d9b34..b8b3ed9 100644 --- a/skills/neuron-nki-docs/references/debugging/error-codes/EVRF024.md +++ b/skills/neuron-nki-docs/references/debugging/error-codes/EVRF024.md @@ -7,13 +7,12 @@ NCC_EVRF024 There are two potential ways to fix this issue: -* Simply reduce the batch/tensor size if possible +- Simply reduce the batch/tensor size if possible -* Utilize tensor parallelism via neuronx-distributed +- Utilize tensor parallelism via neuronx-distributed Short snippet of tensor parallelism: - ```python class ParallelSelfAttention(transformers.models.bert.modeling_bert.BertSelfAttention): def __init__(self, config, position_embedding_type=None): @@ -36,7 +35,6 @@ class ParallelSelfAttention(transformers.models.bert.modeling_bert.BertSelfAtten self.all_head_size = self.all_head_size // tp_size ``` - For more information: [https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/neuronx-distributed/activation_memory_reduction.html](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/neuronx-distributed/activation_memory_reduction.html) -**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` \ No newline at end of file +**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` diff --git a/skills/neuron-nki-docs/references/debugging/error-codes/EVRF031.md b/skills/neuron-nki-docs/references/debugging/error-codes/EVRF031.md index 9e2365a..00906a9 100644 --- a/skills/neuron-nki-docs/references/debugging/error-codes/EVRF031.md +++ b/skills/neuron-nki-docs/references/debugging/error-codes/EVRF031.md @@ -7,7 +7,6 @@ NCC_EVRF031 Erroneous code example: - ```python # size 3 in dimension 0 operand = jnp.zeros((3, 4), dtype=jnp.float32) @@ -30,10 +29,8 @@ result = lax.scatter( ) ``` - Ensure that the iota size matches the operand dimension size: - ```python N = 3 D = 4 @@ -58,5 +55,4 @@ result = lax.scatter( ) ``` - -**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` \ No newline at end of file +**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` diff --git a/skills/neuron-nki-docs/references/debugging/error-codes/EXSP001.md b/skills/neuron-nki-docs/references/debugging/error-codes/EXSP001.md index aece544..860a173 100644 --- a/skills/neuron-nki-docs/references/debugging/error-codes/EXSP001.md +++ b/skills/neuron-nki-docs/references/debugging/error-codes/EXSP001.md @@ -8,13 +8,12 @@ NCC_EXSP001 There are several ways to potentially fix this issue. -* Simply reduce the batch/tensor size if possible +- Simply reduce the batch/tensor size if possible -* Utilize pipeline/tensor parallelism via neuronx-distributed +- Utilize pipeline/tensor parallelism via neuronx-distributed Short snippet of tensor parallelism: - ```python class ParallelSelfAttention(transformers.models.bert.modeling_bert.BertSelfAttention): def __init__(self, config, position_embedding_type=None): @@ -37,11 +36,10 @@ class ParallelSelfAttention(transformers.models.bert.modeling_bert.BertSelfAtten self.all_head_size = self.all_head_size // tp_size ``` - For more information: -* [https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/neuronx-distributed/activation_memory_reduction.html](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/neuronx-distributed/activation_memory_reduction.html) +- [https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/neuronx-distributed/activation_memory_reduction.html](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/neuronx-distributed/activation_memory_reduction.html) -* [https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/nxd-training/app_notes/nxd-training-pp-appnote.html](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/nxd-training/app_notes/nxd-training-pp-appnote.html) +- [https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/nxd-training/app_notes/nxd-training-pp-appnote.html](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/nxd-training/app_notes/nxd-training-pp-appnote.html) -**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` \ No newline at end of file +**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` diff --git a/skills/neuron-nki-docs/references/debugging/error-codes/EXTP004.md b/skills/neuron-nki-docs/references/debugging/error-codes/EXTP004.md index 1502f9a..a32429a 100644 --- a/skills/neuron-nki-docs/references/debugging/error-codes/EXTP004.md +++ b/skills/neuron-nki-docs/references/debugging/error-codes/EXTP004.md @@ -9,8 +9,8 @@ Consider applying model parallelism as partitioning the model will help break la For more information: -* [https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/neuronx-distributed/api_guide.html#api-guide](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/neuronx-distributed/api_guide.html#api-guide) +- [https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/neuronx-distributed/api_guide.html#api-guide](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/neuronx-distributed/api_guide.html#api-guide) -* [https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/nxd-training/app_notes/nxd-training-pp-appnote.html](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/nxd-training/app_notes/nxd-training-pp-appnote.html) +- [https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/nxd-training/app_notes/nxd-training-pp-appnote.html](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/nxd-training/app_notes/nxd-training-pp-appnote.html) -**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` \ No newline at end of file +**This document is relevant for**: `Inf1`, `Inf2`, `Trn1`, `Trn2`, `Trn3` diff --git a/skills/neuron-nki-docs/references/indices/hierarchical-toc.md b/skills/neuron-nki-docs/references/indices/hierarchical-toc.md index c8e62b3..eeaa322 100644 --- a/skills/neuron-nki-docs/references/indices/hierarchical-toc.md +++ b/skills/neuron-nki-docs/references/indices/hierarchical-toc.md @@ -17,6 +17,7 @@ This document provides a comprehensive hierarchical index of all NKI (Neuron Ker - [NKI Architecture Guides Overview](../architecture/nki_arch_guides.md) - Introduction to architecture documentation ### 2.1 NeuronDevice Architectures + - [Trainium/Inferentia2 Architecture](../architecture/trainium_inferentia2_arch.md) - NeuronCore-v2 architecture details - NeuronCore-v2 Compute Engines (Tensor, Vector, Scalar, GpSimd) - Memory Hierarchy (HBM, SBUF, PSUM) @@ -29,6 +30,7 @@ This document provides a comprehensive hierarchical index of all NKI (Neuron Ker ## 3. Programming Guide ### 3.1 Core Concepts + - [NKI Language Guide](../programming/nki-language-guide.md) - Comprehensive language syntax guide - [NKI Compiler Documentation](../programming/nki-compiler.md) - Compiler integration and usage - [Memory Hierarchy Overview](../programming/memory-hierarchy-overview.md) - Understanding HBM, SBUF, and PSUM @@ -38,19 +40,23 @@ This document provides a comprehensive hierarchical index of all NKI (Neuron Ker - [DMA Overview](../programming/nki-dma-overview.md) - Data movement operations ### 3.2 Advanced Topics + - [NKI APS](../programming/nki-aps.md) - Advanced programming systems - [Logical NeuronCore (LNC)](../programming/lnc.md) - Logical NeuronCore configurations - [Framework Custom Operators](../programming/framework_custom_op.md) - PyTorch and JAX integration - [Using Prebuilt Kernels](../programming/tutorial-use-a-prebuilt-kernel.md) - Working with NKI Library kernels ### 3.3 API Reference + - [API Reference Index](../programming/api/index.md) - Complete API documentation index - [API Overview](../programming/api/api-overview.md) - High-level API organization #### nki Module + - [nki Module](../programming/api/nki.md) - Top-level nki module #### nki.language Module + - [nki.language Module Overview](../programming/api/nki.language.md) - Language-level APIs - [nki.language - Creation Operations](../programming/api/api-nki-language-creation.md) - Tensor creation (ndarray, zeros) - [nki.language - Memory Operations](../programming/api/api-nki-language-memory.md) - Memory buffers (sbuf, psum, hbm) @@ -60,6 +66,7 @@ This document provides a comprehensive hierarchical index of all NKI (Neuron Ker - [nki.language - Miscellaneous](../programming/api/api-nki-language-misc.md) - Other language APIs #### nki.isa Module + - [nki.isa Module Overview](../programming/api/nki.isa.md) - Instruction Set Architecture APIs - [nki.isa - Tensor Engine](../programming/api/api-nki-isa-tensor.md) - Tensor Engine instructions (nc_matmul, nc_transpose) - [nki.isa - Vector Engine](../programming/api/api-nki-isa-vector.md) - Vector Engine instructions (bn_stats, bn_aggr) @@ -70,29 +77,37 @@ This document provides a comprehensive hierarchical index of all NKI (Neuron Ker - [nki.isa - Local Collective (LNC)](../programming/api/api-nki-isa-local-collective.md) - core_barrier, sendrecv #### nki.collectives Module + - [nki.collectives](../programming/api/api-nki-collectives.md) - Collective communication (all_gather, all_reduce, all_to_all, collective_permute, reduce_scatter, ...) #### Shared APIs + - [nki.api.shared](../programming/api/nki.api.shared.md) - Shared data types and operators #### NkiTensor + - [NkiTensor View Methods](../programming/api/api-nki-tensor.md) - Composable tensor view methods (slice, select, permute, reshape, ap, etc.) #### Tools + - [NKI Tools](../programming/api/api-nki-tools.md) - Development and debugging tools ### 3.4 Tutorials + - [NKI Tutorials Index](../programming/tutorials/tutorials.md) - All tutorials overview #### Basic Operations + - [Matrix Multiplication](../programming/tutorials/matrix_multiplication.md) - Matmul implementation and optimization - [2D Transpose](../programming/tutorials/transpose2d.md) - Efficient transpose operations - [Average Pooling 2D](../programming/tutorials/average_pool2d.md) - Pooling kernel implementation #### Advanced Kernels + - [Fused Mamba](../programming/tutorials/fused_mamba.md) - State space model kernel #### Performance + - [Kernel Optimization](../programming/tutorials/kernel-optimization.md) - Optimization techniques --- @@ -113,33 +128,41 @@ This document provides a comprehensive hierarchical index of all NKI (Neuron Ker ## 5. Reference ### 5.1 General Reference + - [NKI FAQ](../reference/nki_faq.md) - Frequently asked questions - [NKI Release Notes](../reference/nki_rn.md) - Version history and changes ### 5.2 Migration Guides + - [NKI 0.3.0 Update Guide](../reference/migration/nki-030-update-guide.md) - Updating from NKI 0.2.0 to 0.3.0 (GA) - [NKI Migration Guide (Beta 1 to Beta 2)](../reference/migration/nki-migration-guide.md) - Upgrading from Beta 1 to Beta 2 - [Block Dimension Migration Guide](../reference/migration/nki_block_dimension_migration_guide.md) - Block dimension changes - [NKI 0.6.0 Dynamic Loop Migration Guide](../reference/migration/nki-060-dynamic-loop-migration-guide.md) - Migrate dynamic_range / bare `while reg:` to nl.fori_loop / nl.while_loop (Parsing→Tracing) ### 5.3 NKI Library Kernels + Pre-built reference kernels for common operations: #### Normalization and Quantization + - [RMSNorm-Quant Kernel](../reference/library/rmsnorm-quant.md) - RMS normalization with quantization - [RMSNorm-Quant Design](../reference/library/design-rmsnorm-quant.md) - Design documentation #### QKV Projection + - [QKV Kernel](../reference/library/qkv.md) - Query-Key-Value projection #### Attention Kernels + - [Attention CTE Kernel](../reference/library/attention-cte.md) - Context encoding attention - [Attention TKG Kernel](../reference/library/attention-tkg.md) - Token generation attention #### MLP Kernels + - [MLP Kernel](../reference/library/mlp.md) - Multi-Layer Perceptron #### Output Projection + - [Output Projection CTE Kernel](../reference/library/output-projection-cte.md) - Context encoding output projection - [Output Projection TKG Kernel](../reference/library/output-projection-tkg.md) - Token generation output projection @@ -150,48 +173,49 @@ Pre-built reference kernels for common operations: - [Compiler Error Codes Index](../debugging/error-codes-index.md) - All error codes overview ### 6.1 Error Code Reference + Individual error code documentation: -| Error Code | Description | -|------------|-------------| -| [EARG001](../debugging/error-codes/EARG001.md) | Unsupported LNC configuration | -| [EBVF030](../debugging/error-codes/EBVF030.md) | Instruction limit exceeded | -| [EHCA005](../debugging/error-codes/EHCA005.md) | Unrecognized custom call target | -| [EOOM001](../debugging/error-codes/EOOM001.md) | Activation memory limit exceeded | -| [EOOM002](../debugging/error-codes/EOOM002.md) | Activation memory limit exceeded (variant) | -| [ESFH002](../debugging/error-codes/ESFH002.md) | 64-bit to 32-bit conversion error | -| [ESPP004](../debugging/error-codes/ESPP004.md) | Unsupported data type | -| [ESPP047](../debugging/error-codes/ESPP047.md) | Unsupported 8-bit floating-point type | -| [EUOC002](../debugging/error-codes/EUOC002.md) | Unsupported operator | -| [EVRF001](../debugging/error-codes/EVRF001.md) | Unsupported operator (verification) | -| [EVRF004](../debugging/error-codes/EVRF004.md) | Complex data types unsupported | -| [EVRF005](../debugging/error-codes/EVRF005.md) | Unsupported FP8 variants | -| [EVRF006](../debugging/error-codes/EVRF006.md) | RNG algorithm error | -| [EVRF007](../debugging/error-codes/EVRF007.md) | Instruction limit exceeded (verification) | -| [EVRF009](../debugging/error-codes/EVRF009.md) | Memory limit exceeded (verification) | -| [EVRF010](../debugging/error-codes/EVRF010.md) | Simultaneous dilation unsupported | -| [EVRF011](../debugging/error-codes/EVRF011.md) | Strided convolution with dilated input | -| [EVRF013](../debugging/error-codes/EVRF013.md) | TopK integer input unsupported | +| Error Code | Description | +| ---------------------------------------------- | ---------------------------------------------- | +| [EARG001](../debugging/error-codes/EARG001.md) | Unsupported LNC configuration | +| [EBVF030](../debugging/error-codes/EBVF030.md) | Instruction limit exceeded | +| [EHCA005](../debugging/error-codes/EHCA005.md) | Unrecognized custom call target | +| [EOOM001](../debugging/error-codes/EOOM001.md) | Activation memory limit exceeded | +| [EOOM002](../debugging/error-codes/EOOM002.md) | Activation memory limit exceeded (variant) | +| [ESFH002](../debugging/error-codes/ESFH002.md) | 64-bit to 32-bit conversion error | +| [ESPP004](../debugging/error-codes/ESPP004.md) | Unsupported data type | +| [ESPP047](../debugging/error-codes/ESPP047.md) | Unsupported 8-bit floating-point type | +| [EUOC002](../debugging/error-codes/EUOC002.md) | Unsupported operator | +| [EVRF001](../debugging/error-codes/EVRF001.md) | Unsupported operator (verification) | +| [EVRF004](../debugging/error-codes/EVRF004.md) | Complex data types unsupported | +| [EVRF005](../debugging/error-codes/EVRF005.md) | Unsupported FP8 variants | +| [EVRF006](../debugging/error-codes/EVRF006.md) | RNG algorithm error | +| [EVRF007](../debugging/error-codes/EVRF007.md) | Instruction limit exceeded (verification) | +| [EVRF009](../debugging/error-codes/EVRF009.md) | Memory limit exceeded (verification) | +| [EVRF010](../debugging/error-codes/EVRF010.md) | Simultaneous dilation unsupported | +| [EVRF011](../debugging/error-codes/EVRF011.md) | Strided convolution with dilated input | +| [EVRF013](../debugging/error-codes/EVRF013.md) | TopK integer input unsupported | | [EVRF015](../debugging/error-codes/EVRF015.md) | Unrecognized custom call target (verification) | -| [EVRF016](../debugging/error-codes/EVRF016.md) | Scatter-reduce data type error | -| [EVRF017](../debugging/error-codes/EVRF017.md) | Reduce-window base dilation | -| [EVRF018](../debugging/error-codes/EVRF018.md) | Reduce-window window dilation | -| [EVRF019](../debugging/error-codes/EVRF019.md) | Reduce-window operand count | -| [EVRF022](../debugging/error-codes/EVRF022.md) | Shift-right-arithmetic bit width | -| [EVRF024](../debugging/error-codes/EVRF024.md) | Output tensor size limit | -| [EVRF031](../debugging/error-codes/EVRF031.md) | Scatter out-of-bounds | -| [EXSP001](../debugging/error-codes/EXSP001.md) | Memory limit exceeded (expansion) | -| [EXTP004](../debugging/error-codes/EXTP004.md) | Instruction limit exceeded (expansion) | +| [EVRF016](../debugging/error-codes/EVRF016.md) | Scatter-reduce data type error | +| [EVRF017](../debugging/error-codes/EVRF017.md) | Reduce-window base dilation | +| [EVRF018](../debugging/error-codes/EVRF018.md) | Reduce-window window dilation | +| [EVRF019](../debugging/error-codes/EVRF019.md) | Reduce-window operand count | +| [EVRF022](../debugging/error-codes/EVRF022.md) | Shift-right-arithmetic bit width | +| [EVRF024](../debugging/error-codes/EVRF024.md) | Output tensor size limit | +| [EVRF031](../debugging/error-codes/EVRF031.md) | Scatter out-of-bounds | +| [EXSP001](../debugging/error-codes/EXSP001.md) | Memory limit exceeded (expansion) | +| [EXTP004](../debugging/error-codes/EXTP004.md) | Instruction limit exceeded (expansion) | --- ## Quick Navigation -| Category | Key Documents | -|----------|--------------| -| **New Users** | [Quickstart](../programming/quickstart-implement-run-kernel.md), [Language Guide](../programming/nki-language-guide.md) | -| **Architecture** | [Trainium/Inf2](../architecture/trainium_inferentia2_arch.md), [Trainium2](../architecture/trainium2_arch.md) | -| **API Reference** | [nki.language](../programming/api/nki.language.md), [nki.isa](../programming/api/nki.isa.md) | -| **Tutorials** | [Matrix Multiplication](../programming/tutorials/matrix_multiplication.md), [Fused Mamba](../programming/tutorials/fused_mamba.md) | -| **Performance** | [Performance Guide](../optimization/nki_perf_guide.md), [Profiling](../optimization/use-neuron-profile.md) | -| **Debugging** | [Error Codes](../debugging/error-codes-index.md), [FAQ](../reference/nki_faq.md) | +| Category | Key Documents | +| ----------------- | ---------------------------------------------------------------------------------------------------------------------------------- | +| **New Users** | [Quickstart](../programming/quickstart-implement-run-kernel.md), [Language Guide](../programming/nki-language-guide.md) | +| **Architecture** | [Trainium/Inf2](../architecture/trainium_inferentia2_arch.md), [Trainium2](../architecture/trainium2_arch.md) | +| **API Reference** | [nki.language](../programming/api/nki.language.md), [nki.isa](../programming/api/nki.isa.md) | +| **Tutorials** | [Matrix Multiplication](../programming/tutorials/matrix_multiplication.md), [Fused Mamba](../programming/tutorials/fused_mamba.md) | +| **Performance** | [Performance Guide](../optimization/nki_perf_guide.md), [Profiling](../optimization/use-neuron-profile.md) | +| **Debugging** | [Error Codes](../debugging/error-codes-index.md), [FAQ](../reference/nki_faq.md) | diff --git a/skills/neuron-nki-docs/references/indices/symbol-lookup.md b/skills/neuron-nki-docs/references/indices/symbol-lookup.md index 5d6f815..d7d594c 100644 --- a/skills/neuron-nki-docs/references/indices/symbol-lookup.md +++ b/skills/neuron-nki-docs/references/indices/symbol-lookup.md @@ -6,372 +6,372 @@ Quick reference for finding NKI API function and symbol documentation. Symbols a ## Quick Module Reference -| Module | Description | Documentation | -|--------|-------------|---------------| -| `nki` | Top-level NKI module | [nki](../programming/api/nki.md) | -| `nki.language` | High-level language APIs | [nki.language](../programming/api/nki.language.md) | -| `nki.isa` | Low-level ISA instructions | [nki.isa](../programming/api/nki.isa.md) | +| Module | Description | Documentation | +| ----------------- | ------------------------------------------------------ | ------------------------------------------------------------ | +| `nki` | Top-level NKI module | [nki](../programming/api/nki.md) | +| `nki.language` | High-level language APIs | [nki.language](../programming/api/nki.language.md) | +| `nki.isa` | Low-level ISA instructions | [nki.isa](../programming/api/nki.isa.md) | | `nki.collectives` | Collective communication (all_gather, all_reduce, ...) | [nki.collectives](../programming/api/api-nki-collectives.md) | -| `nki.api.shared` | Shared data types and operators | [nki.api.shared](../programming/api/nki.api.shared.md) | +| `nki.api.shared` | Shared data types and operators | [nki.api.shared](../programming/api/nki.api.shared.md) | --- ## A -| Symbol | Module | Description | Documentation | -|--------|--------|-------------|---------------| -| `abs` | nki.language | Op specifier for abs. | [nki.language.abs](../programming/api/api-nki-language-operators.md#nki-language-abs) | -| `abs_max` | nki.language | Element-wise absolute maximum (trn3 only) | [nki.api.shared](../programming/api/nki.api.shared.md) | -| `abs_min` | nki.language | Element-wise absolute minimum (trn3 only) | [nki.api.shared](../programming/api/nki.api.shared.md) | -| `activate2` | nki.isa | Two-stage tensor-scalar + activation in one instruction (trn3 only) | [nki.isa.activate2](../programming/api/api-nki-isa-scalar.md#nki-isa-activation) | -| `activation` | nki.isa | Apply activation function with optional scale/bias | [nki.isa.activation](../programming/api/api-nki-isa-scalar.md#nki-isa-activation) | -| `activation_reduce` | nki.isa | Activation with free-dimension reduction | [nki.isa.activation_reduce](../programming/api/api-nki-isa-scalar.md#nki-isa-activation_reduce) | -| `add` | nki.language | Op specifier for add. | [nki.language.add](../programming/api/api-nki-language-operators.md#nki-language-add) | -| `affine_range` | nki.language | Loop iterator (legacy alias for `range`) | [nki.language.affine_range](../programming/api/nki.language.md) | -| `affine_select` | nki.isa | Select elements using affine predicate | [nki.isa.affine_select](../programming/api/api-nki-isa-utility.md#nki-isa-affine_select) | -| `all` | nki.language | Whether all elements along the specified axis (or axes) evaluate to True. | [nki.language.all](../programming/api/api-nki-language-misc.md#nki-language-all) | -| `all_gather` | nki.collectives | Perform an all-gather on the given replica group and input/output tensors. | [nki.collectives.all_gather](../programming/api/api-nki-collectives.md#nki-collectives-all_gather) | -| `all_gather_v` | nki.collectives | Perform a variable-length all-gather on the given replica group. | [nki.collectives.all_gather_v](../programming/api/api-nki-collectives.md#nki-collectives-all_gather_v) | -| `all_reduce` | nki.collectives | Perform an all-reduce on the given replica group and input/output tensors. | [nki.collectives.all_reduce](../programming/api/api-nki-collectives.md#nki-collectives-all_reduce) | -| `all_to_all` | nki.collectives | Perform an all-to-all on the given replica group and input/output tensors. | [nki.collectives.all_to_all](../programming/api/api-nki-collectives.md#nki-collectives-all_to_all) | -| `all_to_all_v` | nki.collectives | Executes an all-to-all collective where each rank can send | [nki.collectives.all_to_all_v](../programming/api/api-nki-collectives.md#nki-collectives-all_to_all_v) | -| `ap` | nki.tensor | Low-level access pattern override (escape hatch). | [NkiTensor.ap](../programming/api/api-nki-tensor.md#nki-tensor-ap) | -| `arctan` | nki.language | Op specifier for arctan. | [nki.language.arctan](../programming/api/api-nki-language-operators.md#nki-language-arctan) | -| `average` | nki.language | Op specifier for average. | [nki.language.average](../programming/api/api-nki-language-operators.md#nki-language-average) | +| Symbol | Module | Description | Documentation | +| ------------------- | --------------- | -------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------ | +| `abs` | nki.language | Op specifier for abs. | [nki.language.abs](../programming/api/api-nki-language-operators.md#nki-language-abs) | +| `abs_max` | nki.language | Element-wise absolute maximum (trn3 only) | [nki.api.shared](../programming/api/nki.api.shared.md) | +| `abs_min` | nki.language | Element-wise absolute minimum (trn3 only) | [nki.api.shared](../programming/api/nki.api.shared.md) | +| `activate2` | nki.isa | Two-stage tensor-scalar + activation in one instruction (trn3 only) | [nki.isa.activate2](../programming/api/api-nki-isa-scalar.md#nki-isa-activation) | +| `activation` | nki.isa | Apply activation function with optional scale/bias | [nki.isa.activation](../programming/api/api-nki-isa-scalar.md#nki-isa-activation) | +| `activation_reduce` | nki.isa | Activation with free-dimension reduction | [nki.isa.activation_reduce](../programming/api/api-nki-isa-scalar.md#nki-isa-activation_reduce) | +| `add` | nki.language | Op specifier for add. | [nki.language.add](../programming/api/api-nki-language-operators.md#nki-language-add) | +| `affine_range` | nki.language | Loop iterator (legacy alias for `range`) | [nki.language.affine_range](../programming/api/nki.language.md) | +| `affine_select` | nki.isa | Select elements using affine predicate | [nki.isa.affine_select](../programming/api/api-nki-isa-utility.md#nki-isa-affine_select) | +| `all` | nki.language | Whether all elements along the specified axis (or axes) evaluate to True. | [nki.language.all](../programming/api/api-nki-language-misc.md#nki-language-all) | +| `all_gather` | nki.collectives | Perform an all-gather on the given replica group and input/output tensors. | [nki.collectives.all_gather](../programming/api/api-nki-collectives.md#nki-collectives-all_gather) | +| `all_gather_v` | nki.collectives | Perform a variable-length all-gather on the given replica group. | [nki.collectives.all_gather_v](../programming/api/api-nki-collectives.md#nki-collectives-all_gather_v) | +| `all_reduce` | nki.collectives | Perform an all-reduce on the given replica group and input/output tensors. | [nki.collectives.all_reduce](../programming/api/api-nki-collectives.md#nki-collectives-all_reduce) | +| `all_to_all` | nki.collectives | Perform an all-to-all on the given replica group and input/output tensors. | [nki.collectives.all_to_all](../programming/api/api-nki-collectives.md#nki-collectives-all_to_all) | +| `all_to_all_v` | nki.collectives | Executes an all-to-all collective where each rank can send | [nki.collectives.all_to_all_v](../programming/api/api-nki-collectives.md#nki-collectives-all_to_all_v) | +| `ap` | nki.tensor | Low-level access pattern override (escape hatch). | [NkiTensor.ap](../programming/api/api-nki-tensor.md#nki-tensor-ap) | +| `arctan` | nki.language | Op specifier for arctan. | [nki.language.arctan](../programming/api/api-nki-language-operators.md#nki-language-arctan) | +| `average` | nki.language | Op specifier for average. | [nki.language.average](../programming/api/api-nki-language-operators.md#nki-language-average) | --- ## B -| Symbol | Module | Description | Documentation | -|--------|--------|-------------|---------------| -| `bfloat16` | nki.language | BF16 data type (1S,8E,7M) | [nki.language.bfloat16](../programming/api/api-nki-language-types.md#nki-language-bfloat16) | -| `bitwise_and` | nki.language | Op specifier for bitwise_and. | [nki.language.bitwise_and](../programming/api/api-nki-language-operators.md#nki-language-bitwise_and) | -| `bitwise_or` | nki.language | Op specifier for bitwise_or. | [nki.language.bitwise_or](../programming/api/api-nki-language-operators.md#nki-language-bitwise_or) | -| `bitwise_xor` | nki.language | Op specifier for bitwise_xor. | [nki.language.bitwise_xor](../programming/api/api-nki-language-operators.md#nki-language-bitwise_xor) | -| `bn_aggr` | nki.isa | Aggregate batch norm statistics | [nki.isa.bn_aggr](../programming/api/api-nki-isa-vector.md#nki-isa-bn_aggr) | -| `bn_stats` | nki.isa | Compute batch norm statistics | [nki.isa.bn_stats](../programming/api/api-nki-isa-vector.md#nki-isa-bn_stats) | -| `bool_` | nki.language | Boolean data type | [nki.language.bool_](../programming/api/api-nki-language-types.md#nki-language-bool_) | -| `broadcast` | nki.tensor | Expand a size-1 dimension to `size` by repeating elements. | [NkiTensor.broadcast](../programming/api/api-nki-tensor.md#nki-tensor-broadcast) | -| `broadcast_to` | nki.language | Broadcast a tile to a new shape following numpy broadcasting rules. | [nki.language.broadcast_to](../programming/api/api-nki-language-misc.md#nki-language-broadcast_to) | -| `bypass` | nki.language | Op specifier for bypass. | [nki.language.bypass](../programming/api/api-nki-language-operators.md#nki-language-bypass) | +| Symbol | Module | Description | Documentation | +| -------------- | ------------ | ------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------- | +| `bfloat16` | nki.language | BF16 data type (1S,8E,7M) | [nki.language.bfloat16](../programming/api/api-nki-language-types.md#nki-language-bfloat16) | +| `bitwise_and` | nki.language | Op specifier for bitwise_and. | [nki.language.bitwise_and](../programming/api/api-nki-language-operators.md#nki-language-bitwise_and) | +| `bitwise_or` | nki.language | Op specifier for bitwise_or. | [nki.language.bitwise_or](../programming/api/api-nki-language-operators.md#nki-language-bitwise_or) | +| `bitwise_xor` | nki.language | Op specifier for bitwise_xor. | [nki.language.bitwise_xor](../programming/api/api-nki-language-operators.md#nki-language-bitwise_xor) | +| `bn_aggr` | nki.isa | Aggregate batch norm statistics | [nki.isa.bn_aggr](../programming/api/api-nki-isa-vector.md#nki-isa-bn_aggr) | +| `bn_stats` | nki.isa | Compute batch norm statistics | [nki.isa.bn_stats](../programming/api/api-nki-isa-vector.md#nki-isa-bn_stats) | +| `bool_` | nki.language | Boolean data type | [nki.language.bool\_](../programming/api/api-nki-language-types.md#nki-language-bool_) | +| `broadcast` | nki.tensor | Expand a size-1 dimension to `size` by repeating elements. | [NkiTensor.broadcast](../programming/api/api-nki-tensor.md#nki-tensor-broadcast) | +| `broadcast_to` | nki.language | Broadcast a tile to a new shape following numpy broadcasting rules. | [nki.language.broadcast_to](../programming/api/api-nki-language-misc.md#nki-language-broadcast_to) | +| `bypass` | nki.language | Op specifier for bypass. | [nki.language.bypass](../programming/api/api-nki-language-operators.md#nki-language-bypass) | --- ## C -| Symbol | Module | Description | Documentation | -|--------|--------|-------------|---------------| -| `ceil` | nki.language | Op specifier for ceil. | [nki.language.ceil](../programming/api/api-nki-language-operators.md#nki-language-ceil) | -| `collective_permute` | nki.collectives | Send and receive data between ranks based on explicitly defined source-target pa | [nki.collectives.collective_permute](../programming/api/api-nki-collectives.md#nki-collectives-collective_permute) | -| `collective_permute_implicit` | nki.collectives | Send and receive data between ranks in a ring, where sources and destinations ar | [nki.collectives.collective_permute_implicit](../programming/api/api-nki-collectives.md#nki-collectives-collective_permute_implicit) | -| `collective_permute_implicit_current_processing_rank_id` | nki.collectives | Returns the rank ID of the data to be processed in the current ring iteration. | [nki.collectives.collective_permute_implicit_current_processing_rank_id](../programming/api/api-nki-collectives.md#nki-collectives-collective_permute_implicit_current_processing_rank_id) | -| `collective_permute_implicit_reduce` | nki.collectives | Perform an implicit collective permute with reduction in a ring, where sources a | [nki.collectives.collective_permute_implicit_reduce](../programming/api/api-nki-collectives.md#nki-collectives-collective_permute_implicit_reduce) | -| `copy` | nki.language | Op specifier for copy. | [nki.language.copy](../programming/api/api-nki-language-operators.md#nki-language-copy) | -| `core_barrier` | nki.isa | Synchronize across NeuronCores | [nki.isa.core_barrier](../programming/api/api-nki-isa-tensor.md#nki-isa-core_barrier) | -| `cos` | nki.language | Op specifier for cos. | [nki.language.cos](../programming/api/api-nki-language-operators.md#nki-language-cos) | +| Symbol | Module | Description | Documentation | +| -------------------------------------------------------- | --------------- | -------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | +| `ceil` | nki.language | Op specifier for ceil. | [nki.language.ceil](../programming/api/api-nki-language-operators.md#nki-language-ceil) | +| `collective_permute` | nki.collectives | Send and receive data between ranks based on explicitly defined source-target pa | [nki.collectives.collective_permute](../programming/api/api-nki-collectives.md#nki-collectives-collective_permute) | +| `collective_permute_implicit` | nki.collectives | Send and receive data between ranks in a ring, where sources and destinations ar | [nki.collectives.collective_permute_implicit](../programming/api/api-nki-collectives.md#nki-collectives-collective_permute_implicit) | +| `collective_permute_implicit_current_processing_rank_id` | nki.collectives | Returns the rank ID of the data to be processed in the current ring iteration. | [nki.collectives.collective_permute_implicit_current_processing_rank_id](../programming/api/api-nki-collectives.md#nki-collectives-collective_permute_implicit_current_processing_rank_id) | +| `collective_permute_implicit_reduce` | nki.collectives | Perform an implicit collective permute with reduction in a ring, where sources a | [nki.collectives.collective_permute_implicit_reduce](../programming/api/api-nki-collectives.md#nki-collectives-collective_permute_implicit_reduce) | +| `copy` | nki.language | Op specifier for copy. | [nki.language.copy](../programming/api/api-nki-language-operators.md#nki-language-copy) | +| `core_barrier` | nki.isa | Synchronize across NeuronCores | [nki.isa.core_barrier](../programming/api/api-nki-isa-tensor.md#nki-isa-core_barrier) | +| `cos` | nki.language | Op specifier for cos. | [nki.language.cos](../programming/api/api-nki-language-operators.md#nki-language-cos) | --- ## D -| Symbol | Module | Description | Documentation | -|--------|--------|-------------|---------------| -| `device_print` | nki.language | Print debug output from kernel | [nki.language.device_print](../programming/api/nki.language.md) | -| `dge_mode` | nki.isa | DMA Descriptor Generation Engine mode enum | [nki.isa.dge_mode](../programming/api/nki.isa.md) | -| `divide` | nki.language | Op specifier for divide. | [nki.language.divide](../programming/api/api-nki-language-operators.md#nki-language-divide) | -| `dma_compute` | nki.isa | Math operations using DMA engines (replaces dma_copy RMW) | [nki.isa.dma_compute](../programming/api/api-nki-isa-memory.md#nki-isa-dma_compute) | -| `dma_copy` | nki.isa | Copy data using DMA engines | [nki.isa.dma_copy](../programming/api/api-nki-isa-memory.md#nki-isa-dma_copy) | -| `dma_engine` | nki.isa | DMA engine enum (dma, gpsimd_dma) | [nki.isa.dma_engine](../programming/api/nki.isa.md) | -| `dma_transpose` | nki.isa | Transpose using DMA engines | [nki.isa.dma_transpose](../programming/api/api-nki-isa-memory.md#nki-isa-dma_transpose) | -| `dropout` | nki.isa | Apply dropout to tensor | [nki.isa.dropout](../programming/api/api-nki-isa-scalar.md#nki-isa-dropout) | -| `dropout` | nki.language | Randomly zeroes some of the elements of the input tile given a probability rate. | [nki.language.dropout](../programming/api/api-nki-language-misc.md#nki-language-dropout) | -| `ds` | nki.language | Dynamic slice for tensor indexing | [nki.language.ds](../programming/api/nki.language.md) | -| `dynamic_range` | nki.language | Create a sequence for **dynamic** loop iteration. | [nki.language.dynamic_range](../programming/api/api-nki-language-dims.md#nki-language-dynamic_range) | +| Symbol | Module | Description | Documentation | +| --------------- | ------------ | -------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------- | +| `device_print` | nki.language | Print debug output from kernel | [nki.language.device_print](../programming/api/nki.language.md) | +| `dge_mode` | nki.isa | DMA Descriptor Generation Engine mode enum | [nki.isa.dge_mode](../programming/api/nki.isa.md) | +| `divide` | nki.language | Op specifier for divide. | [nki.language.divide](../programming/api/api-nki-language-operators.md#nki-language-divide) | +| `dma_compute` | nki.isa | Math operations using DMA engines (replaces dma_copy RMW) | [nki.isa.dma_compute](../programming/api/api-nki-isa-memory.md#nki-isa-dma_compute) | +| `dma_copy` | nki.isa | Copy data using DMA engines | [nki.isa.dma_copy](../programming/api/api-nki-isa-memory.md#nki-isa-dma_copy) | +| `dma_engine` | nki.isa | DMA engine enum (dma, gpsimd_dma) | [nki.isa.dma_engine](../programming/api/nki.isa.md) | +| `dma_transpose` | nki.isa | Transpose using DMA engines | [nki.isa.dma_transpose](../programming/api/api-nki-isa-memory.md#nki-isa-dma_transpose) | +| `dropout` | nki.isa | Apply dropout to tensor | [nki.isa.dropout](../programming/api/api-nki-isa-scalar.md#nki-isa-dropout) | +| `dropout` | nki.language | Randomly zeroes some of the elements of the input tile given a probability rate. | [nki.language.dropout](../programming/api/api-nki-language-misc.md#nki-language-dropout) | +| `ds` | nki.language | Dynamic slice for tensor indexing | [nki.language.ds](../programming/api/nki.language.md) | +| `dynamic_range` | nki.language | Create a sequence for **dynamic** loop iteration. | [nki.language.dynamic_range](../programming/api/api-nki-language-dims.md#nki-language-dynamic_range) | --- ## E -| Symbol | Module | Description | Documentation | -|--------|--------|-------------|---------------| -| `empty_like` | nki.language | Create a new tensor with the same shape and type as a given tensor. | [nki.language.empty_like](../programming/api/api-nki-language-creation.md#nki-language-empty_like) | -| `engine` | nki.isa | Neuron Device engine enum | [nki.isa.engine](../programming/api/nki.isa.md) | -| `equal` | nki.language | Op specifier for equal. | [nki.language.equal](../programming/api/api-nki-language-operators.md#nki-language-equal) | -| `erf` | nki.language | Op specifier for erf. | [nki.language.erf](../programming/api/api-nki-language-operators.md#nki-language-erf) | -| `erf_dx` | nki.language | Op specifier for erf_dx. | [nki.language.erf_dx](../programming/api/api-nki-language-operators.md#nki-language-erf_dx) | -| `exp` | nki.language | Op specifier for exp. | [nki.language.exp](../programming/api/api-nki-language-operators.md#nki-language-exp) | -| `expand_dim` | nki.tensor | Insert a new dimension of size 1 at position `dim`. | [NkiTensor.expand_dim](../programming/api/api-nki-tensor.md#nki-tensor-expand_dim) | -| `expand_dims` | nki.language | Expand the shape of a tile. | [nki.language.expand_dims](../programming/api/api-nki-language-misc.md#nki-language-expand_dims) | -| `exponential` | nki.isa | Dedicated exponential instruction (Trn3/NeuronCore-v4 only) | [nki.isa.exponential](../programming/api/nki.isa.md) | +| Symbol | Module | Description | Documentation | +| ------------- | ------------ | ------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------- | +| `empty_like` | nki.language | Create a new tensor with the same shape and type as a given tensor. | [nki.language.empty_like](../programming/api/api-nki-language-creation.md#nki-language-empty_like) | +| `engine` | nki.isa | Neuron Device engine enum | [nki.isa.engine](../programming/api/nki.isa.md) | +| `equal` | nki.language | Op specifier for equal. | [nki.language.equal](../programming/api/api-nki-language-operators.md#nki-language-equal) | +| `erf` | nki.language | Op specifier for erf. | [nki.language.erf](../programming/api/api-nki-language-operators.md#nki-language-erf) | +| `erf_dx` | nki.language | Op specifier for erf_dx. | [nki.language.erf_dx](../programming/api/api-nki-language-operators.md#nki-language-erf_dx) | +| `exp` | nki.language | Op specifier for exp. | [nki.language.exp](../programming/api/api-nki-language-operators.md#nki-language-exp) | +| `expand_dim` | nki.tensor | Insert a new dimension of size 1 at position `dim`. | [NkiTensor.expand_dim](../programming/api/api-nki-tensor.md#nki-tensor-expand_dim) | +| `expand_dims` | nki.language | Expand the shape of a tile. | [nki.language.expand_dims](../programming/api/api-nki-language-misc.md#nki-language-expand_dims) | +| `exponential` | nki.isa | Dedicated exponential instruction (Trn3/NeuronCore-v4 only) | [nki.isa.exponential](../programming/api/nki.isa.md) | --- ## F -| Symbol | Module | Description | Documentation | -|--------|--------|-------------|---------------| -| `flatten_dims` | nki.tensor | Merge a contiguous range of dimensions into one. | [NkiTensor.flatten_dims](../programming/api/api-nki-tensor.md#nki-tensor-flatten_dims) | -| `float16` | nki.language | FP16 data type | [nki.language.float16](../programming/api/api-nki-language-types.md#nki-language-float16) | -| `float32` | nki.language | FP32 data type | [nki.language.float32](../programming/api/api-nki-language-types.md#nki-language-float32) | -| `float4_e2m1fn_x4` | nki.language | 4x packed float4 for MXFP matmul | [nki.language.float4_e2m1fn_x4](../programming/api/api-nki-language-types.md#nki-language-float4_e2m1fn_x4) | -| `float8_e4m3` | nki.language | FP8 E4M3 data type | [nki.language.float8_e4m3](../programming/api/api-nki-language-types.md#nki-language-float8_e4m3) | -| `float8_e4m3fn` | nki.language | Data type constant `float8_e4m3fn` for tensor element types. | [nki.language.float8_e4m3fn](../programming/api/api-nki-language-types.md#nki-language-float8_e4m3fn) | -| `float8_e4m3fn_x4` | nki.language | 4x packed FP8 E4M3 for MXFP matmul | [nki.language.float8_e4m3fn_x4](../programming/api/api-nki-language-types.md#nki-language-float8_e4m3fn_x4) | -| `float8_e5m2` | nki.language | FP8 E5M2 data type | [nki.language.float8_e5m2](../programming/api/api-nki-language-types.md#nki-language-float8_e5m2) | -| `float8_e5m2_x4` | nki.language | 4x packed FP8 E5M2 for MXFP matmul | [nki.language.float8_e5m2_x4](../programming/api/api-nki-language-types.md#nki-language-float8_e5m2_x4) | -| `floor` | nki.language | Op specifier for floor. | [nki.language.floor](../programming/api/api-nki-language-operators.md#nki-language-floor) | -| `fmod` | nki.language | Op specifier for fmod. | [nki.language.fmod](../programming/api/api-nki-language-operators.md#nki-language-fmod) | -| `fori_loop` | nki.language | Structured for loop with dynamic bounds. | [nki.language.fori_loop](../programming/api/api-nki-language-dims.md#nki-language-fori_loop) | +| Symbol | Module | Description | Documentation | +| ------------------ | ------------ | ------------------------------------------------------------ | ----------------------------------------------------------------------------------------------------------- | +| `flatten_dims` | nki.tensor | Merge a contiguous range of dimensions into one. | [NkiTensor.flatten_dims](../programming/api/api-nki-tensor.md#nki-tensor-flatten_dims) | +| `float16` | nki.language | FP16 data type | [nki.language.float16](../programming/api/api-nki-language-types.md#nki-language-float16) | +| `float32` | nki.language | FP32 data type | [nki.language.float32](../programming/api/api-nki-language-types.md#nki-language-float32) | +| `float4_e2m1fn_x4` | nki.language | 4x packed float4 for MXFP matmul | [nki.language.float4_e2m1fn_x4](../programming/api/api-nki-language-types.md#nki-language-float4_e2m1fn_x4) | +| `float8_e4m3` | nki.language | FP8 E4M3 data type | [nki.language.float8_e4m3](../programming/api/api-nki-language-types.md#nki-language-float8_e4m3) | +| `float8_e4m3fn` | nki.language | Data type constant `float8_e4m3fn` for tensor element types. | [nki.language.float8_e4m3fn](../programming/api/api-nki-language-types.md#nki-language-float8_e4m3fn) | +| `float8_e4m3fn_x4` | nki.language | 4x packed FP8 E4M3 for MXFP matmul | [nki.language.float8_e4m3fn_x4](../programming/api/api-nki-language-types.md#nki-language-float8_e4m3fn_x4) | +| `float8_e5m2` | nki.language | FP8 E5M2 data type | [nki.language.float8_e5m2](../programming/api/api-nki-language-types.md#nki-language-float8_e5m2) | +| `float8_e5m2_x4` | nki.language | 4x packed FP8 E5M2 for MXFP matmul | [nki.language.float8_e5m2_x4](../programming/api/api-nki-language-types.md#nki-language-float8_e5m2_x4) | +| `floor` | nki.language | Op specifier for floor. | [nki.language.floor](../programming/api/api-nki-language-operators.md#nki-language-floor) | +| `fmod` | nki.language | Op specifier for fmod. | [nki.language.fmod](../programming/api/api-nki-language-operators.md#nki-language-fmod) | +| `fori_loop` | nki.language | Structured for loop with dynamic bounds. | [nki.language.fori_loop](../programming/api/api-nki-language-dims.md#nki-language-fori_loop) | --- ## G -| Symbol | Module | Description | Documentation | -|--------|--------|-------------|---------------| -| `gather_flattened` | nki.language | Gather elements from data tensor using indices after flattening. | [nki.language.gather_flattened](../programming/api/api-nki-language-creation.md#nki-language-gather_flattened) | -| `gelu` | nki.language | Op specifier for gelu. | [nki.language.gelu](../programming/api/api-nki-language-operators.md#nki-language-gelu) | -| `gelu_apprx_sigmoid` | nki.language | Op specifier for gelu_apprx_sigmoid. | [nki.language.gelu_apprx_sigmoid](../programming/api/api-nki-language-operators.md#nki-language-gelu_apprx_sigmoid) | -| `gelu_apprx_sigmoid_dx` | nki.language | Op specifier for gelu_apprx_sigmoid_dx. | [nki.language.gelu_apprx_sigmoid_dx](../programming/api/api-nki-language-operators.md#nki-language-gelu_apprx_sigmoid_dx) | -| `gelu_apprx_tanh` | nki.language | Op specifier for gelu_apprx_tanh. | [nki.language.gelu_apprx_tanh](../programming/api/api-nki-language-operators.md#nki-language-gelu_apprx_tanh) | -| `gelu_dx` | nki.language | Op specifier for gelu_dx. | [nki.language.gelu_dx](../programming/api/api-nki-language-operators.md#nki-language-gelu_dx) | -| `get_nc_version` | nki.isa | Get NeuronCore version | [nki.isa.get_nc_version](../programming/api/api-nki-isa-tensor.md#nki-isa-get_nc_version) | -| `get_pattern` | nki.tensor | Return the view's access pattern as `[[stride, count], . | [NkiTensor.get_pattern](../programming/api/api-nki-tensor.md#nki-tensor-get_pattern) | -| `greater` | nki.language | Op specifier for greater. | [nki.language.greater](../programming/api/api-nki-language-operators.md#nki-language-greater) | -| `greater_equal` | nki.language | Op specifier for greater_equal. | [nki.language.greater_equal](../programming/api/api-nki-language-operators.md#nki-language-greater_equal) | +| Symbol | Module | Description | Documentation | +| ----------------------- | ------------ | ---------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------- | +| `gather_flattened` | nki.language | Gather elements from data tensor using indices after flattening. | [nki.language.gather_flattened](../programming/api/api-nki-language-creation.md#nki-language-gather_flattened) | +| `gelu` | nki.language | Op specifier for gelu. | [nki.language.gelu](../programming/api/api-nki-language-operators.md#nki-language-gelu) | +| `gelu_apprx_sigmoid` | nki.language | Op specifier for gelu_apprx_sigmoid. | [nki.language.gelu_apprx_sigmoid](../programming/api/api-nki-language-operators.md#nki-language-gelu_apprx_sigmoid) | +| `gelu_apprx_sigmoid_dx` | nki.language | Op specifier for gelu_apprx_sigmoid_dx. | [nki.language.gelu_apprx_sigmoid_dx](../programming/api/api-nki-language-operators.md#nki-language-gelu_apprx_sigmoid_dx) | +| `gelu_apprx_tanh` | nki.language | Op specifier for gelu_apprx_tanh. | [nki.language.gelu_apprx_tanh](../programming/api/api-nki-language-operators.md#nki-language-gelu_apprx_tanh) | +| `gelu_dx` | nki.language | Op specifier for gelu_dx. | [nki.language.gelu_dx](../programming/api/api-nki-language-operators.md#nki-language-gelu_dx) | +| `get_nc_version` | nki.isa | Get NeuronCore version | [nki.isa.get_nc_version](../programming/api/api-nki-isa-tensor.md#nki-isa-get_nc_version) | +| `get_pattern` | nki.tensor | Return the view's access pattern as `[[stride, count], . | [NkiTensor.get_pattern](../programming/api/api-nki-tensor.md#nki-tensor-get_pattern) | +| `greater` | nki.language | Op specifier for greater. | [nki.language.greater](../programming/api/api-nki-language-operators.md#nki-language-greater) | +| `greater_equal` | nki.language | Op specifier for greater_equal. | [nki.language.greater_equal](../programming/api/api-nki-language-operators.md#nki-language-greater_equal) | --- ## H -| Symbol | Module | Description | Documentation | -|--------|--------|-------------|---------------| -| `hbm` | nki.language | HBM memory buffer (alias of private_hbm) | [nki.language.hbm](../programming/api/api-nki-language-memory.md#nki-language-hbm) | +| Symbol | Module | Description | Documentation | +| ------ | ------------ | ---------------------------------------- | ---------------------------------------------------------------------------------- | +| `hbm` | nki.language | HBM memory buffer (alias of private_hbm) | [nki.language.hbm](../programming/api/api-nki-language-memory.md#nki-language-hbm) | --- ## I -| Symbol | Module | Description | Documentation | -|--------|--------|-------------|---------------| -| `indirect` | nki.tensor | Create an indirect tensor view for Tensor Indirection (TI). | [NkiTensor.indirect](../programming/api/api-nki-tensor.md#nki-tensor-indirect) | -| `int16` | nki.language | 16-bit signed integer | [nki.language.int16](../programming/api/api-nki-language-types.md#nki-language-int16) | -| `int32` | nki.language | 32-bit signed integer | [nki.language.int32](../programming/api/api-nki-language-types.md#nki-language-int32) | -| `int8` | nki.language | 8-bit signed integer | [nki.language.int8](../programming/api/api-nki-language-types.md#nki-language-int8) | -| `invert` | nki.language | Op specifier for invert. | [nki.language.invert](../programming/api/api-nki-language-operators.md#nki-language-invert) | -| `iota` | nki.isa | Generate constant literal pattern | [nki.isa.iota](../programming/api/api-nki-isa-utility.md#nki-isa-iota) | -| `is_contiguous` | nki.tensor | Return True if the view covers storage contiguously (row-major order). | [NkiTensor.is_contiguous](../programming/api/api-nki-tensor.md#nki-tensor-is_contiguous) | -| `is_indirect` | nki.tensor | Return True if this view already uses indirect addressing. | [NkiTensor.is_indirect](../programming/api/api-nki-tensor.md#nki-tensor-is_indirect) | +| Symbol | Module | Description | Documentation | +| --------------- | ------------ | ---------------------------------------------------------------------- | ------------------------------------------------------------------------------------------- | +| `indirect` | nki.tensor | Create an indirect tensor view for Tensor Indirection (TI). | [NkiTensor.indirect](../programming/api/api-nki-tensor.md#nki-tensor-indirect) | +| `int16` | nki.language | 16-bit signed integer | [nki.language.int16](../programming/api/api-nki-language-types.md#nki-language-int16) | +| `int32` | nki.language | 32-bit signed integer | [nki.language.int32](../programming/api/api-nki-language-types.md#nki-language-int32) | +| `int8` | nki.language | 8-bit signed integer | [nki.language.int8](../programming/api/api-nki-language-types.md#nki-language-int8) | +| `invert` | nki.language | Op specifier for invert. | [nki.language.invert](../programming/api/api-nki-language-operators.md#nki-language-invert) | +| `iota` | nki.isa | Generate constant literal pattern | [nki.isa.iota](../programming/api/api-nki-isa-utility.md#nki-isa-iota) | +| `is_contiguous` | nki.tensor | Return True if the view covers storage contiguously (row-major order). | [NkiTensor.is_contiguous](../programming/api/api-nki-tensor.md#nki-tensor-is_contiguous) | +| `is_indirect` | nki.tensor | Return True if this view already uses indirect addressing. | [NkiTensor.is_indirect](../programming/api/api-nki-tensor.md#nki-tensor-is_indirect) | --- ## L -| Symbol | Module | Description | Documentation | -|--------|--------|-------------|---------------| -| `left_shift` | nki.language | Op specifier for left_shift. | [nki.language.left_shift](../programming/api/api-nki-language-operators.md#nki-language-left_shift) | -| `less` | nki.language | Op specifier for less. | [nki.language.less](../programming/api/api-nki-language-operators.md#nki-language-less) | -| `less_equal` | nki.language | Op specifier for less_equal. | [nki.language.less_equal](../programming/api/api-nki-language-operators.md#nki-language-less_equal) | -| `load` | nki.language | Load a tensor from device memory (HBM) into on-chip memory (SBUF). | [nki.language.load](../programming/api/api-nki-language-creation.md#nki-language-load) | +| Symbol | Module | Description | Documentation | +| ------------------ | ------------ | ------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------- | +| `left_shift` | nki.language | Op specifier for left_shift. | [nki.language.left_shift](../programming/api/api-nki-language-operators.md#nki-language-left_shift) | +| `less` | nki.language | Op specifier for less. | [nki.language.less](../programming/api/api-nki-language-operators.md#nki-language-less) | +| `less_equal` | nki.language | Op specifier for less_equal. | [nki.language.less_equal](../programming/api/api-nki-language-operators.md#nki-language-less_equal) | +| `load` | nki.language | Load a tensor from device memory (HBM) into on-chip memory (SBUF). | [nki.language.load](../programming/api/api-nki-language-creation.md#nki-language-load) | | `load_transpose2d` | nki.language | Load a tensor from device memory (HBM) and 2D-transpose the data before storing | [nki.language.load_transpose2d](../programming/api/api-nki-language-creation.md#nki-language-load_transpose2d) | -| `local_gather` | nki.isa | Gather SBUF data using indices | [nki.isa.local_gather](../programming/api/api-nki-isa-memory.md#nki-isa-local_gather) | -| `log` | nki.language | Op specifier for log. | [nki.language.log](../programming/api/api-nki-language-operators.md#nki-language-log) | -| `logical_and` | nki.language | Op specifier for logical_and. | [nki.language.logical_and](../programming/api/api-nki-language-operators.md#nki-language-logical_and) | -| `logical_not` | nki.language | Op specifier for logical_not. | [nki.language.logical_not](../programming/api/api-nki-language-operators.md#nki-language-logical_not) | -| `logical_or` | nki.language | Op specifier for logical_or. | [nki.language.logical_or](../programming/api/api-nki-language-operators.md#nki-language-logical_or) | -| `logical_xor` | nki.language | Op specifier for logical_xor. | [nki.language.logical_xor](../programming/api/api-nki-language-operators.md#nki-language-logical_xor) | +| `local_gather` | nki.isa | Gather SBUF data using indices | [nki.isa.local_gather](../programming/api/api-nki-isa-memory.md#nki-isa-local_gather) | +| `log` | nki.language | Op specifier for log. | [nki.language.log](../programming/api/api-nki-language-operators.md#nki-language-log) | +| `logical_and` | nki.language | Op specifier for logical_and. | [nki.language.logical_and](../programming/api/api-nki-language-operators.md#nki-language-logical_and) | +| `logical_not` | nki.language | Op specifier for logical_not. | [nki.language.logical_not](../programming/api/api-nki-language-operators.md#nki-language-logical_not) | +| `logical_or` | nki.language | Op specifier for logical_or. | [nki.language.logical_or](../programming/api/api-nki-language-operators.md#nki-language-logical_or) | +| `logical_xor` | nki.language | Op specifier for logical_xor. | [nki.language.logical_xor](../programming/api/api-nki-language-operators.md#nki-language-logical_xor) | --- ## M -| Symbol | Module | Description | Documentation | -|--------|--------|-------------|---------------| -| `matmul` | nki.language | x @ y matrix multiplication of x and y. | [nki.language.matmul](../programming/api/api-nki-language-misc.md#nki-language-matmul) | -| `max` | nki.language | Maximum of elements along the specified axis (or axes) of the input. | [nki.language.max](../programming/api/api-nki-language-misc.md#nki-language-max) | -| `max8` | nki.isa | Find 8 largest values per partition | [nki.isa.max8](../programming/api/api-nki-isa-utility.md#nki-isa-max8) | -| `maximum` | nki.language | Op specifier for maximum. | [nki.language.maximum](../programming/api/api-nki-language-operators.md#nki-language-maximum) | -| `mean` | nki.language | Arithmetic mean along the specified axis (or axes) of the input. | [nki.language.mean](../programming/api/api-nki-language-misc.md#nki-language-mean) | -| `memset` | nki.isa | Initialize tensor with constant value | [nki.isa.memset](../programming/api/api-nki-isa-memory.md#nki-isa-memset) | -| `min` | nki.language | Minimum of elements along the specified axis (or axes) of the input. | [nki.language.min](../programming/api/api-nki-language-misc.md#nki-language-min) | -| `minimum` | nki.language | Op specifier for minimum. | [nki.language.minimum](../programming/api/api-nki-language-operators.md#nki-language-minimum) | -| `mish` | nki.language | Op specifier for mish. | [nki.language.mish](../programming/api/api-nki-language-operators.md#nki-language-mish) | -| `mod` | nki.language | Op specifier for mod. | [nki.language.mod](../programming/api/api-nki-language-operators.md#nki-language-mod) | -| `multiply` | nki.language | Op specifier for multiply. | [nki.language.multiply](../programming/api/api-nki-language-operators.md#nki-language-multiply) | +| Symbol | Module | Description | Documentation | +| ---------- | ------------ | -------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------- | +| `matmul` | nki.language | x @ y matrix multiplication of x and y. | [nki.language.matmul](../programming/api/api-nki-language-misc.md#nki-language-matmul) | +| `max` | nki.language | Maximum of elements along the specified axis (or axes) of the input. | [nki.language.max](../programming/api/api-nki-language-misc.md#nki-language-max) | +| `max8` | nki.isa | Find 8 largest values per partition | [nki.isa.max8](../programming/api/api-nki-isa-utility.md#nki-isa-max8) | +| `maximum` | nki.language | Op specifier for maximum. | [nki.language.maximum](../programming/api/api-nki-language-operators.md#nki-language-maximum) | +| `mean` | nki.language | Arithmetic mean along the specified axis (or axes) of the input. | [nki.language.mean](../programming/api/api-nki-language-misc.md#nki-language-mean) | +| `memset` | nki.isa | Initialize tensor with constant value | [nki.isa.memset](../programming/api/api-nki-isa-memory.md#nki-isa-memset) | +| `min` | nki.language | Minimum of elements along the specified axis (or axes) of the input. | [nki.language.min](../programming/api/api-nki-language-misc.md#nki-language-min) | +| `minimum` | nki.language | Op specifier for minimum. | [nki.language.minimum](../programming/api/api-nki-language-operators.md#nki-language-minimum) | +| `mish` | nki.language | Op specifier for mish. | [nki.language.mish](../programming/api/api-nki-language-operators.md#nki-language-mish) | +| `mod` | nki.language | Op specifier for mod. | [nki.language.mod](../programming/api/api-nki-language-operators.md#nki-language-mod) | +| `multiply` | nki.language | Op specifier for multiply. | [nki.language.multiply](../programming/api/api-nki-language-operators.md#nki-language-multiply) | --- ## N -| Symbol | Module | Description | Documentation | -|--------|--------|-------------|---------------| -| `nc_find_index8` | nki.isa | Find indices of 8 values in data | [nki.isa.nc_find_index8](../programming/api/api-nki-isa-tensor.md#nki-isa-nc_find_index8) | -| `nc_match_replace8` | nki.isa | Replace values and optionally return indices | [nki.isa.nc_match_replace8](../programming/api/api-nki-isa-tensor.md#nki-isa-nc_match_replace8) | -| `nc_matmul` | nki.isa | Matrix multiplication on Tensor Engine | [nki.isa.nc_matmul](../programming/api/api-nki-isa-tensor.md#nki-isa-nc_matmul) | -| `nc_matmul_mx` | nki.isa | MXFP quantized matrix multiplication | [nki.isa.nc_matmul_mx](../programming/api/api-nki-isa-tensor.md#nki-isa-nc_matmul_mx) | -| `nc_n_gather` | nki.isa | Gather elements using indices | [nki.isa.nc_n_gather](../programming/api/api-nki-isa-tensor.md#nki-isa-nc_n_gather) | -| `nc_stream_shuffle` | nki.isa | Cross-partition data shuffle | [nki.isa.nc_stream_shuffle](../programming/api/api-nki-isa-tensor.md#nki-isa-nc_stream_shuffle) | -| `nc_transpose` | nki.isa | 2D transpose between P and F axes | [nki.isa.nc_transpose](../programming/api/api-nki-isa-tensor.md#nki-isa-nc_transpose) | -| `nc_version` | nki.isa | NeuronCore version enum | [nki.isa.nc_version](../programming/api/api-nki-isa-tensor.md#nki-isa-nc_version) | -| `ndarray` | nki.language | Create tensor on specified buffer | [nki.language.ndarray](../programming/api/nki.language.md) | -| `negative` | nki.language | Op specifier for negative. | [nki.language.negative](../programming/api/api-nki-language-operators.md#nki-language-negative) | -| `nonzero_with_count` | nki.isa | Find indices of nonzero elements and count (NeuronCore-v3+) | [nki.isa.nonzero_with_count](../programming/api/api-nki-isa-misc.md#nki-isa-nonzero_with_count) | -| `not_equal` | nki.language | Op specifier for not_equal. | [nki.language.not_equal](../programming/api/api-nki-language-operators.md#nki-language-not_equal) | -| `num_programs` | nki.language | Number of SPMD programs in grid | [nki.language.num_programs](../programming/api/nki.language.md) | +| Symbol | Module | Description | Documentation | +| -------------------- | ------------ | ----------------------------------------------------------- | ------------------------------------------------------------------------------------------------- | +| `nc_find_index8` | nki.isa | Find indices of 8 values in data | [nki.isa.nc_find_index8](../programming/api/api-nki-isa-tensor.md#nki-isa-nc_find_index8) | +| `nc_match_replace8` | nki.isa | Replace values and optionally return indices | [nki.isa.nc_match_replace8](../programming/api/api-nki-isa-tensor.md#nki-isa-nc_match_replace8) | +| `nc_matmul` | nki.isa | Matrix multiplication on Tensor Engine | [nki.isa.nc_matmul](../programming/api/api-nki-isa-tensor.md#nki-isa-nc_matmul) | +| `nc_matmul_mx` | nki.isa | MXFP quantized matrix multiplication | [nki.isa.nc_matmul_mx](../programming/api/api-nki-isa-tensor.md#nki-isa-nc_matmul_mx) | +| `nc_n_gather` | nki.isa | Gather elements using indices | [nki.isa.nc_n_gather](../programming/api/api-nki-isa-tensor.md#nki-isa-nc_n_gather) | +| `nc_stream_shuffle` | nki.isa | Cross-partition data shuffle | [nki.isa.nc_stream_shuffle](../programming/api/api-nki-isa-tensor.md#nki-isa-nc_stream_shuffle) | +| `nc_transpose` | nki.isa | 2D transpose between P and F axes | [nki.isa.nc_transpose](../programming/api/api-nki-isa-tensor.md#nki-isa-nc_transpose) | +| `nc_version` | nki.isa | NeuronCore version enum | [nki.isa.nc_version](../programming/api/api-nki-isa-tensor.md#nki-isa-nc_version) | +| `ndarray` | nki.language | Create tensor on specified buffer | [nki.language.ndarray](../programming/api/nki.language.md) | +| `negative` | nki.language | Op specifier for negative. | [nki.language.negative](../programming/api/api-nki-language-operators.md#nki-language-negative) | +| `nonzero_with_count` | nki.isa | Find indices of nonzero elements and count (NeuronCore-v3+) | [nki.isa.nonzero_with_count](../programming/api/api-nki-isa-misc.md#nki-isa-nonzero_with_count) | +| `not_equal` | nki.language | Op specifier for not_equal. | [nki.language.not_equal](../programming/api/api-nki-language-operators.md#nki-language-not_equal) | +| `num_programs` | nki.language | Number of SPMD programs in grid | [nki.language.num_programs](../programming/api/nki.language.md) | --- ## O -| Symbol | Module | Description | Documentation | -|--------|--------|-------------|---------------| +| Symbol | Module | Description | Documentation | +| ------ | ------------ | -------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------- | | `ones` | nki.language | Create a new tensor of given shape and dtype on the specified buffer, filled wit | [nki.language.ones](../programming/api/api-nki-language-creation.md#nki-language-ones) | --- ## P -| Symbol | Module | Description | Documentation | -|--------|--------|-------------|---------------| -| `permute` | nki.tensor | Reorder tensor dimensions. | [NkiTensor.permute](../programming/api/api-nki-tensor.md#nki-tensor-permute) | -| `power` | nki.language | Op specifier for power. | [nki.language.power](../programming/api/api-nki-language-operators.md#nki-language-power) | -| `prelu` | nki.language | Op specifier for prelu. | [nki.language.prelu](../programming/api/api-nki-language-operators.md#nki-language-prelu) | -| `private_hbm` | nki.language | Private HBM memory buffer | [nki.language.private_hbm](../programming/api/api-nki-language-memory.md#nki-language-private_hbm) | -| `prod` | nki.language | Product of elements along the specified axis (or axes) of the input. | [nki.language.prod](../programming/api/api-nki-language-misc.md#nki-language-prod) | -| `program_id` | nki.language | Index of current SPMD program | [nki.language.program_id](../programming/api/nki.language.md) | -| `program_ndim` | nki.language | Number of dimensions in SPMD grid | [nki.language.program_ndim](../programming/api/nki.language.md) | -| `psum` | nki.language | PSUM memory buffer | [nki.language.psum](../programming/api/api-nki-language-memory.md#nki-language-psum) | +| Symbol | Module | Description | Documentation | +| -------------- | ------------ | -------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------- | +| `permute` | nki.tensor | Reorder tensor dimensions. | [NkiTensor.permute](../programming/api/api-nki-tensor.md#nki-tensor-permute) | +| `power` | nki.language | Op specifier for power. | [nki.language.power](../programming/api/api-nki-language-operators.md#nki-language-power) | +| `prelu` | nki.language | Op specifier for prelu. | [nki.language.prelu](../programming/api/api-nki-language-operators.md#nki-language-prelu) | +| `private_hbm` | nki.language | Private HBM memory buffer | [nki.language.private_hbm](../programming/api/api-nki-language-memory.md#nki-language-private_hbm) | +| `prod` | nki.language | Product of elements along the specified axis (or axes) of the input. | [nki.language.prod](../programming/api/api-nki-language-misc.md#nki-language-prod) | +| `program_id` | nki.language | Index of current SPMD program | [nki.language.program_id](../programming/api/nki.language.md) | +| `program_ndim` | nki.language | Number of dimensions in SPMD grid | [nki.language.program_ndim](../programming/api/nki.language.md) | +| `psum` | nki.language | PSUM memory buffer | [nki.language.psum](../programming/api/api-nki-language-memory.md#nki-language-psum) | --- ## Q -| Symbol | Module | Description | Documentation | -|--------|--------|-------------|---------------| +| Symbol | Module | Description | Documentation | +| ------------- | ------- | ------------------------ | ----------------------------------------------------------------------------------- | | `quantize_mx` | nki.isa | Quantize to MXFP8 format | [nki.isa.quantize_mx](../programming/api/api-nki-isa-tensor.md#nki-isa-quantize_mx) | --- ## R -| Symbol | Module | Description | Documentation | -|--------|--------|-------------|---------------| -| `rand` | nki.language | Create a new tensor of given shape and dtype on the specified buffer, filled wit | [nki.language.rand](../programming/api/api-nki-language-creation.md#nki-language-rand) | -| `rand2` | nki.isa | Generate uniform random numbers | [nki.isa.rand2](../programming/api/nki.isa.md) | -| `rand_get_state` | nki.isa | Get PRNG state from engine | [nki.isa.rand_get_state](../programming/api/nki.isa.md) | -| `rand_set_state` | nki.isa | Set PRNG state in engine | [nki.isa.rand_set_state](../programming/api/nki.isa.md) | -| `range_select` | nki.isa | Select elements based on range comparison | [nki.isa.range_select](../programming/api/api-nki-isa-utility.md#nki-isa-range_select) | -| `rank_id` | nki.collectives | Get the rank ID of the current rank. | [nki.collectives.rank_id](../programming/api/api-nki-collectives.md#nki-collectives-rank_id) | -| `rearrange` | nki.tensor | Rearrange tensor dimensions using einops-style patterns. | [NkiTensor.rearrange](../programming/api/api-nki-tensor.md#nki-tensor-rearrange) | -| `reciprocal` | nki.isa | Compute element-wise 1/x | [nki.isa.reciprocal](../programming/api/api-nki-isa-scalar.md#nki-isa-reciprocal) | -| `reciprocal` | nki.language | Op specifier for reciprocal. | [nki.language.reciprocal](../programming/api/api-nki-language-operators.md#nki-language-reciprocal) | -| `reduce_cmd` | nki.isa | Engine register reduce commands enum | [nki.isa.reduce_cmd](../programming/api/nki.isa.md) | -| `reduce_scatter` | nki.collectives | Perform a reduce-scatter on the given replica group and input/output tensors. | [nki.collectives.reduce_scatter](../programming/api/api-nki-collectives.md#nki-collectives-reduce_scatter) | -| `register_alloc` | nki.isa | Allocate virtual register | [nki.isa.register_alloc](../programming/api/nki.isa.md) | -| `register_load` | nki.isa | Load scalar from memory to register | [nki.isa.register_load](../programming/api/nki.isa.md) | -| `register_move` | nki.isa | Move value from source register to destination register | [nki.isa.register_move](../programming/api/nki.isa.md) | -| `register_store` | nki.isa | Store register value to memory | [nki.isa.register_store](../programming/api/nki.isa.md) | -| `relu` | nki.language | Op specifier for relu. | [nki.language.relu](../programming/api/api-nki-language-operators.md#nki-language-relu) | -| `reshape` | nki.tensor | Reshape the tensor to a new shape without copying data. | [NkiTensor.reshape](../programming/api/api-nki-tensor.md#nki-tensor-reshape) | -| `reshape_dim` | nki.tensor | Split a single dimension into multiple dimensions. | [NkiTensor.reshape_dim](../programming/api/api-nki-tensor.md#nki-tensor-reshape_dim) | -| `right_shift` | nki.language | Op specifier for right_shift. | [nki.language.right_shift](../programming/api/api-nki-language-operators.md#nki-language-right_shift) | -| `rms_norm` | nki.language | Apply Root Mean Square Layer Normalization. | [nki.language.rms_norm](../programming/api/api-nki-language-misc.md#nki-language-rms_norm) | -| `rng` | nki.isa | Generate pseudo random numbers | [nki.isa.rng](../programming/api/nki.isa.md) | -| `rsqrt` | nki.language | Op specifier for rsqrt. | [nki.language.rsqrt](../programming/api/api-nki-language-operators.md#nki-language-rsqrt) | +| Symbol | Module | Description | Documentation | +| ---------------- | --------------- | -------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------- | +| `rand` | nki.language | Create a new tensor of given shape and dtype on the specified buffer, filled wit | [nki.language.rand](../programming/api/api-nki-language-creation.md#nki-language-rand) | +| `rand2` | nki.isa | Generate uniform random numbers | [nki.isa.rand2](../programming/api/nki.isa.md) | +| `rand_get_state` | nki.isa | Get PRNG state from engine | [nki.isa.rand_get_state](../programming/api/nki.isa.md) | +| `rand_set_state` | nki.isa | Set PRNG state in engine | [nki.isa.rand_set_state](../programming/api/nki.isa.md) | +| `range_select` | nki.isa | Select elements based on range comparison | [nki.isa.range_select](../programming/api/api-nki-isa-utility.md#nki-isa-range_select) | +| `rank_id` | nki.collectives | Get the rank ID of the current rank. | [nki.collectives.rank_id](../programming/api/api-nki-collectives.md#nki-collectives-rank_id) | +| `rearrange` | nki.tensor | Rearrange tensor dimensions using einops-style patterns. | [NkiTensor.rearrange](../programming/api/api-nki-tensor.md#nki-tensor-rearrange) | +| `reciprocal` | nki.isa | Compute element-wise 1/x | [nki.isa.reciprocal](../programming/api/api-nki-isa-scalar.md#nki-isa-reciprocal) | +| `reciprocal` | nki.language | Op specifier for reciprocal. | [nki.language.reciprocal](../programming/api/api-nki-language-operators.md#nki-language-reciprocal) | +| `reduce_cmd` | nki.isa | Engine register reduce commands enum | [nki.isa.reduce_cmd](../programming/api/nki.isa.md) | +| `reduce_scatter` | nki.collectives | Perform a reduce-scatter on the given replica group and input/output tensors. | [nki.collectives.reduce_scatter](../programming/api/api-nki-collectives.md#nki-collectives-reduce_scatter) | +| `register_alloc` | nki.isa | Allocate virtual register | [nki.isa.register_alloc](../programming/api/nki.isa.md) | +| `register_load` | nki.isa | Load scalar from memory to register | [nki.isa.register_load](../programming/api/nki.isa.md) | +| `register_move` | nki.isa | Move value from source register to destination register | [nki.isa.register_move](../programming/api/nki.isa.md) | +| `register_store` | nki.isa | Store register value to memory | [nki.isa.register_store](../programming/api/nki.isa.md) | +| `relu` | nki.language | Op specifier for relu. | [nki.language.relu](../programming/api/api-nki-language-operators.md#nki-language-relu) | +| `reshape` | nki.tensor | Reshape the tensor to a new shape without copying data. | [NkiTensor.reshape](../programming/api/api-nki-tensor.md#nki-tensor-reshape) | +| `reshape_dim` | nki.tensor | Split a single dimension into multiple dimensions. | [NkiTensor.reshape_dim](../programming/api/api-nki-tensor.md#nki-tensor-reshape_dim) | +| `right_shift` | nki.language | Op specifier for right_shift. | [nki.language.right_shift](../programming/api/api-nki-language-operators.md#nki-language-right_shift) | +| `rms_norm` | nki.language | Apply Root Mean Square Layer Normalization. | [nki.language.rms_norm](../programming/api/api-nki-language-misc.md#nki-language-rms_norm) | +| `rng` | nki.isa | Generate pseudo random numbers | [nki.isa.rng](../programming/api/nki.isa.md) | +| `rsqrt` | nki.language | Op specifier for rsqrt. | [nki.language.rsqrt](../programming/api/api-nki-language-operators.md#nki-language-rsqrt) | --- ## S -| Symbol | Module | Description | Documentation | -|--------|--------|-------------|---------------| -| `sbuf` | nki.language | State Buffer memory | [nki.language.sbuf](../programming/api/api-nki-language-memory.md#nki-language-sbuf) | -| `scalar_tensor_tensor` | nki.isa | Two-op sequence with scalar broadcast | [nki.isa.scalar_tensor_tensor](../programming/api/api-nki-isa-tensor.md#nki-isa-scalar_tensor_tensor) | -| `select` | nki.tensor | Select a single element along a dimension, removing it. | [NkiTensor.select](../programming/api/api-nki-tensor.md#nki-tensor-select) | -| `select_reduce` | nki.isa | Conditional copy with optional reduction | [nki.isa.select_reduce](../programming/api/api-nki-isa-utility.md#nki-isa-select_reduce) | -| `sendrecv` | nki.isa | Point-to-point NeuronCore communication | [nki.isa.sendrecv](../programming/api/nki.isa.md) | -| `sequence_bounds` | nki.isa | Compute sequence bounds from segment IDs | [nki.isa.sequence_bounds](../programming/api/api-nki-isa-utility.md#nki-isa-sequence_bounds) | -| `sequential_range` | nki.language | Loop iterator (legacy alias for `range`) | [nki.language.sequential_range](../programming/api/nki.language.md) | -| `set_rng_seed` | nki.isa | Seed Vector Engine PRNG | [nki.isa.set_rng_seed](../programming/api/nki.isa.md) | -| `shared_constant` | nki.language | Create a tensor in shared HBM initialized with constant data. | [nki.language.shared_constant](../programming/api/api-nki-language-creation.md#nki-language-shared_constant) | -| `shared_hbm` | nki.language | Shared HBM across kernel instances | [nki.language.shared_hbm](../programming/api/api-nki-language-memory.md#nki-language-shared_hbm) | -| `sigmoid` | nki.language | Op specifier for sigmoid. | [nki.language.sigmoid](../programming/api/api-nki-language-operators.md#nki-language-sigmoid) | -| `sign` | nki.language | Op specifier for sign. | [nki.language.sign](../programming/api/api-nki-language-operators.md#nki-language-sign) | -| `silu` | nki.language | Op specifier for silu. | [nki.language.silu](../programming/api/api-nki-language-operators.md#nki-language-silu) | -| `silu_dx` | nki.language | Op specifier for silu_dx. | [nki.language.silu_dx](../programming/api/api-nki-language-operators.md#nki-language-silu_dx) | -| `simulate` | nki | Run NKI kernel on CPU without NeuronDevice (experimental) | [nki.simulate](../programming/api/api-nki-tools.md#nki-simulate) | -| `sin` | nki.language | Op specifier for sin. | [nki.language.sin](../programming/api/api-nki-language-operators.md#nki-language-sin) | -| `slice` | nki.tensor | Slice along a single dimension. | [NkiTensor.slice](../programming/api/api-nki-tensor.md#nki-tensor-slice) | -| `softmax` | nki.language | Softmax activation function on the input, element-wise. | [nki.language.softmax](../programming/api/api-nki-language-misc.md#nki-language-softmax) | -| `softplus` | nki.language | Op specifier for softplus. | [nki.language.softplus](../programming/api/api-nki-language-operators.md#nki-language-softplus) | -| `sqrt` | nki.language | Op specifier for sqrt. | [nki.language.sqrt](../programming/api/api-nki-language-operators.md#nki-language-sqrt) | -| `square` | nki.language | Op specifier for square. | [nki.language.square](../programming/api/api-nki-language-operators.md#nki-language-square) | -| `squeeze_dim` | nki.tensor | Remove a dimension of size 1. | [NkiTensor.squeeze_dim](../programming/api/api-nki-tensor.md#nki-tensor-squeeze_dim) | -| `static_range` | nki.language | Loop iterator (legacy alias for `range`) | [nki.language.static_range](../programming/api/nki.language.md) | -| `store` | nki.language | Store into a tensor on device memory (HBM) from on-chip memory (SBUF). | [nki.language.store](../programming/api/api-nki-language-creation.md#nki-language-store) | -| `subtract` | nki.language | Op specifier for subtract. | [nki.language.subtract](../programming/api/api-nki-language-operators.md#nki-language-subtract) | -| `sum` | nki.language | Sum of elements along the specified axis (or axes) of the input. | [nki.language.sum](../programming/api/api-nki-language-misc.md#nki-language-sum) | +| Symbol | Module | Description | Documentation | +| ---------------------- | ------------ | ---------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------ | +| `sbuf` | nki.language | State Buffer memory | [nki.language.sbuf](../programming/api/api-nki-language-memory.md#nki-language-sbuf) | +| `scalar_tensor_tensor` | nki.isa | Two-op sequence with scalar broadcast | [nki.isa.scalar_tensor_tensor](../programming/api/api-nki-isa-tensor.md#nki-isa-scalar_tensor_tensor) | +| `select` | nki.tensor | Select a single element along a dimension, removing it. | [NkiTensor.select](../programming/api/api-nki-tensor.md#nki-tensor-select) | +| `select_reduce` | nki.isa | Conditional copy with optional reduction | [nki.isa.select_reduce](../programming/api/api-nki-isa-utility.md#nki-isa-select_reduce) | +| `sendrecv` | nki.isa | Point-to-point NeuronCore communication | [nki.isa.sendrecv](../programming/api/nki.isa.md) | +| `sequence_bounds` | nki.isa | Compute sequence bounds from segment IDs | [nki.isa.sequence_bounds](../programming/api/api-nki-isa-utility.md#nki-isa-sequence_bounds) | +| `sequential_range` | nki.language | Loop iterator (legacy alias for `range`) | [nki.language.sequential_range](../programming/api/nki.language.md) | +| `set_rng_seed` | nki.isa | Seed Vector Engine PRNG | [nki.isa.set_rng_seed](../programming/api/nki.isa.md) | +| `shared_constant` | nki.language | Create a tensor in shared HBM initialized with constant data. | [nki.language.shared_constant](../programming/api/api-nki-language-creation.md#nki-language-shared_constant) | +| `shared_hbm` | nki.language | Shared HBM across kernel instances | [nki.language.shared_hbm](../programming/api/api-nki-language-memory.md#nki-language-shared_hbm) | +| `sigmoid` | nki.language | Op specifier for sigmoid. | [nki.language.sigmoid](../programming/api/api-nki-language-operators.md#nki-language-sigmoid) | +| `sign` | nki.language | Op specifier for sign. | [nki.language.sign](../programming/api/api-nki-language-operators.md#nki-language-sign) | +| `silu` | nki.language | Op specifier for silu. | [nki.language.silu](../programming/api/api-nki-language-operators.md#nki-language-silu) | +| `silu_dx` | nki.language | Op specifier for silu_dx. | [nki.language.silu_dx](../programming/api/api-nki-language-operators.md#nki-language-silu_dx) | +| `simulate` | nki | Run NKI kernel on CPU without NeuronDevice (experimental) | [nki.simulate](../programming/api/api-nki-tools.md#nki-simulate) | +| `sin` | nki.language | Op specifier for sin. | [nki.language.sin](../programming/api/api-nki-language-operators.md#nki-language-sin) | +| `slice` | nki.tensor | Slice along a single dimension. | [NkiTensor.slice](../programming/api/api-nki-tensor.md#nki-tensor-slice) | +| `softmax` | nki.language | Softmax activation function on the input, element-wise. | [nki.language.softmax](../programming/api/api-nki-language-misc.md#nki-language-softmax) | +| `softplus` | nki.language | Op specifier for softplus. | [nki.language.softplus](../programming/api/api-nki-language-operators.md#nki-language-softplus) | +| `sqrt` | nki.language | Op specifier for sqrt. | [nki.language.sqrt](../programming/api/api-nki-language-operators.md#nki-language-sqrt) | +| `square` | nki.language | Op specifier for square. | [nki.language.square](../programming/api/api-nki-language-operators.md#nki-language-square) | +| `squeeze_dim` | nki.tensor | Remove a dimension of size 1. | [NkiTensor.squeeze_dim](../programming/api/api-nki-tensor.md#nki-tensor-squeeze_dim) | +| `static_range` | nki.language | Loop iterator (legacy alias for `range`) | [nki.language.static_range](../programming/api/nki.language.md) | +| `store` | nki.language | Store into a tensor on device memory (HBM) from on-chip memory (SBUF). | [nki.language.store](../programming/api/api-nki-language-creation.md#nki-language-store) | +| `subtract` | nki.language | Op specifier for subtract. | [nki.language.subtract](../programming/api/api-nki-language-operators.md#nki-language-subtract) | +| `sum` | nki.language | Sum of elements along the specified axis (or axes) of the input. | [nki.language.sum](../programming/api/api-nki-language-misc.md#nki-language-sum) | --- ## T -| Symbol | Module | Description | Documentation | -|--------|--------|-------------|---------------| -| `tan` | nki.language | Op specifier for tan. | [nki.language.tan](../programming/api/api-nki-language-operators.md#nki-language-tan) | -| `tanh` | nki.language | Op specifier for tanh. | [nki.language.tanh](../programming/api/api-nki-language-operators.md#nki-language-tanh) | -| `tensor_copy` | nki.isa | Copy tensor within on-chip SRAM | [nki.isa.tensor_copy](../programming/api/api-nki-isa-tensor.md#nki-isa-tensor_copy) | -| `tensor_copy_predicated` | nki.isa | Conditional element copy | [nki.isa.tensor_copy_predicated](../programming/api/api-nki-isa-tensor.md#nki-isa-tensor_copy_predicated) | -| `tensor_partition_reduce` | nki.isa | Reduce across partitions | [nki.isa.tensor_partition_reduce](../programming/api/api-nki-isa-tensor.md#nki-isa-tensor_partition_reduce) | -| `tensor_reduce` | nki.isa | Reduce along free axes | [nki.isa.tensor_reduce](../programming/api/api-nki-isa-tensor.md#nki-isa-tensor_reduce) | -| `tensor_scalar` | nki.isa | Tensor-scalar operations with broadcasting | [nki.isa.tensor_scalar](../programming/api/api-nki-isa-tensor.md#nki-isa-tensor_scalar) | -| `tensor_scalar_cumulative` | nki.isa | Tensor-scalar with cumulative reduction | [nki.isa.tensor_scalar_cumulative](../programming/api/api-nki-isa-tensor.md#nki-isa-tensor_scalar_cumulative) | -| `tensor_scalar_reduce` | nki.isa | Tensor-scalar with free-dim reduction | [nki.isa.tensor_scalar_reduce](../programming/api/api-nki-isa-tensor.md#nki-isa-tensor_scalar_reduce) | -| `tensor_tensor` | nki.isa | Element-wise operation on two tensors | [nki.isa.tensor_tensor](../programming/api/api-nki-isa-tensor.md#nki-isa-tensor_tensor) | -| `tensor_tensor_scan` | nki.isa | Scan operation on two tensors | [nki.isa.tensor_tensor_scan](../programming/api/api-nki-isa-tensor.md#nki-isa-tensor_tensor_scan) | -| `tfloat32` | nki.language | TF32 data type (1S,8E,10M) | [nki.language.tfloat32](../programming/api/api-nki-language-types.md#nki-language-tfloat32) | -| `tile_size` | nki.language | Tile size constants | [nki.language.tile_size](../programming/api/nki.language.md) | -| `topk` | nki.isa | Find the K largest values and their indices from a source tile using GpSIMD Engi | [nki.isa.topk](../programming/api/api-nki-isa-tensor.md#nki-isa-topk) | -| `transpose` | nki.language | Transposes a 2D tile between its partition and free dimension. | [nki.language.transpose](../programming/api/api-nki-language-misc.md#nki-language-transpose) | -| `trunc` | nki.language | Op specifier for trunc. | [nki.language.trunc](../programming/api/api-nki-language-operators.md#nki-language-trunc) | +| Symbol | Module | Description | Documentation | +| -------------------------- | ------------ | -------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------- | +| `tan` | nki.language | Op specifier for tan. | [nki.language.tan](../programming/api/api-nki-language-operators.md#nki-language-tan) | +| `tanh` | nki.language | Op specifier for tanh. | [nki.language.tanh](../programming/api/api-nki-language-operators.md#nki-language-tanh) | +| `tensor_copy` | nki.isa | Copy tensor within on-chip SRAM | [nki.isa.tensor_copy](../programming/api/api-nki-isa-tensor.md#nki-isa-tensor_copy) | +| `tensor_copy_predicated` | nki.isa | Conditional element copy | [nki.isa.tensor_copy_predicated](../programming/api/api-nki-isa-tensor.md#nki-isa-tensor_copy_predicated) | +| `tensor_partition_reduce` | nki.isa | Reduce across partitions | [nki.isa.tensor_partition_reduce](../programming/api/api-nki-isa-tensor.md#nki-isa-tensor_partition_reduce) | +| `tensor_reduce` | nki.isa | Reduce along free axes | [nki.isa.tensor_reduce](../programming/api/api-nki-isa-tensor.md#nki-isa-tensor_reduce) | +| `tensor_scalar` | nki.isa | Tensor-scalar operations with broadcasting | [nki.isa.tensor_scalar](../programming/api/api-nki-isa-tensor.md#nki-isa-tensor_scalar) | +| `tensor_scalar_cumulative` | nki.isa | Tensor-scalar with cumulative reduction | [nki.isa.tensor_scalar_cumulative](../programming/api/api-nki-isa-tensor.md#nki-isa-tensor_scalar_cumulative) | +| `tensor_scalar_reduce` | nki.isa | Tensor-scalar with free-dim reduction | [nki.isa.tensor_scalar_reduce](../programming/api/api-nki-isa-tensor.md#nki-isa-tensor_scalar_reduce) | +| `tensor_tensor` | nki.isa | Element-wise operation on two tensors | [nki.isa.tensor_tensor](../programming/api/api-nki-isa-tensor.md#nki-isa-tensor_tensor) | +| `tensor_tensor_scan` | nki.isa | Scan operation on two tensors | [nki.isa.tensor_tensor_scan](../programming/api/api-nki-isa-tensor.md#nki-isa-tensor_tensor_scan) | +| `tfloat32` | nki.language | TF32 data type (1S,8E,10M) | [nki.language.tfloat32](../programming/api/api-nki-language-types.md#nki-language-tfloat32) | +| `tile_size` | nki.language | Tile size constants | [nki.language.tile_size](../programming/api/nki.language.md) | +| `topk` | nki.isa | Find the K largest values and their indices from a source tile using GpSIMD Engi | [nki.isa.topk](../programming/api/api-nki-isa-tensor.md#nki-isa-topk) | +| `transpose` | nki.language | Transposes a 2D tile between its partition and free dimension. | [nki.language.transpose](../programming/api/api-nki-language-misc.md#nki-language-transpose) | +| `trunc` | nki.language | Op specifier for trunc. | [nki.language.trunc](../programming/api/api-nki-language-operators.md#nki-language-trunc) | --- ## U -| Symbol | Module | Description | Documentation | -|--------|--------|-------------|---------------| +| Symbol | Module | Description | Documentation | +| -------- | ------------ | ----------------------- | --------------------------------------------------------------------------------------- | | `uint16` | nki.language | 16-bit unsigned integer | [nki.language.uint16](../programming/api/api-nki-language-types.md#nki-language-uint16) | | `uint32` | nki.language | 32-bit unsigned integer | [nki.language.uint32](../programming/api/api-nki-language-types.md#nki-language-uint32) | -| `uint8` | nki.language | 8-bit unsigned integer | [nki.language.uint8](../programming/api/api-nki-language-types.md#nki-language-uint8) | +| `uint8` | nki.language | 8-bit unsigned integer | [nki.language.uint8](../programming/api/api-nki-language-types.md#nki-language-uint8) | --- ## V -| Symbol | Module | Description | Documentation | -|--------|--------|-------------|---------------| -| `var` | nki.language | Variance along the specified axis (or axes) of the input. | [nki.language.var](../programming/api/api-nki-language-misc.md#nki-language-var) | -| `vector_select` | nki.tensor | Per-partition indirect addressing using a vector of offsets. | [NkiTensor.vector_select](../programming/api/api-nki-tensor.md#nki-tensor-vector_select) | -| `view` | nki.tensor | Reinterpret the tensor's data as a different dtype. | [NkiTensor.view](../programming/api/api-nki-tensor.md#nki-tensor-view) | +| Symbol | Module | Description | Documentation | +| --------------- | ------------ | ------------------------------------------------------------ | ---------------------------------------------------------------------------------------- | +| `var` | nki.language | Variance along the specified axis (or axes) of the input. | [nki.language.var](../programming/api/api-nki-language-misc.md#nki-language-var) | +| `vector_select` | nki.tensor | Per-partition indirect addressing using a vector of offsets. | [NkiTensor.vector_select](../programming/api/api-nki-tensor.md#nki-tensor-vector_select) | +| `view` | nki.tensor | Reinterpret the tensor's data as a different dtype. | [NkiTensor.view](../programming/api/api-nki-tensor.md#nki-tensor-view) | --- ## W -| Symbol | Module | Description | Documentation | -|--------|--------|-------------|---------------| -| `where` | nki.language | Return elements chosen from x or y depending on condition. | [nki.language.where](../programming/api/api-nki-language-misc.md#nki-language-where) | -| `while_loop` | nki.language | Structured while loop with a register condition. | [nki.language.while_loop](../programming/api/api-nki-language-dims.md#nki-language-while_loop) | +| Symbol | Module | Description | Documentation | +| ------------ | ------------ | ---------------------------------------------------------- | ---------------------------------------------------------------------------------------------- | +| `where` | nki.language | Return elements chosen from x or y depending on condition. | [nki.language.where](../programming/api/api-nki-language-misc.md#nki-language-where) | +| `while_loop` | nki.language | Structured while loop with a register condition. | [nki.language.while_loop](../programming/api/api-nki-language-dims.md#nki-language-while_loop) | --- ## Z -| Symbol | Module | Description | Documentation | -|--------|--------|-------------|---------------| -| `zeros` | nki.language | Create zero-filled tensor | [nki.language.zeros](../programming/api/nki.language.md) | +| Symbol | Module | Description | Documentation | +| ------------ | ------------ | ---------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------- | +| `zeros` | nki.language | Create zero-filled tensor | [nki.language.zeros](../programming/api/nki.language.md) | | `zeros_like` | nki.language | Create a new tensor of zeros with the same shape and type as a given tensor. | [nki.language.zeros_like](../programming/api/api-nki-language-misc.md#nki-language-zeros_like) | --- @@ -379,145 +379,161 @@ Quick reference for finding NKI API function and symbol documentation. Symbols a ## Symbols by Category ### Tensor Creation -| Symbol | Documentation | -|--------|---------------| + +| Symbol | Documentation | +| ---------------------- | ------------------------------------------ | | `nki.language.ndarray` | [Link](../programming/api/nki.language.md) | -| `nki.language.zeros` | [Link](../programming/api/nki.language.md) | +| `nki.language.zeros` | [Link](../programming/api/nki.language.md) | ### Memory Buffers -| Symbol | Documentation | -|--------|---------------| -| `nki.language.sbuf` | [Link](../programming/api/api-nki-language-memory.md#nki-language-sbuf) | -| `nki.language.psum` | [Link](../programming/api/api-nki-language-memory.md#nki-language-psum) | -| `nki.language.hbm` | [Link](../programming/api/api-nki-language-memory.md#nki-language-hbm) | + +| Symbol | Documentation | +| -------------------------- | ------------------------------------------------------------------------------ | +| `nki.language.sbuf` | [Link](../programming/api/api-nki-language-memory.md#nki-language-sbuf) | +| `nki.language.psum` | [Link](../programming/api/api-nki-language-memory.md#nki-language-psum) | +| `nki.language.hbm` | [Link](../programming/api/api-nki-language-memory.md#nki-language-hbm) | | `nki.language.private_hbm` | [Link](../programming/api/api-nki-language-memory.md#nki-language-private_hbm) | -| `nki.language.shared_hbm` | [Link](../programming/api/api-nki-language-memory.md#nki-language-shared_hbm) | +| `nki.language.shared_hbm` | [Link](../programming/api/api-nki-language-memory.md#nki-language-shared_hbm) | ### Loop Iterators -| Symbol | Documentation | -|--------|---------------| -| `range` (recommended) | Standard Python range | -| `nki.language.static_range` | [Link](../programming/api/nki.language.md) (legacy alias for `range`) | -| `nki.language.affine_range` | [Link](../programming/api/nki.language.md) (legacy alias for `range`) | + +| Symbol | Documentation | +| ------------------------------- | --------------------------------------------------------------------- | +| `range` (recommended) | Standard Python range | +| `nki.language.static_range` | [Link](../programming/api/nki.language.md) (legacy alias for `range`) | +| `nki.language.affine_range` | [Link](../programming/api/nki.language.md) (legacy alias for `range`) | | `nki.language.sequential_range` | [Link](../programming/api/nki.language.md) (legacy alias for `range`) | ### Data Types -| Symbol | Documentation | -|--------|---------------| -| `nki.language.bool_` | [Link](../programming/api/api-nki-language-types.md) | -| `nki.language.int8` | [Link](../programming/api/api-nki-language-types.md) | -| `nki.language.int16` | [Link](../programming/api/api-nki-language-types.md) | -| `nki.language.int32` | [Link](../programming/api/api-nki-language-types.md) | -| `nki.language.uint8` | [Link](../programming/api/api-nki-language-types.md) | -| `nki.language.uint16` | [Link](../programming/api/api-nki-language-types.md) | -| `nki.language.uint32` | [Link](../programming/api/api-nki-language-types.md) | -| `nki.language.float16` | [Link](../programming/api/api-nki-language-types.md) | -| `nki.language.float32` | [Link](../programming/api/api-nki-language-types.md) | -| `nki.language.bfloat16` | [Link](../programming/api/api-nki-language-types.md) | -| `nki.language.tfloat32` | [Link](../programming/api/api-nki-language-types.md) | + +| Symbol | Documentation | +| -------------------------- | ---------------------------------------------------- | +| `nki.language.bool_` | [Link](../programming/api/api-nki-language-types.md) | +| `nki.language.int8` | [Link](../programming/api/api-nki-language-types.md) | +| `nki.language.int16` | [Link](../programming/api/api-nki-language-types.md) | +| `nki.language.int32` | [Link](../programming/api/api-nki-language-types.md) | +| `nki.language.uint8` | [Link](../programming/api/api-nki-language-types.md) | +| `nki.language.uint16` | [Link](../programming/api/api-nki-language-types.md) | +| `nki.language.uint32` | [Link](../programming/api/api-nki-language-types.md) | +| `nki.language.float16` | [Link](../programming/api/api-nki-language-types.md) | +| `nki.language.float32` | [Link](../programming/api/api-nki-language-types.md) | +| `nki.language.bfloat16` | [Link](../programming/api/api-nki-language-types.md) | +| `nki.language.tfloat32` | [Link](../programming/api/api-nki-language-types.md) | | `nki.language.float8_e4m3` | [Link](../programming/api/api-nki-language-types.md) | | `nki.language.float8_e5m2` | [Link](../programming/api/api-nki-language-types.md) | ### Matrix Operations (Tensor Engine) -| Symbol | Documentation | -|--------|---------------| -| `nki.isa.nc_matmul` | [Link](../programming/api/api-nki-isa-tensor.md#nki-isa-nc_matmul) | + +| Symbol | Documentation | +| ---------------------- | --------------------------------------------------------------------- | +| `nki.isa.nc_matmul` | [Link](../programming/api/api-nki-isa-tensor.md#nki-isa-nc_matmul) | | `nki.isa.nc_matmul_mx` | [Link](../programming/api/api-nki-isa-tensor.md#nki-isa-nc_matmul_mx) | | `nki.isa.nc_transpose` | [Link](../programming/api/api-nki-isa-tensor.md#nki-isa-nc_transpose) | ### Vector Operations (Vector Engine) -| Symbol | Documentation | -|--------|---------------| -| `nki.isa.tensor_tensor` | [Link](../programming/api/api-nki-isa-tensor.md#nki-isa-tensor_tensor) | + +| Symbol | Documentation | +| ---------------------------- | --------------------------------------------------------------------------- | +| `nki.isa.tensor_tensor` | [Link](../programming/api/api-nki-isa-tensor.md#nki-isa-tensor_tensor) | | `nki.isa.tensor_tensor_scan` | [Link](../programming/api/api-nki-isa-tensor.md#nki-isa-tensor_tensor_scan) | -| `nki.isa.tensor_scalar` | [Link](../programming/api/api-nki-isa-tensor.md#nki-isa-tensor_scalar) | -| `nki.isa.tensor_reduce` | [Link](../programming/api/api-nki-isa-tensor.md#nki-isa-tensor_reduce) | -| `nki.isa.bn_stats` | [Link](../programming/api/api-nki-isa-vector.md#nki-isa-bn_stats) | -| `nki.isa.bn_aggr` | [Link](../programming/api/api-nki-isa-vector.md#nki-isa-bn_aggr) | -| `nki.isa.reciprocal` | [Link](../programming/api/api-nki-isa-scalar.md#nki-isa-reciprocal) | +| `nki.isa.tensor_scalar` | [Link](../programming/api/api-nki-isa-tensor.md#nki-isa-tensor_scalar) | +| `nki.isa.tensor_reduce` | [Link](../programming/api/api-nki-isa-tensor.md#nki-isa-tensor_reduce) | +| `nki.isa.bn_stats` | [Link](../programming/api/api-nki-isa-vector.md#nki-isa-bn_stats) | +| `nki.isa.bn_aggr` | [Link](../programming/api/api-nki-isa-vector.md#nki-isa-bn_aggr) | +| `nki.isa.reciprocal` | [Link](../programming/api/api-nki-isa-scalar.md#nki-isa-reciprocal) | ### Scalar Operations (Scalar Engine) -| Symbol | Documentation | -|--------|---------------| -| `nki.isa.activation` | [Link](../programming/api/api-nki-isa-scalar.md#nki-isa-activation) | + +| Symbol | Documentation | +| --------------------------- | -------------------------------------------------------------------------- | +| `nki.isa.activation` | [Link](../programming/api/api-nki-isa-scalar.md#nki-isa-activation) | | `nki.isa.activation_reduce` | [Link](../programming/api/api-nki-isa-scalar.md#nki-isa-activation_reduce) | -| `nki.isa.dropout` | [Link](../programming/api/api-nki-isa-scalar.md#nki-isa-dropout) | +| `nki.isa.dropout` | [Link](../programming/api/api-nki-isa-scalar.md#nki-isa-dropout) | ### DMA Operations -| Symbol | Documentation | -|--------|---------------| -| `nki.isa.dma_copy` | [Link](../programming/api/api-nki-isa-memory.md#nki-isa-dma_copy) | + +| Symbol | Documentation | +| ----------------------- | ---------------------------------------------------------------------- | +| `nki.isa.dma_copy` | [Link](../programming/api/api-nki-isa-memory.md#nki-isa-dma_copy) | | `nki.isa.dma_transpose` | [Link](../programming/api/api-nki-isa-memory.md#nki-isa-dma_transpose) | -| `nki.isa.dma_compute` | [Link](../programming/api/api-nki-isa-memory.md#nki-isa-dma_compute) | +| `nki.isa.dma_compute` | [Link](../programming/api/api-nki-isa-memory.md#nki-isa-dma_compute) | ### Copy Operations -| Symbol | Documentation | -|--------|---------------| -| `nki.isa.tensor_copy` | [Link](../programming/api/api-nki-isa-tensor.md#nki-isa-tensor_copy) | + +| Symbol | Documentation | +| -------------------------------- | ------------------------------------------------------------------------------- | +| `nki.isa.tensor_copy` | [Link](../programming/api/api-nki-isa-tensor.md#nki-isa-tensor_copy) | | `nki.isa.tensor_copy_predicated` | [Link](../programming/api/api-nki-isa-tensor.md#nki-isa-tensor_copy_predicated) | ### Utility Functions -| Symbol | Documentation | -|--------|---------------| -| `nki.isa.iota` | [Link](../programming/api/api-nki-isa-utility.md#nki-isa-iota) | -| `nki.isa.memset` | [Link](../programming/api/api-nki-isa-memory.md#nki-isa-memset) | -| `nki.isa.affine_select` | [Link](../programming/api/api-nki-isa-utility.md#nki-isa-affine_select) | -| `nki.isa.range_select` | [Link](../programming/api/api-nki-isa-utility.md#nki-isa-range_select) | -| `nki.isa.select_reduce` | [Link](../programming/api/api-nki-isa-utility.md#nki-isa-select_reduce) | -| `nki.isa.max8` | [Link](../programming/api/api-nki-isa-utility.md#nki-isa-max8) | + +| Symbol | Documentation | +| ------------------------- | ------------------------------------------------------------------------- | +| `nki.isa.iota` | [Link](../programming/api/api-nki-isa-utility.md#nki-isa-iota) | +| `nki.isa.memset` | [Link](../programming/api/api-nki-isa-memory.md#nki-isa-memset) | +| `nki.isa.affine_select` | [Link](../programming/api/api-nki-isa-utility.md#nki-isa-affine_select) | +| `nki.isa.range_select` | [Link](../programming/api/api-nki-isa-utility.md#nki-isa-range_select) | +| `nki.isa.select_reduce` | [Link](../programming/api/api-nki-isa-utility.md#nki-isa-select_reduce) | +| `nki.isa.max8` | [Link](../programming/api/api-nki-isa-utility.md#nki-isa-max8) | | `nki.isa.sequence_bounds` | [Link](../programming/api/api-nki-isa-utility.md#nki-isa-sequence_bounds) | ### Gather/Scatter Operations -| Symbol | Documentation | -|--------|---------------| -| `nki.isa.local_gather` | [Link](../programming/api/api-nki-isa-memory.md#nki-isa-local_gather) | -| `nki.isa.nc_n_gather` | [Link](../programming/api/api-nki-isa-tensor.md#nki-isa-nc_n_gather) | -| `nki.isa.nc_find_index8` | [Link](../programming/api/api-nki-isa-tensor.md#nki-isa-nc_find_index8) | + +| Symbol | Documentation | +| --------------------------- | -------------------------------------------------------------------------- | +| `nki.isa.local_gather` | [Link](../programming/api/api-nki-isa-memory.md#nki-isa-local_gather) | +| `nki.isa.nc_n_gather` | [Link](../programming/api/api-nki-isa-tensor.md#nki-isa-nc_n_gather) | +| `nki.isa.nc_find_index8` | [Link](../programming/api/api-nki-isa-tensor.md#nki-isa-nc_find_index8) | | `nki.isa.nc_match_replace8` | [Link](../programming/api/api-nki-isa-tensor.md#nki-isa-nc_match_replace8) | ### Quantization -| Symbol | Documentation | -|--------|---------------| + +| Symbol | Documentation | +| --------------------- | -------------------------------------------------------------------- | | `nki.isa.quantize_mx` | [Link](../programming/api/api-nki-isa-tensor.md#nki-isa-quantize_mx) | ### Random Number Generation -| Symbol | Documentation | -|--------|---------------| -| `nki.isa.rng` | [Link](../programming/api/nki.isa.md) | -| `nki.isa.rand2` | [Link](../programming/api/nki.isa.md) | + +| Symbol | Documentation | +| ------------------------ | ------------------------------------- | +| `nki.isa.rng` | [Link](../programming/api/nki.isa.md) | +| `nki.isa.rand2` | [Link](../programming/api/nki.isa.md) | | `nki.isa.rand_set_state` | [Link](../programming/api/nki.isa.md) | | `nki.isa.rand_get_state` | [Link](../programming/api/nki.isa.md) | -| `nki.isa.set_rng_seed` | [Link](../programming/api/nki.isa.md) | +| `nki.isa.set_rng_seed` | [Link](../programming/api/nki.isa.md) | ### Multi-Core/SPMD -| Symbol | Documentation | -|--------|---------------| -| `nki.language.program_id` | [Link](../programming/api/nki.language.md) | -| `nki.language.num_programs` | [Link](../programming/api/nki.language.md) | -| `nki.language.program_ndim` | [Link](../programming/api/nki.language.md) | -| `nki.isa.core_barrier` | [Link](../programming/api/api-nki-isa-tensor.md#nki-isa-core_barrier) | -| `nki.isa.sendrecv` | [Link](../programming/api/nki.isa.md) | + +| Symbol | Documentation | +| --------------------------- | -------------------------------------------------------------------------- | +| `nki.language.program_id` | [Link](../programming/api/nki.language.md) | +| `nki.language.num_programs` | [Link](../programming/api/nki.language.md) | +| `nki.language.program_ndim` | [Link](../programming/api/nki.language.md) | +| `nki.isa.core_barrier` | [Link](../programming/api/api-nki-isa-tensor.md#nki-isa-core_barrier) | +| `nki.isa.sendrecv` | [Link](../programming/api/nki.isa.md) | | `nki.isa.nc_stream_shuffle` | [Link](../programming/api/api-nki-isa-tensor.md#nki-isa-nc_stream_shuffle) | ### Register Operations -| Symbol | Documentation | -|--------|---------------| + +| Symbol | Documentation | +| ------------------------ | ------------------------------------- | | `nki.isa.register_alloc` | [Link](../programming/api/nki.isa.md) | -| `nki.isa.register_load` | [Link](../programming/api/nki.isa.md) | -| `nki.isa.register_move` | [Link](../programming/api/nki.isa.md) | +| `nki.isa.register_load` | [Link](../programming/api/nki.isa.md) | +| `nki.isa.register_move` | [Link](../programming/api/nki.isa.md) | | `nki.isa.register_store` | [Link](../programming/api/nki.isa.md) | ### Enums and Constants -| Symbol | Documentation | -|--------|---------------| -| `nki.isa.engine` | [Link](../programming/api/nki.isa.md) | -| `nki.isa.reduce_cmd` | [Link](../programming/api/nki.isa.md) | -| `nki.isa.dge_mode` | [Link](../programming/api/nki.isa.md) | -| `nki.isa.dma_engine` | [Link](../programming/api/nki.isa.md) | -| `nki.isa.oob_mode` | [Link](../programming/api/nki.isa.md) | -| `nki.isa.nc_version` | [Link](../programming/api/api-nki-isa-tensor.md#nki-isa-nc_version) | + +| Symbol | Documentation | +| ------------------------ | ----------------------------------------------------------------------- | +| `nki.isa.engine` | [Link](../programming/api/nki.isa.md) | +| `nki.isa.reduce_cmd` | [Link](../programming/api/nki.isa.md) | +| `nki.isa.dge_mode` | [Link](../programming/api/nki.isa.md) | +| `nki.isa.dma_engine` | [Link](../programming/api/nki.isa.md) | +| `nki.isa.oob_mode` | [Link](../programming/api/nki.isa.md) | +| `nki.isa.nc_version` | [Link](../programming/api/api-nki-isa-tensor.md#nki-isa-nc_version) | | `nki.isa.get_nc_version` | [Link](../programming/api/api-nki-isa-tensor.md#nki-isa-get_nc_version) | -| `nki.language.tile_size` | [Link](../programming/api/nki.language.md) | +| `nki.language.tile_size` | [Link](../programming/api/nki.language.md) | --- diff --git a/skills/neuron-nki-docs/references/indices/task-routing.md b/skills/neuron-nki-docs/references/indices/task-routing.md index 2adb2a1..2e13e23 100644 --- a/skills/neuron-nki-docs/references/indices/task-routing.md +++ b/skills/neuron-nki-docs/references/indices/task-routing.md @@ -8,170 +8,170 @@ This index maps common user goals and tasks to the relevant NKI documentation. U ### Get Started with NKI -| Task | Recommended Documentation | -|------|--------------------------| -| **Install NKI and set up my environment** | [Set Up Your Environment](../programming/setup-env.md) | -| **Write my first NKI kernel** | [Quickstart: Implement and Run Your First Kernel](../programming/quickstart-implement-run-kernel.md) | -| **Understand what NKI is and when to use it** | [Introduction to NKI](../programming/nki-introduction.md), [FAQ](../reference/nki_faq.md) | -| **Learn NKI syntax and programming model** | [NKI Language Guide](../programming/nki-language-guide.md) | -| **Understand NKI compilation process** | [NKI Compiler Documentation](../programming/nki-compiler.md) | +| Task | Recommended Documentation | +| --------------------------------------------- | ---------------------------------------------------------------------------------------------------- | +| **Install NKI and set up my environment** | [Set Up Your Environment](../programming/setup-env.md) | +| **Write my first NKI kernel** | [Quickstart: Implement and Run Your First Kernel](../programming/quickstart-implement-run-kernel.md) | +| **Understand what NKI is and when to use it** | [Introduction to NKI](../programming/nki-introduction.md), [FAQ](../reference/nki_faq.md) | +| **Learn NKI syntax and programming model** | [NKI Language Guide](../programming/nki-language-guide.md) | +| **Understand NKI compilation process** | [NKI Compiler Documentation](../programming/nki-compiler.md) | --- ### Understand Hardware Architecture -| Task | Recommended Documentation | -|------|--------------------------| -| **Learn about Trainium/Inferentia2 architecture** | [Trainium/Inferentia2 Architecture](../architecture/trainium_inferentia2_arch.md) | -| **Learn about Trainium2 architecture** | [Trainium2 Architecture](../architecture/trainium2_arch.md) | -| **Learn about Trainium3 architecture** | [Trainium3 Architecture](../architecture/trainium3_arch.md) | -| **Understand NeuronCore compute engines** | [Trainium/Inferentia2 Architecture - Compute Engines](../architecture/trainium_inferentia2_arch.md#neuroncore-v2-compute-engines) | -| **Understand memory hierarchy (HBM, SBUF, PSUM)** | [Memory Hierarchy Overview](../programming/memory-hierarchy-overview.md), [Architecture Guide](../architecture/trainium_inferentia2_arch.md#data-movement) | -| **Learn about Tensor Engine capabilities** | [Architecture Guide - Tensor Engine](../architecture/trainium_inferentia2_arch.md#tensor-engine) | -| **Learn about Vector/Scalar/GpSimd Engines** | [Architecture Guide - Vector Engine](../architecture/trainium_inferentia2_arch.md#vector-engine), [Scalar Engine](../architecture/trainium_inferentia2_arch.md#scalar-engine), [GpSimd Engine](../architecture/trainium_inferentia2_arch.md#gpsimd-engine) | +| Task | Recommended Documentation | +| ------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| **Learn about Trainium/Inferentia2 architecture** | [Trainium/Inferentia2 Architecture](../architecture/trainium_inferentia2_arch.md) | +| **Learn about Trainium2 architecture** | [Trainium2 Architecture](../architecture/trainium2_arch.md) | +| **Learn about Trainium3 architecture** | [Trainium3 Architecture](../architecture/trainium3_arch.md) | +| **Understand NeuronCore compute engines** | [Trainium/Inferentia2 Architecture - Compute Engines](../architecture/trainium_inferentia2_arch.md#neuroncore-v2-compute-engines) | +| **Understand memory hierarchy (HBM, SBUF, PSUM)** | [Memory Hierarchy Overview](../programming/memory-hierarchy-overview.md), [Architecture Guide](../architecture/trainium_inferentia2_arch.md#data-movement) | +| **Learn about Tensor Engine capabilities** | [Architecture Guide - Tensor Engine](../architecture/trainium_inferentia2_arch.md#tensor-engine) | +| **Learn about Vector/Scalar/GpSimd Engines** | [Architecture Guide - Vector Engine](../architecture/trainium_inferentia2_arch.md#vector-engine), [Scalar Engine](../architecture/trainium_inferentia2_arch.md#scalar-engine), [GpSimd Engine](../architecture/trainium_inferentia2_arch.md#gpsimd-engine) | --- ### Implement Common Operations -| Task | Recommended Documentation | -|------|--------------------------| -| **Implement matrix multiplication** | [Matrix Multiplication Tutorial](../programming/tutorials/matrix_multiplication.md), [nki.isa.nc_matmul](../programming/api/api-nki-isa-tensor.md#nki-isa-nc_matmul) | -| **Implement 2D transpose** | [2D Transpose Tutorial](../programming/tutorials/transpose2d.md), [nki.isa.nc_transpose](../programming/api/api-nki-isa-tensor.md#nki-isa-nc_transpose) | -| **Implement average pooling** | [Average Pooling 2D Tutorial](../programming/tutorials/average_pool2d.md) | -| **Implement normalization (LayerNorm, RMSNorm)** | [RMSNorm-Quant Kernel](../reference/library/rmsnorm-quant.md), [bn_stats/bn_aggr](../programming/api/api-nki-isa-vector.md) | -| **Implement attention mechanism** | [Attention CTE Kernel](../reference/library/attention-cte.md), [Attention TKG Kernel](../reference/library/attention-tkg.md) | -| **Implement MLP layers** | [MLP Kernel](../reference/library/mlp.md) | -| **Implement state space models (Mamba)** | [Fused Mamba Tutorial](../programming/tutorials/fused_mamba.md) | -| **Implement QKV projection** | [QKV Kernel](../reference/library/qkv.md) | -| **Implement output projection** | [Output Projection CTE](../reference/library/output-projection-cte.md), [Output Projection TKG](../reference/library/output-projection-tkg.md) | +| Task | Recommended Documentation | +| ------------------------------------------------ | -------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| **Implement matrix multiplication** | [Matrix Multiplication Tutorial](../programming/tutorials/matrix_multiplication.md), [nki.isa.nc_matmul](../programming/api/api-nki-isa-tensor.md#nki-isa-nc_matmul) | +| **Implement 2D transpose** | [2D Transpose Tutorial](../programming/tutorials/transpose2d.md), [nki.isa.nc_transpose](../programming/api/api-nki-isa-tensor.md#nki-isa-nc_transpose) | +| **Implement average pooling** | [Average Pooling 2D Tutorial](../programming/tutorials/average_pool2d.md) | +| **Implement normalization (LayerNorm, RMSNorm)** | [RMSNorm-Quant Kernel](../reference/library/rmsnorm-quant.md), [bn_stats/bn_aggr](../programming/api/api-nki-isa-vector.md) | +| **Implement attention mechanism** | [Attention CTE Kernel](../reference/library/attention-cte.md), [Attention TKG Kernel](../reference/library/attention-tkg.md) | +| **Implement MLP layers** | [MLP Kernel](../reference/library/mlp.md) | +| **Implement state space models (Mamba)** | [Fused Mamba Tutorial](../programming/tutorials/fused_mamba.md) | +| **Implement QKV projection** | [QKV Kernel](../reference/library/qkv.md) | +| **Implement output projection** | [Output Projection CTE](../reference/library/output-projection-cte.md), [Output Projection TKG](../reference/library/output-projection-tkg.md) | --- ### Work with Data Movement -| Task | Recommended Documentation | -|------|--------------------------| -| **Load data from HBM to SBUF** | [nki.isa.dma_copy](../programming/api/api-nki-isa-memory.md#nki-isa-dma_copy), [DMA Overview](../programming/nki-dma-overview.md) | -| **Store data from SBUF to HBM** | [nki.isa.dma_copy](../programming/api/api-nki-isa-memory.md#nki-isa-dma_copy) | -| **Transpose data during DMA** | [nki.isa.dma_transpose](../programming/api/api-nki-isa-memory.md#nki-isa-dma_transpose) | -| **Copy data within on-chip memory** | [nki.isa.tensor_copy](../programming/api/api-nki-isa-tensor.md#nki-isa-tensor_copy) | -| **Perform gather operations** | [nki.isa.local_gather](../programming/api/api-nki-isa-memory.md#nki-isa-local_gather), [nki.isa.nc_n_gather](../programming/api/api-nki-isa-tensor.md#nki-isa-nc_n_gather) | -| **Initialize memory with constant value** | [nki.isa.memset](../programming/api/api-nki-isa-memory.md#nki-isa-memset) | +| Task | Recommended Documentation | +| ----------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| **Load data from HBM to SBUF** | [nki.isa.dma_copy](../programming/api/api-nki-isa-memory.md#nki-isa-dma_copy), [DMA Overview](../programming/nki-dma-overview.md) | +| **Store data from SBUF to HBM** | [nki.isa.dma_copy](../programming/api/api-nki-isa-memory.md#nki-isa-dma_copy) | +| **Transpose data during DMA** | [nki.isa.dma_transpose](../programming/api/api-nki-isa-memory.md#nki-isa-dma_transpose) | +| **Copy data within on-chip memory** | [nki.isa.tensor_copy](../programming/api/api-nki-isa-tensor.md#nki-isa-tensor_copy) | +| **Perform gather operations** | [nki.isa.local_gather](../programming/api/api-nki-isa-memory.md#nki-isa-local_gather), [nki.isa.nc_n_gather](../programming/api/api-nki-isa-tensor.md#nki-isa-nc_n_gather) | +| **Initialize memory with constant value** | [nki.isa.memset](../programming/api/api-nki-isa-memory.md#nki-isa-memset) | --- ### Perform Tensor Operations -| Task | Recommended Documentation | -|------|--------------------------| -| **Element-wise operations between tensors** | [nki.isa.tensor_tensor](../programming/api/api-nki-isa-tensor.md#nki-isa-tensor_tensor) | -| **Tensor-scalar operations with broadcasting** | [nki.isa.tensor_scalar](../programming/api/api-nki-isa-tensor.md#nki-isa-tensor_scalar) | -| **Reduce tensor along axes (sum, max, etc.)** | [nki.isa.tensor_reduce](../programming/api/api-nki-isa-tensor.md#nki-isa-tensor_reduce) | -| **Reduce across partitions** | [nki.isa.tensor_partition_reduce](../programming/api/api-nki-isa-tensor.md#nki-isa-tensor_partition_reduce) | -| **Apply activation functions (relu, gelu, exp, etc.)** | [nki.isa.activation](../programming/api/api-nki-isa-scalar.md#nki-isa-activation) | -| **Compute reciprocal (1/x)** | [nki.isa.reciprocal](../programming/api/api-nki-isa-scalar.md#nki-isa-reciprocal) | -| **Perform scan operations** | [nki.isa.tensor_tensor_scan](../programming/api/api-nki-isa-tensor.md#nki-isa-tensor_tensor_scan) | -| **Generate index patterns (iota)** | [nki.isa.iota](../programming/api/api-nki-isa-utility.md#nki-isa-iota) | -| **Apply causal masking** | [nki.isa.affine_select](../programming/api/api-nki-isa-utility.md#nki-isa-affine_select) | -| **Find top-k values** | [nki.isa.max8](../programming/api/api-nki-isa-utility.md#nki-isa-max8) | +| Task | Recommended Documentation | +| ------------------------------------------------------ | ----------------------------------------------------------------------------------------------------------- | +| **Element-wise operations between tensors** | [nki.isa.tensor_tensor](../programming/api/api-nki-isa-tensor.md#nki-isa-tensor_tensor) | +| **Tensor-scalar operations with broadcasting** | [nki.isa.tensor_scalar](../programming/api/api-nki-isa-tensor.md#nki-isa-tensor_scalar) | +| **Reduce tensor along axes (sum, max, etc.)** | [nki.isa.tensor_reduce](../programming/api/api-nki-isa-tensor.md#nki-isa-tensor_reduce) | +| **Reduce across partitions** | [nki.isa.tensor_partition_reduce](../programming/api/api-nki-isa-tensor.md#nki-isa-tensor_partition_reduce) | +| **Apply activation functions (relu, gelu, exp, etc.)** | [nki.isa.activation](../programming/api/api-nki-isa-scalar.md#nki-isa-activation) | +| **Compute reciprocal (1/x)** | [nki.isa.reciprocal](../programming/api/api-nki-isa-scalar.md#nki-isa-reciprocal) | +| **Perform scan operations** | [nki.isa.tensor_tensor_scan](../programming/api/api-nki-isa-tensor.md#nki-isa-tensor_tensor_scan) | +| **Generate index patterns (iota)** | [nki.isa.iota](../programming/api/api-nki-isa-utility.md#nki-isa-iota) | +| **Apply causal masking** | [nki.isa.affine_select](../programming/api/api-nki-isa-utility.md#nki-isa-affine_select) | +| **Find top-k values** | [nki.isa.max8](../programming/api/api-nki-isa-utility.md#nki-isa-max8) | --- ### Work with Quantization -| Task | Recommended Documentation | -|------|--------------------------| -| **Quantize to MXFP8/MXFP4** | [nki.isa.quantize_mx](../programming/api/api-nki-isa-tensor.md#nki-isa-quantize_mx) | -| **Perform MXFP matmul** | [nki.isa.nc_matmul_mx](../programming/api/api-nki-isa-tensor.md#nki-isa-nc_matmul_mx), [MXFP Matmul Guide](../optimization/mxfp-matmul.md) | -| **Use FP8 data types** | [Data Types](../programming/api/api-nki-language-types.md), [Architecture Guide](../architecture/trainium2_arch.md) | +| Task | Recommended Documentation | +| --------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------ | +| **Quantize to MXFP8/MXFP4** | [nki.isa.quantize_mx](../programming/api/api-nki-isa-tensor.md#nki-isa-quantize_mx) | +| **Perform MXFP matmul** | [nki.isa.nc_matmul_mx](../programming/api/api-nki-isa-tensor.md#nki-isa-nc_matmul_mx), [MXFP Matmul Guide](../optimization/mxfp-matmul.md) | +| **Use FP8 data types** | [Data Types](../programming/api/api-nki-language-types.md), [Architecture Guide](../architecture/trainium2_arch.md) | --- ### Implement Distributed/Multi-Core Kernels -| Task | Recommended Documentation | -|------|--------------------------| -| **Write SPMD kernels** | [Logical NeuronCore](../programming/lnc.md), `nl.program_id`, `nl.num_programs` | -| **Use multiple NeuronCores** | [Logical NeuronCore](../programming/lnc.md) | -| **Synchronize across cores** | [nki.isa.core_barrier](../programming/api/api-nki-isa-tensor.md#nki-isa-core_barrier) | -| **Send/receive data between cores** | [nki.isa.sendrecv](../programming/api/api-nki-isa-tensor.md#nki-isa-sendrecv) | -| **Get program/core ID** | [nki.language.program_id](../programming/api/nki.language.md), [nki.language.num_programs](../programming/api/nki.language.md) | +| Task | Recommended Documentation | +| ----------------------------------- | ------------------------------------------------------------------------------------------------------------------------------ | +| **Write SPMD kernels** | [Logical NeuronCore](../programming/lnc.md), `nl.program_id`, `nl.num_programs` | +| **Use multiple NeuronCores** | [Logical NeuronCore](../programming/lnc.md) | +| **Synchronize across cores** | [nki.isa.core_barrier](../programming/api/api-nki-isa-tensor.md#nki-isa-core_barrier) | +| **Send/receive data between cores** | [nki.isa.sendrecv](../programming/api/api-nki-isa-tensor.md#nki-isa-sendrecv) | +| **Get program/core ID** | [nki.language.program_id](../programming/api/nki.language.md), [nki.language.num_programs](../programming/api/nki.language.md) | --- ### Optimize Performance -| Task | Recommended Documentation | -|------|--------------------------| -| **Profile my NKI kernel** | [Profiling with Neuron Profile](../optimization/use-neuron-profile.md) | -| **Optimize overall kernel performance** | [NKI Performance Guide](../optimization/nki_perf_guide.md) | -| **Improve arithmetic intensity** | [Performance Guide - Arithmetic Intensity](../optimization/nki_perf_guide.md#improving-arithmetic-intensity) | -| **Optimize compute efficiency** | [Performance Guide - Compute Efficiency](../optimization/nki_perf_guide.md#optimizing-compute-efficiency) | -| **Optimize data movement** | [Performance Guide - Data Movement](../optimization/nki_perf_guide.md#optimizing-data-movement-efficiency) | -| **Reduce tensor transposes** | [Performance Guide - Opt #8](../optimization/nki_perf_guide.md#opt-8-tensore-only-mitigating-overhead-from-tensor-transposes) | -| **Overlap compute and data loading** | [Performance Guide - Opt #4](../optimization/nki_perf_guide.md#opt-4-overlap-data-loading-with-computation) | -| **Enable engine pipelining** | [Performance Guide - Opt #3](../optimization/nki_perf_guide.md#opt-3-overlap-execution-across-compute-engines-through-pipelining) | -| **Optimize tile sizes** | [Tiling Overview](../programming/tiling-overview.md), [Matrix Multiplication Tutorial](../programming/tutorials/matrix_multiplication.md) | -| **Combine instructions** | [Performance Guide - Opt #6](../optimization/nki_perf_guide.md#opt-6-combine-instructions) | +| Task | Recommended Documentation | +| --------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------- | +| **Profile my NKI kernel** | [Profiling with Neuron Profile](../optimization/use-neuron-profile.md) | +| **Optimize overall kernel performance** | [NKI Performance Guide](../optimization/nki_perf_guide.md) | +| **Improve arithmetic intensity** | [Performance Guide - Arithmetic Intensity](../optimization/nki_perf_guide.md#improving-arithmetic-intensity) | +| **Optimize compute efficiency** | [Performance Guide - Compute Efficiency](../optimization/nki_perf_guide.md#optimizing-compute-efficiency) | +| **Optimize data movement** | [Performance Guide - Data Movement](../optimization/nki_perf_guide.md#optimizing-data-movement-efficiency) | +| **Reduce tensor transposes** | [Performance Guide - Opt #8](../optimization/nki_perf_guide.md#opt-8-tensore-only-mitigating-overhead-from-tensor-transposes) | +| **Overlap compute and data loading** | [Performance Guide - Opt #4](../optimization/nki_perf_guide.md#opt-4-overlap-data-loading-with-computation) | +| **Enable engine pipelining** | [Performance Guide - Opt #3](../optimization/nki_perf_guide.md#opt-3-overlap-execution-across-compute-engines-through-pipelining) | +| **Optimize tile sizes** | [Tiling Overview](../programming/tiling-overview.md), [Matrix Multiplication Tutorial](../programming/tutorials/matrix_multiplication.md) | +| **Combine instructions** | [Performance Guide - Opt #6](../optimization/nki_perf_guide.md#opt-6-combine-instructions) | --- ### Integrate with ML Frameworks -| Task | Recommended Documentation | -|------|--------------------------| -| **Use NKI with PyTorch** | [Framework Custom Operators - PyTorch](../programming/framework_custom_op.md#nki-framework-custom-op-pytorch) | -| **Use NKI with JAX** | [Framework Custom Operators - JAX](../programming/framework_custom_op.md#nki-framework-custom-op-jax) | +| Task | Recommended Documentation | +| ------------------------------------ | -------------------------------------------------------------------------------------------------------------------- | +| **Use NKI with PyTorch** | [Framework Custom Operators - PyTorch](../programming/framework_custom_op.md#nki-framework-custom-op-pytorch) | +| **Use NKI with JAX** | [Framework Custom Operators - JAX](../programming/framework_custom_op.md#nki-framework-custom-op-jax) | | **Use prebuilt NKI Library kernels** | [Using Prebuilt Kernels](../programming/tutorial-use-a-prebuilt-kernel.md), [API Index](../programming/api/index.md) | --- ### Debug and Troubleshoot -| Task | Recommended Documentation | -|------|--------------------------| -| **Understand compiler error messages** | [Compiler Error Codes Index](../debugging/error-codes-index.md) | -| **Fix out-of-memory errors (EOOM)** | [EOOM001](../debugging/error-codes/EOOM001.md), [EOOM002](../debugging/error-codes/EOOM002.md) | -| **Fix unsupported operator errors** | [EVRF001](../debugging/error-codes/EVRF001.md), [EUOC002](../debugging/error-codes/EUOC002.md) | -| **Fix instruction limit errors** | [EBVF030](../debugging/error-codes/EBVF030.md), [EVRF007](../debugging/error-codes/EVRF007.md) | -| **Fix data type errors** | [ESPP004](../debugging/error-codes/ESPP004.md), [EVRF004](../debugging/error-codes/EVRF004.md) | -| **Debug numerical issues** | [FAQ - Debugging](../reference/nki_faq.md#how-can-i-debug-numerical-issues-in-nki-kernels) | -| **Print debug output from kernel** | [nki.language.device_print](../programming/api/nki.language.md) | +| Task | Recommended Documentation | +| -------------------------------------- | ---------------------------------------------------------------------------------------------- | +| **Understand compiler error messages** | [Compiler Error Codes Index](../debugging/error-codes-index.md) | +| **Fix out-of-memory errors (EOOM)** | [EOOM001](../debugging/error-codes/EOOM001.md), [EOOM002](../debugging/error-codes/EOOM002.md) | +| **Fix unsupported operator errors** | [EVRF001](../debugging/error-codes/EVRF001.md), [EUOC002](../debugging/error-codes/EUOC002.md) | +| **Fix instruction limit errors** | [EBVF030](../debugging/error-codes/EBVF030.md), [EVRF007](../debugging/error-codes/EVRF007.md) | +| **Fix data type errors** | [ESPP004](../debugging/error-codes/ESPP004.md), [EVRF004](../debugging/error-codes/EVRF004.md) | +| **Debug numerical issues** | [FAQ - Debugging](../reference/nki_faq.md#how-can-i-debug-numerical-issues-in-nki-kernels) | +| **Print debug output from kernel** | [nki.language.device_print](../programming/api/nki.language.md) | --- ### Work with Random Numbers -| Task | Recommended Documentation | -|------|--------------------------| -| **Generate random numbers** | [nki.isa.rng](../programming/api/api-nki-isa-tensor.md#nki-isa-rng), [nki.isa.rand2](../programming/api/api-nki-isa-tensor.md#nki-isa-rand2) | -| **Implement dropout** | [nki.isa.dropout](../programming/api/api-nki-isa-scalar.md#nki-isa-dropout) | -| **Set/get RNG state** | [nki.isa.rand_set_state](../programming/api/api-nki-isa-tensor.md#nki-isa-rand_set_state), [nki.isa.rand_get_state](../programming/api/api-nki-isa-tensor.md#nki-isa-rand_get_state) | -| **Seed random number generator** | [nki.isa.set_rng_seed](../programming/api/api-nki-isa-tensor.md#nki-isa-set_rng_seed) | +| Task | Recommended Documentation | +| -------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | +| **Generate random numbers** | [nki.isa.rng](../programming/api/api-nki-isa-tensor.md#nki-isa-rng), [nki.isa.rand2](../programming/api/api-nki-isa-tensor.md#nki-isa-rand2) | +| **Implement dropout** | [nki.isa.dropout](../programming/api/api-nki-isa-scalar.md#nki-isa-dropout) | +| **Set/get RNG state** | [nki.isa.rand_set_state](../programming/api/api-nki-isa-tensor.md#nki-isa-rand_set_state), [nki.isa.rand_get_state](../programming/api/api-nki-isa-tensor.md#nki-isa-rand_get_state) | +| **Seed random number generator** | [nki.isa.set_rng_seed](../programming/api/api-nki-isa-tensor.md#nki-isa-set_rng_seed) | --- ### Understand Tiling and Layout -| Task | Recommended Documentation | -|------|--------------------------| -| **Understand tiling concepts** | [Tiling Overview](../programming/tiling-overview.md) | +| Task | Recommended Documentation | +| ------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------ | +| **Understand tiling concepts** | [Tiling Overview](../programming/tiling-overview.md) | | **Understand partition vs free dimensions** | [Tiling Overview](../programming/tiling-overview.md), [Architecture Guide](../architecture/trainium_inferentia2_arch.md) | -| **Handle tensors larger than tile limits** | [Matrix Multiplication Tutorial - Tiling](../programming/tutorials/matrix_multiplication.md) | -| **Understand indexing** | [Indexing Overview](../programming/indexing-overview.md) | +| **Handle tensors larger than tile limits** | [Matrix Multiplication Tutorial - Tiling](../programming/tutorials/matrix_multiplication.md) | +| **Understand indexing** | [Indexing Overview](../programming/indexing-overview.md) | --- ### Migrate or Update NKI Code -| Task | Recommended Documentation | -|------|--------------------------| -| **Update from NKI 0.2.0 to 0.3.0 (GA)** | [NKI 0.3.0 Update Guide](../reference/migration/nki-030-update-guide.md) | -| **Migrate from Beta 1 to Beta 2** | [NKI Migration Guide](../reference/migration/nki-migration-guide.md) | -| **Update block dimension usage** | [Block Dimension Migration Guide](../reference/migration/nki_block_dimension_migration_guide.md) | -| **Check version information** | [NKI Versions](../optimization/nki-beta-versions.md) | -| **Review release notes** | [NKI Release Notes](../reference/nki_rn.md) | +| Task | Recommended Documentation | +| ----------------------------------------------- | ---------------------------------------------------------------------------------------------------------- | +| **Update from NKI 0.2.0 to 0.3.0 (GA)** | [NKI 0.3.0 Update Guide](../reference/migration/nki-030-update-guide.md) | +| **Migrate from Beta 1 to Beta 2** | [NKI Migration Guide](../reference/migration/nki-migration-guide.md) | +| **Update block dimension usage** | [Block Dimension Migration Guide](../reference/migration/nki_block_dimension_migration_guide.md) | +| **Check version information** | [NKI Versions](../optimization/nki-beta-versions.md) | +| **Review release notes** | [NKI Release Notes](../reference/nki_rn.md) | | **Run kernel without hardware (CPU simulator)** | [NKI 0.3.0 Update Guide - CPU Simulator](../reference/migration/nki-030-update-guide.md#nki-cpu-simulator) | --- @@ -179,6 +179,7 @@ This index maps common user goals and tasks to the relevant NKI documentation. U ## Common Workflows ### New Kernel Development Workflow + 1. [Set up environment](../programming/setup-env.md) 2. [Follow quickstart](../programming/quickstart-implement-run-kernel.md) 3. [Study language guide](../programming/nki-language-guide.md) @@ -187,12 +188,14 @@ This index maps common user goals and tasks to the relevant NKI documentation. U 6. [Profile and optimize](../optimization/nki_perf_guide.md) ### Performance Optimization Workflow + 1. [Profile kernel](../optimization/use-neuron-profile.md) 2. [Identify bottlenecks](../optimization/nki_perf_guide.md) 3. [Apply relevant optimizations](../optimization/nki_perf_guide.md) 4. [Re-profile to verify improvements](../optimization/use-neuron-profile.md) ### Debugging Workflow + 1. [Check error code](../debugging/error-codes-index.md) 2. [Review FAQ](../reference/nki_faq.md) 3. [Use device_print for debugging](../programming/api/nki.language.md) diff --git a/skills/neuron-nki-docs/references/optimization/deep-dives-overview.md b/skills/neuron-nki-docs/references/optimization/deep-dives-overview.md index 01b4682..27c473d 100644 --- a/skills/neuron-nki-docs/references/optimization/deep-dives-overview.md +++ b/skills/neuron-nki-docs/references/optimization/deep-dives-overview.md @@ -27,4 +27,4 @@ Migrate NKI kernels to use block dimensions for improved performance and resourc [NKI Beta Versions](nki-beta-versions.md) -[NKI Beta Migration Guide](../reference/migration/nki-migration-guide.md) \ No newline at end of file +[NKI Beta Migration Guide](../reference/migration/nki-migration-guide.md) diff --git a/skills/neuron-nki-docs/references/optimization/mxfp-matmul.md b/skills/neuron-nki-docs/references/optimization/mxfp-matmul.md index 82062fb..2368f61 100644 --- a/skills/neuron-nki-docs/references/optimization/mxfp-matmul.md +++ b/skills/neuron-nki-docs/references/optimization/mxfp-matmul.md @@ -5,13 +5,12 @@ In this guide, you’ll learn how to perform MXFP4/8 matrix multiplication, quan ## Before You start -* Read the MX-related sections of the [Trainium 3 Architecture Guide for NKI](../architecture/trainium3_arch.md#trainium3-arch) and become familiar with basic matrix multiplication concepts on Neuron in the [Matrix Multiplication tutorial](../programming/tutorials/matrix_multiplication.md). +- Read the MX-related sections of the [Trainium 3 Architecture Guide for NKI](../architecture/trainium3_arch.md#trainium3-arch) and become familiar with basic matrix multiplication concepts on Neuron in the [Matrix Multiplication tutorial](../programming/tutorials/matrix_multiplication.md). > **Note** > > Note -> -> +> > The code snippets in this guide are taken from the [tutorial code package](https://github.com/aws-neuron/aws-neuron-sdk/tree/master/nki/deep-dives/src/mxfp-matmul) which demonstrates how to execute all MX kernel examples from Torch. We recommend you browse and run the code as you read the tutorial. ### What is MXFP4/8 Matrix Multiplication? @@ -30,13 +29,12 @@ Compared to BF16/FP32 matrix multiplication, the performance uplift from Matmul- First, let’s examine the tile-size constraints for MX so we can allocate the correct space for tensors. MX data is represented in NKI using quad (x4) packed data types ([float8_e5m2_x4](../programming/api/api-nki-language-types.md#nki-language-float8_e5m2_x4), [float8_e4m3fn_x4](../programming/api/api-nki-language-types.md#nki-language-float8_e4m3fn_x4), and [float4_e2m1fn_x4](../programming/api/api-nki-language-misc.md#nki-language-float4_e2m1fn_x4), herein referred to collectively as `MXFP_x4`). The `float8_*_x4` types are 32-bits wide and physically contain four `float8` elements. The `float4_*_x4` type is 16-bits wide and physically contains four `float4` elements. As expressed in `_x4` elements, the TensorE maximum tile sizes in NKI code continue to be given by the existing hardware constraints, summarized below. - | Matrix Type | Data Type | Implied Physical Size | Max Tile Size in Code | -| --- | --- | --- | --- | -| Stationary | BF16 | [128P, 128F] | [128P, 128F] | -| Stationary | MXFP_x4 | [512P, 128F] | [128P, 128F] | -| Moving | BF16 | [128P, 512F] | [128P, 512F] | -| Moving | MXFP_x4 | [512P, 512F] | [128P, 512F] | +| ----------- | --------- | --------------------- | --------------------- | +| Stationary | BF16 | [128P, 128F] | [128P, 128F] | +| Stationary | MXFP_x4 | [512P, 128F] | [128P, 128F] | +| Moving | BF16 | [128P, 512F] | [128P, 512F] | +| Moving | MXFP_x4 | [512P, 512F] | [128P, 512F] | This means that we will allocate data tensors, of type `MXFP_x4`, in our NKI code with the same shapes as we would for BF16/FP32, but it’s implied they contain 4x more contraction elements as shown in the subsequent diagrams. @@ -44,7 +42,6 @@ Now let’s examine a BF16 tile destined to be quantized into a max-sized moving Since a 4x larger contraction dimension is supported we’ll start with a BF16 tile of size `[512, 512]` as shown below. To help us in the subsequent step we’ll also view it as being sectioned into 4 regions of 128 rows (i.e. reshaped as `[4, 128, 512]`). This view is mathematical (i.e. not residing in any particular memory). - > **Figure: mxfp84 matmul guide 1** > > A diagram showing the structure of the moving matrix in BF16 format for MXFP84 matrix multiplication, divided into four 128-element blocks along the 512-element contraction axis. @@ -60,6 +57,7 @@ Since a 4x larger contraction dimension is supported we’ll start with a BF16 t > This visualization demonstrates how the moving matrix is tiled into manageable blocks for the MXFP84 matmul operation, where data is streamed through the tensor engine in chunks. > > **Key Elements:** +> > - **Moving (BF16)**: Title indicating the moving operand in bfloat16 format > - **512 (width)**: The free dimension size of the matrix > - **512 contraction axis**: The total size of the contraction dimension (left label) @@ -70,6 +68,7 @@ Since a 4x larger contraction dimension is supported we’ll start with a BF16 t As explained in the [Trainium 3 Architecture Guide for NKI](../architecture/trainium3_arch.md) we must take 4 elements originating 128 apart on the contraction axis and pack them together on the SBUF free-dimension as shown below. We’ll call this transformation “interleaving”. ! + > **Figure: mxfp84 matmul guide 2** > > A diagram showing the layout of a Moving (BF16) Unquantized Interleaved Data Tile with 128 partitions and 2048 elements in the free dimension, illustrating how data blocks are interleaved. @@ -83,6 +82,7 @@ As explained in the [Trainium 3 Architecture Guide for NKI](../architecture/trai > The right portion of the tile is shown as white/empty space with a black border, indicating the full extent of the 2048-element free dimension. The interleaved colored portion occupies only a small fraction on the left, visually demonstrating the relationship between the interleaved block data and the total tile size. > > **Key Elements:** +> > - **Title**: "Moving (BF16) Unquantized Interleaved Data Tile" identifying the data format > - **128P**: Labelindicating 128 partitions along the P dimension (left side) > - **1F**: Label marking the start of the free dimension @@ -95,7 +95,6 @@ Notice the SBUF shape has become `[128P, 2048F]`. In a subsequent code example w Next, let’s Quantize-MX this tile, which will preserve the layout but pack groups of 4 free-dimension elements into a single `MXFP_x4` element, as shown below. Note that Quantize-MX does not support an FP4 output but Matmul-MX does support FP4 input. - > **Figure: mxfp84 matmul guide 3** > > A diagram showing the Moving (MXFP_x4) Quantized Data Tile layout with 128 partitions and 512 free dimension elements, demonstrating how data is compactly organized after quantization. @@ -111,6 +110,7 @@ Next, let’s Quantize-MX this tile, which will preserve the layout but pack gro > Below the main tile, a legend shows a small red-bordered empty rectangle followed by the label "[1P,1F] MXFP_x4", indicating that each colored block represents one partition by one free element in the MXFP_x4 format. > > **Key Elements:** +> > - **Title**: "Moving (MXFP_x4) Quantized Data Tile" identifying the quantized format > - **F: 512**: Free dimension size of 512 elements > - **P: 128**: Partition dimension size of 128 @@ -123,19 +123,18 @@ Notice the shape is now `[128P, 512F]` which is the max moving tile size we aime With this understanding we’ll state the space allocation rules for quantized `MXFP_x4` data tiles. - ```text Unquantized Interleaved Data Tile = [P,F] BF16 in SBUF MX Quantized Data Tile = [P, F//4] MXFP_x4 in SBUF ``` - ### Scale Tensor Let’s revisit the BF16 tile with the interleaved SBUF layout but this time with one of the `[8P, 4F]` scaling groups overlaid. ! + > **Figure: mxfp84 matmul guide 4** > > A diagram showing the Moving (BF16) Unquantized Interleaved Data Tile with a highlighted scaling group region, illustrating how 8P by 4F scaling groups are organized within the tile for MXFP quantization. @@ -151,6 +150,7 @@ Let’s revisit the BF16 tile with the interleaved SBUF layout but this time wit > The diagram demonstrates that MXFP quantization organizes data into scaling groups of 8 partitions by 4 free dimension elements, where each scaling group shares a common scale factor for the quantized values. > > **Key Elements:** +> > - **Title**: "Moving (BF16) Unquantized Interleaved Data Tile" identifying the data format > - **2048F**: Free dimension size of the full tile > - **128P**: Partition dimension size (128 partitions) @@ -165,7 +165,6 @@ MX scales are represented using a `UINT8` tile containing one element for each s As explained in the [Trainium 3 Architecture Guide for NKI](../architecture/trainium3_arch.md), we view the partition-dimension of SBUF as being split into 4 quadrants of 32 partitions each. Scales must be placed in the quadrant from which the corresponding scaling group originated, as shown below. - > **Figure: mxfp84 matmul guide 5** > > A diagram showing the MX Scale Tile layout in UINT8 format, illustrating four 4P-height scale data strips separated by 32P gaps within a 128P by 512F tile structure. @@ -177,6 +176,7 @@ As explained in the [Trainium 3 Architecture Guide for NKI](../architecture/trai > Each green strip is labeled "4P" on the right side, indicating that each scale data region occupies 4 partitions in height. The strips span the full 512F width of the tile. > > Between each pair of green strips and below the last strip, double-headed vertical arrows are shown with the label "32P", indicating 32-partition gaps between consecutive scale data regions. This spacing pattern shows: +> > - First green strip (4P) at top > - 32P gap > - Second green strip (4P) @@ -189,6 +189,7 @@ As explained in the [Trainium 3 Architecture Guide for NKI](../architecture/trai > The total adds up to 4 strips of 4P each (16P) plus 4 gaps of 32P each (128P total for gaps), but since the last 32P extends beyond, the structure fits within the 128P allocation. This layout corresponds to how scale factors are organized to match the interleaved data format from previous diagrams. > > **Key Elements:** +> > - **Title**: "MX Scale Tile (UINT8)" identifying the scale factor storage format > - **512F**: Free dimension size of 512 elements > - **128P**: Total partition dimension size of 128 @@ -201,7 +202,6 @@ Notice the allocated shape is `[128P, 512F]` despite the underlying useful shape With this understanding we’ll state the space allocation rules for quantized MX scale tiles. - ```text Unquantized Interleaved Data Tile = [P,F] BF16 in SBUF @@ -214,11 +214,9 @@ If P > 32 (Oversize required) MX Quantized Scale = [P, F//4] UINT8 in SBUF ``` - ## Basic Matmul-MX -This NKI example performs a single Matmul-MX using offline-quantized, max-sized input tiles. For simplicity, it assumes the MX *data* tiles in HBM already satisfy the layout requirements so they may be simply loaded straight into SBUF. The MX *scale* tiles require some shuffling. Note that subsequent examples, instead, show how to establish this layout yourself in SBUF. - +This NKI example performs a single Matmul-MX using offline-quantized, max-sized input tiles. For simplicity, it assumes the MX _data_ tiles in HBM already satisfy the layout requirements so they may be simply loaded straight into SBUF. The MX _scale_ tiles require some shuffling. Note that subsequent examples, instead, show how to establish this layout yourself in SBUF. ```python import os @@ -278,18 +276,16 @@ def kernel_offline_quantized_mx_matmul(stationary_mx_data, stationary_mx_scale, return result_hbm ``` - A few notes about the above example: -* The `MXFP_x4` packed data types are custom to NKI and are not supported in Torch. Therefore, we mimic the packed data using `uint8` in Torch and simply view it as `MXFP_x4` in the kernel, as shown. +- The `MXFP_x4` packed data types are custom to NKI and are not supported in Torch. Therefore, we mimic the packed data using `uint8` in Torch and simply view it as `MXFP_x4` in the kernel, as shown. -* The `load_scales_scattered()` helper function reads contiguously packed offline scales from HBM and spreads them across partition-dim quadrants. +- The `load_scales_scattered()` helper function reads contiguously packed offline scales from HBM and spreads them across partition-dim quadrants. -* The PSUM output tile is allocated with data type BF16 to indicate the desired output data type of the Matmul-MX. Note that Matmul-MX ([nki.isa.nc_matmul](../programming/api/api-nki-isa-tensor.md#nki-isa-nc_matmul_mx)) supports both BF16 and FP32 output dtypes. +- The PSUM output tile is allocated with data type BF16 to indicate the desired output data type of the Matmul-MX. Note that Matmul-MX ([nki.isa.nc_matmul](../programming/api/api-nki-isa-tensor.md#nki-isa-nc_matmul_mx)) supports both BF16 and FP32 output dtypes. Let’s also look at the host code which calls this kernel as all subsequent examples use the same structure. - ```python def run_offline_quantized_matmul_mx_test(quantized_dtype): @@ -326,14 +322,13 @@ def run_offline_quantized_matmul_mx_test(quantized_dtype): compare_and_print_results(output_kernel_np, golden) ``` +- The `generate_stabilized_mx_data()` helper function is used to generate MX data on the host. “Stabilized” means the data is randomly generated but injected with certain properties to allow for lossless quantization/dequantization, including constraining the data to be in the FP4/8 range. It conveniently returns MX data as `ml_dtypes` FP4/FP8, the same data packed into `uint` to mimic the `MXFP_x4` packing (suitable for sending to a NKI kernel), MX scales, and a corresponding unquantized FP32 tensor. The input shape argument specifies the unquantized shape. The unquantized tensor is viewed as being in the required layout for MX operations. Therefore to generate an MX data tile of maximum size we must specify an unquantized free-dimension that is 4x larger. In this example the moving unquantized shape is `[128P, 2048F]` and the function will return a `[128P, 512F]` packed MX data tensor, as desired. -* The `generate_stabilized_mx_data()` helper function is used to generate MX data on the host. “Stabilized” means the data is randomly generated but injected with certain properties to allow for lossless quantization/dequantization, including constraining the data to be in the FP4/8 range. It conveniently returns MX data as `ml_dtypes` FP4/FP8, the same data packed into `uint` to mimic the `MXFP_x4` packing (suitable for sending to a NKI kernel), MX scales, and a corresponding unquantized FP32 tensor. The input shape argument specifies the unquantized shape. The unquantized tensor is viewed as being in the required layout for MX operations. Therefore to generate an MX data tile of maximum size we must specify an unquantized free-dimension that is 4x larger. In this example the moving unquantized shape is `[128P, 2048F]` and the function will return a `[128P, 512F]` packed MX data tensor, as desired. - -* `nc_matmul_mx_golden()` is a utility to mimic the hardware’s Matmul-MX operation and is therefore useful for verifying the hardware output. It assumes the input tensors meet the SBUF layout requirements and the data tensor is packed to mimic `MXFP_x4`. Hence it can directly accept MX data generated by `generate_stabilized_mx_data()`. +- `nc_matmul_mx_golden()` is a utility to mimic the hardware’s Matmul-MX operation and is therefore useful for verifying the hardware output. It assumes the input tensors meet the SBUF layout requirements and the data tensor is packed to mimic `MXFP_x4`. Hence it can directly accept MX data generated by `generate_stabilized_mx_data()`. -* `compare_and_print_results()` uses `numpy.allclose()` to check data correctness and print the tensors to `stdout`. +- `compare_and_print_results()` uses `numpy.allclose()` to check data correctness and print the tensors to `stdout`. -* Although this is a single-tile Matmul-MX, larger MX tensors can be multiplied by using the same tiling techniques shown in the non-MX [Matrix Multiplication tutorial](../programming/tutorials/matrix_multiplication.md). +- Although this is a single-tile Matmul-MX, larger MX tensors can be multiplied by using the same tiling techniques shown in the non-MX [Matrix Multiplication tutorial](../programming/tutorials/matrix_multiplication.md). ## Quantize-MX + Matmul-MX @@ -341,10 +336,9 @@ Next we’ll replace one of the Matmul-MX inputs with a tile that we quantize on The two main changes in this example are: -* The `allocate_mx_tiles()` helper function implements the data and scale tile allocation rules mentioned above. - -* `load_scales_scattered()` is again used for the stationary scales but is unnecessary for the moving scales since Quantize-MX will correctly spread the data across SBUF partition-dim quadrants. +- The `allocate_mx_tiles()` helper function implements the data and scale tile allocation rules mentioned above. +- `load_scales_scattered()` is again used for the stationary scales but is unnecessary for the moving scales since Quantize-MX will correctly spread the data across SBUF partition-dim quadrants. ```python # shape_unquantized represents the 2D unquantized SBUF shape with interleaved @@ -432,7 +426,6 @@ def kernel_on_device_quantize_matmul_mx(stationary_mx_data, stationary_mx_scale, return result_hbm ``` - Please see the code package for the host code that calls this kernel. ## SBUF Layout Using Strided Access @@ -459,7 +452,6 @@ Here we DMA a tensor from HBM to SBUF using a strided access pattern. It’s con This example demonstrates both techniques, selected by the `use_tensor_copy` argument. They are very similar but with slightly different read access patterns. It’s useful to refer to the above layout diagrams as you read this code as the reshapes and access patterns directly correspond. - ```python def copy_data_strided(stationary_hbm, moving_hbm, use_tensor_copy: bool = True): @@ -524,16 +516,15 @@ def copy_data_strided(stationary_hbm, moving_hbm, use_tensor_copy: bool = True): return stationary_sbuf_strided.reshape((P_st, F_st*4)), moving_sbuf_strided.reshape((P_mv, F_mv*4)) ``` - See the code package for an example kernel that calls `copy_data_strided()` to establish the interleaved layout for stationary and moving tiles, quantize both, and perform a Matmul-MX. ## Additional Tips -* It’s important to plan where in your design you’ll pay the cost of interleaving the data. Ideally you minimize the cost by finding existing, prior compute on which you can apply the strided access pattern. Or find existing compute against which you can overlap the interleave process. For offline MX weights prepare the layout offline on CPU so you may load the data to SBUF directly in a contiguous/unstrided fashion. +- It’s important to plan where in your design you’ll pay the cost of interleaving the data. Ideally you minimize the cost by finding existing, prior compute on which you can apply the strided access pattern. Or find existing compute against which you can overlap the interleave process. For offline MX weights prepare the layout offline on CPU so you may load the data to SBUF directly in a contiguous/unstrided fashion. -* As with all compute on Neuron, it’s generally performant to spread it across multiple engines operating in parallel. Given that Quantize-MX runs exclusively on the VectorE a bit more care may be needed to alleviate VectorE contention by becoming familiar with operations that may be relegated other engines, like ScalarE. +- As with all compute on Neuron, it’s generally performant to spread it across multiple engines operating in parallel. Given that Quantize-MX runs exclusively on the VectorE a bit more care may be needed to alleviate VectorE contention by becoming familiar with operations that may be relegated other engines, like ScalarE. -* The TensorE operates at double the clock frequency of VectorE, therefore Matmul-MX produces data at double the rate that Quantize-MX can consume it. It may seem that the TensorE could be back-pressured in a situation where a Matmul-MX quickly feeds a subsequent Matmul-MX (since you must Quantize-MX in between at half the speed), but that only happens for small tensors. Larger tensors require tiled matrix multiplication which inherently reuses input (quantized) tiles, allowing time for prior matmul output data to be quantized. +- The TensorE operates at double the clock frequency of VectorE, therefore Matmul-MX produces data at double the rate that Quantize-MX can consume it. It may seem that the TensorE could be back-pressured in a situation where a Matmul-MX quickly feeds a subsequent Matmul-MX (since you must Quantize-MX in between at half the speed), but that only happens for small tensors. Larger tensors require tiled matrix multiplication which inherently reuses input (quantized) tiles, allowing time for prior matmul output data to be quantized. Matmul-MX supports PE-tiling (row-tiling only) where matmuls with a small (<= 64) contraction-dimension (partition-dimension) may be parallelized on the TensorE. This becomes more relevant for MX since a 4x-larger effective contraction-dimension is supported, meaning it’s useful for an `MXFP_x4` contraction-dimension <= 64 or an equivalent unquantized contraction-dimension <= 256. @@ -541,7 +532,6 @@ Matmul-MX supports PE-tiling (row-tiling only) where matmuls with a small (<= 64 After downloading the [tutorial code package](https://github.com/aws-neuron/aws-neuron-sdk/tree/master/nki/deep-dives/src/mxfp-matmul) to your Trainium3 Neuron environment, simply execute it as follows and observe the sample output. - ```bash $ python3 mx_toplevel.py @@ -657,4 +647,4 @@ Golden: [[ 0.32461044 0.43410686 -0.09810834] ... [ 0.82437325 -2.1703691 0.71522826]] ... [[-0.47003102 -0.733371 0.09745546] ... [ 1.3250915 -1.0969493 -0.32166338]] -``` \ No newline at end of file +``` diff --git a/skills/neuron-nki-docs/references/optimization/nki-beta-versions.md b/skills/neuron-nki-docs/references/optimization/nki-beta-versions.md index b00c44c..d09d082 100644 --- a/skills/neuron-nki-docs/references/optimization/nki-beta-versions.md +++ b/skills/neuron-nki-docs/references/optimization/nki-beta-versions.md @@ -38,4 +38,4 @@ NKI Beta 1 (`neuronxcc.nki.*` namespace) is no longer supported. NKI 0.3.0 does ## NKI Support Information -For support with NKI, file a [GitHub issue](https://github.com/aws-neuron/aws-neuron-sdk/issues) and provide us the details of your experience or issue. Other contact details can be found here: [Contact us](../programming/api/index.md#contact-us). \ No newline at end of file +For support with NKI, file a [GitHub issue](https://github.com/aws-neuron/aws-neuron-sdk/issues) and provide us the details of your experience or issue. Other contact details can be found here: [Contact us](../programming/api/index.md#contact-us). diff --git a/skills/neuron-nki-docs/references/optimization/nki_perf_guide.md b/skills/neuron-nki-docs/references/optimization/nki_perf_guide.md index 93071a5..460ae19 100644 --- a/skills/neuron-nki-docs/references/optimization/nki_perf_guide.md +++ b/skills/neuron-nki-docs/references/optimization/nki_perf_guide.md @@ -31,8 +31,8 @@ engine and data movement efficiency, respectively. ## Improving Arithmetic Intensity Arithmetic intensity of a computation workload is commonly defined as the average number of computation operations performed -per byte of data accessed from memory. In the context of NeuronDevices, the definition refers to data accessed from *device -memory* (HBM), since the on-chip memory (SBUF) has sufficient bandwidth to keep all compute engines busy. +per byte of data accessed from memory. In the context of NeuronDevices, the definition refers to data accessed from _device +memory_ (HBM), since the on-chip memory (SBUF) has sufficient bandwidth to keep all compute engines busy. When arithmetic intensity is overly low, compute engines would be consuming data much faster than DMA engines fetching data from device memory into the on-chip memory SBUF. In this case, the execution is bounded by the available device memory bandwidth. @@ -46,11 +46,11 @@ of an algorithm. Fig. 67 The Roofline Model. -*Algorithmic* arithmetic intensity is an intrinsic characteristic of the particular workload and solely dependent on the -compute algorithm. In reality, due to limited capacity in SBUF, the *achieved* arithmetic intensity of a NKI kernel implementation +_Algorithmic_ arithmetic intensity is an intrinsic characteristic of the particular workload and solely dependent on the +compute algorithm. In reality, due to limited capacity in SBUF, the _achieved_ arithmetic intensity of a NKI kernel implementation of such workload could be lower than the algorithmic arithmetic intensity. This could lead to excessive compute engine idle -time blocked by completion of data movements. The two typical reasons behind this are *input data reloading* and *intermediate -data spillage*. Let’s discuss how to identify their symptoms in `neuron-profile` and how to mitigate these issues to improve +time blocked by completion of data movements. The two typical reasons behind this are _input data reloading_ and _intermediate +data spillage_. Let’s discuss how to identify their symptoms in `neuron-profile` and how to mitigate these issues to improve arithmetic intensity next. ### Opt #1. Exploit temporal locality to minimize input data reloading @@ -108,7 +108,6 @@ a time. As a simple example, assume a chain of operators `op0 → op1` on a larg fit in SBUF all at once. If we were to do the operators one at a time, we will effectively have the following sequence of events: - ```python for tile in kernel_in_hbm: tile_sbuf = load(tile) @@ -122,10 +121,8 @@ for tile in op1_out_device_memory: store(op1_out_sbuf, kernel_out_hbm) ``` - However, if we fuse the operators from above: - ```python for tile in kernel_in_hbm: tile_sbuf = load(tile) @@ -134,7 +131,6 @@ for tile in kernel_in_hbm: store(op1_out_sbuf, kernel_out_hbm) ``` - Inside a NKI kernel, operator fusion is exactly done as the above through explicit loop fusion. One great use of this optimization is the self attention operator commonly found in Transformer models. Self attention performs @@ -150,7 +146,6 @@ Certain code patterns in NKI might lead to unexpected spilling from programmers these in future releases. As an example, buffers sometimes need to be declared within the inner loop to avoid spilling. In other words, instead of: - ```python buf = nl.ndarray((2, 4, nl.par_dim(128), 512), buffer=nl.sbuf) for i0 in range(2): @@ -159,10 +154,8 @@ for i0 in range(2): ... ``` - we need to implement: - ```python for i0 in range(2): for i1 in range(4): @@ -170,7 +163,6 @@ for i0 in range(2): nisa.dma_copy(dst=buf, src=...) ``` - With the above aforementioned optimizations, the kernel execution should achieve an arithmetic intensity that is somewhat close to the algorithmic arithmetic intensity. At this point, you should be able to observe from the execution timeline in `neuron-profile` whether the kernel spends more time in compute or DMA engines. The `engine/dma_active_time_percent` @@ -183,11 +175,11 @@ to understand how to optimize data movement efficiency. Compute efficiency optimizations typically fall into two categories: -* “time” domain engine utilization: reduce engine idle time to keep the compute engine *on critical path* as busy as possible, -such as enabling pipelining among engines. +- “time” domain engine utilization: reduce engine idle time to keep the compute engine _on critical path_ as busy as possible, + such as enabling pipelining among engines. -* “spatial” domain engine utilization: within the engine active periods, increase instruction efficiency to use as many -hardware units within the engine as possible, such as combining multiple instructions into one. +- “spatial” domain engine utilization: within the engine active periods, increase instruction efficiency to use as many + hardware units within the engine as possible, such as combining multiple instructions into one. Let’s dive into each category below. @@ -202,7 +194,7 @@ on the idle gaps on VectorE: Fig. 72 Engine idle gaps. -*Side note*, for faster GUI rendering, neuron-profile enables data sampling by default and “hides” certain instructions +_Side note_, for faster GUI rendering, neuron-profile enables data sampling by default and “hides” certain instructions from the timeline with a large profile. To confirm whether an engine indeed has an idle gap, we recommend zooming into a smaller region of the profile and turn on “Show unsampled data” in `View Edit Settings` to make sure all instructions are rendered: @@ -252,17 +244,17 @@ For example, in Transformer’s self-attention layer, in addition to fusing matm V) in a single kernel to minimize spilling as discussed in [Opt #2](#perf-guide-opt2), we also need to form a complex engine pipeline for the operators to maximize utilization of the compute engines: -* matmul_0/matmul_1: TensorE +- matmul_0/matmul_1: TensorE -* softmax: +- softmax: exponential: ScalarE -* summation: VectorE +- summation: VectorE -* scale by reciprocal of summation: ScalarE +- scale by reciprocal of summation: ScalarE -* for causal self attention, triangular masking: GpSimdE +- for causal self attention, triangular masking: GpSimdE #### Opt #4. Overlap data loading with computation @@ -308,7 +300,7 @@ Fig. 79 DMA and engine timeline with and without overlapping. However, it is also possible that even after maximizing overlapping of compute and data movement the best you can, the data movement duration is still not hidden behind compute even though your kernel has a compute-bound arithmetic intensity. In -these cases, the most common cause is the data movement in your kernel is not using the DMA engines *efficiently*. Refer +these cases, the most common cause is the data movement in your kernel is not using the DMA engines _efficiently_. Refer to a [later section](#perf-guide-memory) to see relevant optimization techniques to improve DMA bandwidth utilization. @@ -403,14 +395,13 @@ trn1/inf2, VectorE cannot run the two independent `nki.isa.tensor_reduce()` inst though the total number of compute lanes required for these instructions does not exceed 128. To improve VectorE utilization in this case, we can: -* The two `nc_matmul()` instructions write to disjoint PSUM partitions: partition 0-63 for the first `nc_matmul` and -partition 64-127 for the second one. +- The two `nc_matmul()` instructions write to disjoint PSUM partitions: partition 0-63 for the first `nc_matmul` and + partition 64-127 for the second one. -* Invoke a single `nki.isa.tensor_reduce()` instruction to process output of both `nki.isa.nc_matmul()` instructions. +- Invoke a single `nki.isa.tensor_reduce()` instruction to process output of both `nki.isa.nc_matmul()` instructions. The below pseudo-code illustrates the above computation without and with partition vectorization. - ```python import nki.isa as nisa import nki.language as nl @@ -446,7 +437,6 @@ mm_tile[i_output1_p, ...] = nki.isa.nc_matmul(...) reduce = nisa.tensor_reduce(mm_tile, ...) ``` - Option #2 above is able to perform the reduction 2x faster, by vectorizing the partition dimension and performing a single reduction instead of two. @@ -471,7 +461,6 @@ For example, below pseudo-code showcase combining three instructions into a sing 2` are functionally equivalent, but `impl 2` is 3x faster in terms of latency by touching the input `data` only once and running all three operations (multiply, add, exp) in a pipeline. - ```python import nki.isa as nisa import nki.language as nl @@ -488,7 +477,6 @@ exp = nisa.activation(nl.exp, data, bias, scale) ``` - Check out [nki.isa APIs](../programming/api/nki.isa.md) to understand low-level ISA API semantics, limitations, engine mapping, and rough estimates of performance cost. @@ -500,9 +488,9 @@ combine matrix-vector multiplication and exponential evaluation in a single `nis **Symptom**: Let’s consider a matrix multiplication between two matrices of shape `[M, K]` and `[K, N]`, with one of the following conditions: -* M is significantly smaller than 128, while N is much larger than 128, or +- M is significantly smaller than 128, while N is much larger than 128, or -* the other way around: N is significantly smaller than 128, while M is much larger than 128 +- the other way around: N is significantly smaller than 128, while M is much larger than 128 In NKI, if the matrix with `min(M, N)` dimension is mapped to the **stationary tensor** (`x` input tensor in `nl.matmul` and `nisa.nc_matmul`) for the TensorE `LoadStationary` instruction (details see [architecture guide](../architecture/trainium_inferentia2_arch.md#arch-guide-tensor-engine) @@ -548,7 +536,6 @@ output tensor will be transposed from the original output. Recall, if there is a difference in initiation interval between `LoadStationary` and `MultiplyMoving`, one of them can end up limiting the throughput of TensorE: - > **Figure: mm bottleneck** > > A timing diagram comparing two execution scenarios for matrix multiplication: MultiplyMoving Bounded (where compute is the bottleneck) and LoadStationary Bounded (where memory loading is the bottleneck). @@ -556,12 +543,14 @@ can end up limiting the throughput of TensorE: > This diagram shows two execution timeline scenarios illustrating different bottleneck conditions in matrix multiplication on NeuronCore, helping developers understand performance limiting factors. > > Part (a) "MultiplyMoving Bounded" (top section) shows two parallel timelines: +> > - **LoadStationary row**: Shows sequential loading operations LS[0], LS[1], LS[2], LS[3], ... with blocks colored in shades of blue/green. These complete relatively quickly with gaps between them. > - **MultiplyMoving row**: Shows sequential computation operations MM[0], MM[1], MM[2], MM[3], ... with blocks colored in shades of blue, green, and purple. These operations are longer and continuous, forming the critical path. > > In this scenario, LoadStationary completes before MultiplyMoving needs the data, indicating compute is the bottleneck. The computation (MultiplyMoving) takes longer than data loading (LoadStationary). > > Part (b) "LoadStationary Bounded" (bottom section) shows two parallel timelines: +> > - **LoadStationary row**: Shows the same LS[0] through LS[3] operations, but now they are longer and form a continuous sequence. > - **MultiplyMoving row**: Shows MM[0] through MM[3] operations with gaps between them, waiting for data to be loaded. > @@ -570,6 +559,7 @@ can end up limiting the throughput of TensorE: > Both timelines have arrows extending to the right with ellipsis (...) indicating the pattern continues. > > **Key Elements:** +> > - **LoadStationary (LS)**: Operations loading the stationary matrix into Tensor Engine > - **MultiplyMoving (MM)**: Matrix multiplication operations with moving matrix > - **LS[0]-LS[3]**: Individual load operations (blue/teal colors) @@ -579,7 +569,6 @@ can end up limiting the throughput of TensorE: > - **Timeline arrows**: Show execution sequence over time > - **Gaps vs continuous**: Visual indication of which operation is bottleneck - Fig. 87 Two possible TensorE performance characteristics. In the above scenarios, we expect TensorE performance to be bound by whichever instruction reads the longer tensor - LoadStationary @@ -588,13 +577,13 @@ in “Short Moving”, and MultiplyMoving in “Short Stationary”. However, wi So in the two above scenarios: -* Short Moving - `LoadStationary` initiation interval is roughly equal to the number of elements divided by 4 (because -of fast LoadStationary), and `MultiplyMoving` initiation interval is dominated TensorE instruction turnaround time `MM_INIT_LATENCY +- Short Moving - `LoadStationary` initiation interval is roughly equal to the number of elements divided by 4 (because + of fast LoadStationary), and `MultiplyMoving` initiation interval is dominated TensorE instruction turnaround time `MM_INIT_LATENCY (64 cycles on trn1)`. Therefore, we have `LS_II ~= 128/4 = 32 cycles` , and `MM_II ~= max(1, MM_INIT_LATENCY=64 cycles)` -which leads to issuing a MM roughly every 64 cycles. + which leads to issuing a MM roughly every 64 cycles. -* Short Stationary - `MultiplyMoving` initiation interval will dominate, which leads to issuing a MM roughly every 128 -cycles. +- Short Stationary - `MultiplyMoving` initiation interval will dominate, which leads to issuing a MM roughly every 128 + cycles. Because of the above, we will prefer to map short tensors to the moving tensor in `MultiplyMoving` instruction in TensorE. @@ -638,11 +627,11 @@ These transposes are most commonly lowered down to Tensor Engine. Broadly speaking, there are 2 different types of tensor transposes, with different root causes: -* IO tensor transpose (abbreviated as IO transpose) +- IO tensor transpose (abbreviated as IO transpose) -* intermediate tensor transpose (abbreviated as intermediate transpose) +- intermediate tensor transpose (abbreviated as intermediate transpose) -**IO transpose.** These transposes are ** done on NKI kernel IO (input/output) tensors, which must reside in device memory +**IO transpose.** These transposes are \*\* done on NKI kernel IO (input/output) tensors, which must reside in device memory in current NKI releases. The transposes are needed when the NKI compute API consuming input tensors or producing the output tensors expect a different layout than their IO layout in device memory. To simplify discussion, we dive into input tensor layout discussion below, but the same reasoning also applies to output tensors. @@ -657,7 +646,7 @@ to transpose the input tensor on the fly in the DMA engine, with a major caveat `nl.load`. `nl.load_transpose2d` could make sense in a compute-bound kernel, but should certainly be avoided in memory-bound kernels. -Either way, an IO transpose is inevitable here *due to* the IO tensor layout choice we made as NKI programmers. In the naive +Either way, an IO transpose is inevitable here _due to_ the IO tensor layout choice we made as NKI programmers. In the naive case scenario where we only care about reaching the best performance for a single kernel, we can carefully decide on the IO tensor layout to make sure it is compatible with the NKI compute API layout requirements. When the input tensor is consumed by multiple compute APIs with conflicting layout requirements, IO-transposes cannot be avoided but should still be minimized @@ -746,7 +735,6 @@ and nl.store. For example, the below data loading will trigger 16 DMA transfers that can be run on all 16 DMA engines, which each transfer loading 8 SBUF partitions’ worth of data with a transfer size of 32KiB: - ```python import nki.language as nl @@ -775,7 +763,6 @@ def load_store_32kib_contiguous(in_tensor, out_tensor): nisa.dma_copy(dst=out_tensor[i_p, i_f], src=data_tile) ``` - ### Opt #10: Minimize use of DMA transposes. **Symptom**: Excessive use of DMA transposes, invoked through `nl.load_transpose2d`, can degrade DMA bandwidth significantly. @@ -792,7 +779,6 @@ inevitable and the kernel is memory bound, we recommend replacing `nl.load_trans For example, if you have an `in_tensor` of shape [8192, 128] in device memory but you would like an SBUF tile of shape [128, 8192] spread across 128 partitions for computation, the following two code snippets can achieve the same functionality: - ```python # Option 1, low DMA bandwidth usage: sbuf_opt1 = nl.load_transpose2d(in_tensor[0:8192, 0:128]) @@ -806,7 +792,6 @@ for i_in_tile in range(8192 // 128): sbuf_opt2[0:128, i_start:i_start+128] = nisa.nc_transpose(current_tile) ``` - Option 2 above is especially great for cases where `nl.load_transpose2d` is slowing down data movement in the critical path and TensorE is otherwise idle. Occasionally Option 1 can still be the right call, when the amount of data to be transposed -is small and the overhead of `nl.load_transpose2d` can be well hidden behind other useful computation. \ No newline at end of file +is small and the overhead of `nl.load_transpose2d` can be well hidden behind other useful computation. diff --git a/skills/neuron-nki-docs/references/optimization/use-neuron-profile.md b/skills/neuron-nki-docs/references/optimization/use-neuron-profile.md index b3cc137..a3b7c05 100644 --- a/skills/neuron-nki-docs/references/optimization/use-neuron-profile.md +++ b/skills/neuron-nki-docs/references/optimization/use-neuron-profile.md @@ -7,17 +7,17 @@ Learn how to profile Neuron Kernel Interface (NKI) kernels using Neuron Explorer Ensure that you have the latest version of the `aws-neuronx-tools` package installed as Neuron Explorer comes with this package. The `aws-neuronx-tools` package is pre-installed on Neuron DLAMIs. -* For detailed installation instructions, see: [How to Get Started with Neuron Explorer](../../tools/neuron-explorer/get-started.md#new-neuron-profiler-setup). +- For detailed installation instructions, see: [How to Get Started with Neuron Explorer](../../tools/neuron-explorer/get-started.md#new-neuron-profiler-setup). ## Profile a NKI Kernel Profiling NKI (Neuron Kernel Interface) kernels helps you understand hardware level performance characteristics of your kernels running on AWS Trainium and Inferentia devices. When you write or optimize custom NKI kernels, profiling allows you to: -* **Identify bottlenecks**: Determine if your kernel is compute-bound, memory-bound, or limited by data movement. +- **Identify bottlenecks**: Determine if your kernel is compute-bound, memory-bound, or limited by data movement. -* **Optimize performance**: Analyze kernel-level execution time, investigate compute engine utilization, look for opportunities to implement operator fusion to fine-tune performance. +- **Optimize performance**: Analyze kernel-level execution time, investigate compute engine utilization, look for opportunities to implement operator fusion to fine-tune performance. -* **Compare implementations**: Benchmark different kernel implementations or configurations to pick the most efficient kernel. +- **Compare implementations**: Benchmark different kernel implementations or configurations to pick the most efficient kernel. You can profile NKI kernels using several approaches. In this guide, you’ll learn two primary methods for profiling NKI kernels. @@ -25,11 +25,11 @@ You can profile NKI kernels using several approaches. In this guide, you’ll le To profile an NKI kernel using neuron-profile capture, follow these three steps: -* Set the environment variable `NEURON_FRAMEWORK_DEBUG=1` to instruct the compiler to save the NEFF (Neuron Executable File Format) file. +- Set the environment variable `NEURON_FRAMEWORK_DEBUG=1` to instruct the compiler to save the NEFF (Neuron Executable File Format) file. -* Execute the NKI kernel to generate the NEFF file. +- Execute the NKI kernel to generate the NEFF file. -* Run `neuron-profile capture` to create an Neuron Trace File Format (NTFF) file for performance analysis. +- Run `neuron-profile capture` to create an Neuron Trace File Format (NTFF) file for performance analysis. Each of these steps is explained in detail below. @@ -37,7 +37,6 @@ Each of these steps is explained in detail below. We will profile a 3-layer MLP model that fuses matrix multiplications with ReLU activation functions and uses a NKI matrix multiplication kernel. The rest of this tutorial will use a performance profile generated from this example. Here is the implementation of `mlp_with_mm_kernel.py`. Save this file before moving on to the next step: - ```python """ Example 3-layer MLP with matrix multiplication kernel to demonstrate Neuron Profile. @@ -223,58 +222,48 @@ if __name__ == "__main__": main() ``` - As you can see, at the very top we have added the following flags: - ```python os.environ["NEURON_FRAMEWORK_DEBUG"] = "1" os.environ["XLA_IR_DEBUG"] = "1" os.environ["XLA_HLO_DEBUG"] = "1" ``` - The `NEURON_FRAMEWORK_DEBUG` environment variable enables Neuron debug output. This will trigger the Neuron compiler to save the Neuron Executable File Format (NEFF) artifact to the current directory after compilation of your NKI kernel. The NEFF contains all hardware instructions required to execute your NKI kernel on a NeuronDevice, as well as metadata and debug info needed for profiling. To enable source code linking to framework code (ex. PyTorch) set the environment variables `XLA_IR_DEBUG=1` and `XLA_HLO_DEBUG=1`. #### Step 2: Compile Your NKI Kernel Compile your NKI kernel to create a NEFF in your current directory: - ```python $ python3 mlp_with_mm_kernel.py ``` - > **Note** > > Note -> -> +> > Find your NEFF file, which will be named something like `MODULE_SyncTensorsGraph.81_690876920003119736.neff`. #### Step 3: Profile the Generated NEFF The last step is profiling the generated NEFF. This step executes the NEFF on the NeuronDevice and records a raw execution trace into a NTFF artifact: - ```python $ neuron-explorer capture -n -s profile.ntff --profile-nth-exec=2 --enable-dge-notifs ``` - This will save your NTFF profile to `profile_exec_2.ntff`. important: - ```python The ``--profile-nth-exec=2`` option will profile your NEFF twice on the NeuronDevice and output a NTFF profile for the second iteration. This is recommended to avoid one-time warmup delays which can be seen in the first iteration of execution. The ``--enable-dge-notifs`` option enables the capture of DGE DMA events but has known issues where it may overflow the status notification queue and cause execution timeouts when there are many DGE instructions. ``` - ## View the Neuron Explorer UI This section assumes you’ve completed the previous step and have already generated both the NEFF and NTFF files, and downloaded them on your local machine. @@ -283,46 +272,37 @@ Neuron Explorer includes an interactive, web-based UI for exploring execution tr To view the Neuron Profile Web UI, execute the view command to start Web UI, replacing `` with a path to a folder to store your profiling artifacts: - ```python $ neuron-explorer view --data-path ./ ``` - `` is a path that neuron profile will use for storing and managing profiles. The above command should print a URL that you can click to open the web UI: - ```python View a list of profiles at http://localhost:3001/ ``` - ### Port Forwarding for Remote Instances If `neuron-profile view` is run on a remote instance, you may need to use port forwarding to access the web UI. By default, neuron-profile creates a web server on port 3001 and the API server on port 3002. To enable connection to your browser in you local computer, we will need to establish an ssh tunnel to both of the ports. For example: - ```python ssh -L 3001:localhost:3001 -L 3002:localhost:3002 @ -fN ``` - If you created an EC2 instance with `pem` credentials, include it in the `ssh` tunnel below: - ```python ssh -i ~/my-ec2.pem -L 3001:localhost:3001 -L 3002:localhost:3002 ubuntu@[PUBLIC_IP_ADDRESS] -fN ``` - ### Using the Profile UI -* Once the ssh tunnel is setup, you can now open a browser and navigate to [http://localhost:3001](http://localhost:3001). - +- Once the ssh tunnel is setup, you can now open a browser and navigate to [http://localhost:3001](http://localhost:3001). > **Figure: nki profiler 1** > @@ -333,28 +313,33 @@ ssh -i ~/my-ec2.pem -L 3001:localhost:3001 -L 3002:localhost:3002 ubuntu@[PUBLIC > The top header bar shows "Neuron Profiler" as the application title on the left, and a user menu showing "myself" with a dropdown on the right. > > The left sidebar navigation panel (dark background) contains: +> > - "Profile" section header with a collapse arrow > - "Profile Manager" (highlighted/selected) > - "Profile" link > - "Summary" link > > The main content area displays the "Profile Manager" page: +> > - Title: "Profile Manager" > - Subtitle: "Times are displayed in America/Toronto time" > - Blue "Upload Profile" button in the top right corner > > Below the title are four navigation tabs: +> > - "User uploaded" (currently selected, underlined) > - "User favorite" > - "Search Profile" > - "View History" > > The profiles section shows: +> > - Header: "Profiles (0)" with subtitle "My Uploaded Profiles" > - Pagination controls showing "1" with navigation arrows > - A settings gear icon > > A data table with column headers: +> > - Status (with filter dropdown) > - Profile Name (with filter dropdown) > - Complete P... (truncated, with filter dropdown) @@ -366,6 +351,7 @@ ssh -i ~/my-ec2.pem -L 3001:localhost:3001 -L 3002:localhost:3002 ubuntu@[PUBLIC > The table body shows an empty state with a search/magnifying glass icon and the message "No profiles found - No profiles available for the selected type." > > **Key Elements:** +> > - **Neuron Profiler**: Application title in header > - **Profile Manager**: Main page for managing uploaded profiles > - **Upload Profile button**: Blue button to upload new profile data @@ -374,9 +360,7 @@ ssh -i ~/my-ec2.pem -L 3001:localhost:3001 -L 3002:localhost:3002 ubuntu@[PUBLIC > - **Column filters**: Sortable/filterable columns for profile management > - **User menu**: "myself" dropdown for user account options - -* Click on the button “Upload Profile” to upload NEFF and NTFF files, and give a meaningful name to your profile. Selecting a source code folder for code linking is optional. - +- Click on the button “Upload Profile” to upload NEFF and NTFF files, and give a meaningful name to your profile. Selecting a source code folder for code linking is optional. > **Figure: nki profiler 2** > @@ -389,9 +373,11 @@ ssh -i ~/my-ec2.pem -L 3001:localhost:3001 -L 3002:localhost:3002 ubuntu@[PUBLIC > The dialog contains several input sections: > > **Profile Name Section:** +> > - Text input field containing "mlp_with_mm_kernel" as the profile name > > **NEFF File Section (Required):** +> > - Header: "NEFF File" with "Required" label > - Upload area with upload icon and text "Drop NEFF file" / "Drag .neff file or browse" > - "Browse Files" button @@ -401,7 +387,8 @@ ssh -i ~/my-ec2.pem -L 3001:localhost:3001 -L 3002:localhost:3002 ubuntu@[PUBLIC > - X button to remove the file > > **NTFF File Section (Required):** -> - Header: "NTFF File" with "Required" label +> +> - Header: "NTFF File" with "Required" label > - Upload area with upload icon and text "Drop NTFF file" / "Drag .ntff file or browse" > - "Browse Files" button > - "Selected File" subsection showing: @@ -410,6 +397,7 @@ ssh -i ~/my-ec2.pem -L 3001:localhost:3001 -L 3002:localhost:3002 ubuntu@[PUBLIC > - X button to remove the file > > **Source Code Section:** +> > - Header: "Source Code" > - Upload area with upload icon and text "Drop source code files or folders" / "Drag files/folders or browse" > - Two buttons: "Browse Files" and "Browse Folders" @@ -419,13 +407,16 @@ ssh -i ~/my-ec2.pem -L 3001:localhost:3001 -L 3002:localhost:3002 ubuntu@[PUBLIC > - X button to remove the file > > **Options Section:** +> > - Checkbox (unchecked): "Force upload (overwrite existing profile with same NEFF and NTFF)" > > **Action Buttons:** +> > - "Cancel" button (gray) > - "Upload" button (blue) > > **Key Elements:** +> > - **Profile Name**: Text field for naming the profile > - **NEFF File upload**: Required compiled model file (graph.neff selected) > - **NTFF File upload**: Required profiling trace file (profile_exec_2.ntff selected) @@ -434,9 +425,7 @@ ssh -i ~/my-ec2.pem -L 3001:localhost:3001 -L 3002:localhost:3002 ubuntu@[PUBLIC > - **File badges**: NEFF, NTFF, and script type indicators > - **Cancel/Upload buttons**: Dialog action buttons - -* After the files are uploaded and processed, you will be able to open the profile from the list. - +- After the files are uploaded and processed, you will be able to open the profile from the list. > **Figure: nki profiler 3** > @@ -447,23 +436,27 @@ ssh -i ~/my-ec2.pem -L 3001:localhost:3001 -L 3002:localhost:3002 ubuntu@[PUBLIC > The header shows "Neuron Profiler" on the left and "myself" user menu on the right with a dropdown arrow. > > The left sidebar navigation shows: +> > - "Profile" section with collapse arrow > - "Profile Manager" (highlighted/selected) > - "Profile" link > - "Summary" link > > The main content area shows: +> > - Title: "Profile Manager" > - Subtitle: "Times are displayed in America/Toronto time" > - Blue "Upload Profile" button in the top right > > Tab navigation shows four tabs: +> > - "User uploaded" (selected, underlined) > - "User favorite" > - "Search Profile" > - "View History" > > The Profiles section displays: +> > - Header: "Profiles (1)" indicating one profile > - Subtitle: "My Uploaded Profiles" > - Search box with placeholder "Filter profiles..." @@ -472,6 +465,7 @@ ssh -i ~/my-ec2.pem -L 3001:localhost:3001 -L 3002:localhost:3002 ubuntu@[PUBLIC > - Settings gear icon > > The data table shows one row with columns: +> > - **Status**: Green checkmark icon with "PROCESSED" label > - **Profile Name**: "mlp_with_mm_kernel" (clickable link) > - **Complete P...**: (truncated column) @@ -481,6 +475,7 @@ ssh -i ~/my-ec2.pem -L 3001:localhost:3001 -L 3002:localhost:3002 ubuntu@[PUBLIC > - **Actions**: Star icon (favorite) and pin icon > > **Key Elements:** +> > - **PROCESSED status**: Green indicator showing successful profile processing > - **mlp_with_mm_kernel**: Profile name for the uploaded kernel > - **Upload timestamp**: 11/12/2025, 15:xx @@ -489,9 +484,7 @@ ssh -i ~/my-ec2.pem -L 3001:localhost:3001 -L 3002:localhost:3002 ubuntu@[PUBLIC > - **Filter search**: Search box to filter profiles by name > - **Single profile**: Profile count shows (1) - -* If you click on the name of your profile in Profile Name column, it will navigate to profile page - +- If you click on the name of your profile in Profile Name column, it will navigate to profile page > **Figure: nki profiler 4** > @@ -500,6 +493,7 @@ ssh -i ~/my-ec2.pem -L 3001:localhost:3001 -L 3002:localhost:3002 ubuntu@[PUBLIC > This screenshot displays the Neuron Explorer view within the Neuron Profiler application, providing comprehensive profiling analysis for the "mlp_with_mm_kernel" profile. The interface is divided into a timeline visualization at the top and a data table at the bottom. > > **Header and Navigation:** +> > - Title: "Neuron Explorer" in the header bar > - Profile name: "mlp_with_mm_kernel" displayed below > - Left sidebar shows: Profile Manager, Profile (selected), Summary @@ -507,9 +501,10 @@ ssh -i ~/my-ec2.pem -L 3001:localhost:3001 -L 3002:localhost:3002 ubuntu@[PUBLIC > - User menu shows "myself" > > **Search/Filter Controls:** +> > - Search field with category selector > - "Select category" dropdown -> - "Select field" dropdown +> - "Select field" dropdown > - Text input field ("Enter value") > - "Submit" and "Clear result" buttons > @@ -534,7 +529,8 @@ ssh -i ~/my-ec2.pem -L 3001:localhost:3001 -L 3002:localhost:3002 ubuntu@[PUBLIC > > **Operator Table:** > Below the timeline, a data table shows operator-level performance metrics with columns: -> - Node Name: xla__op+locals+CallImpl_custom-call.5, .4, .3 +> +> - Node Name: xla\_\_op+locals+CallImpl_custom-call.5, .4, .3 > - subgraph_id | subop_id: 0|45, 0|44, 0|43 > - MFU: 26.22%, 34.74%, 37.80% > - HFU: 26.22%, 35.82%, 38.99% @@ -544,10 +540,12 @@ ssh -i ~/my-ec2.pem -L 3001:localhost:3001 -L 3002:localhost:3002 ubuntu@[PUBLIC > - Instructions Ve... (truncated) > > Additional rows show: +> > - aten_relu_maximum.51: 0|28, 0.00% MFU/HFU, 4096 Channels Vector > - aten_relu_maximum.32: 0|22, 0.00% MFU/HFU, 8192 Channels Vector > > **Tab Selectors:** +> > - Operator Table (selected) > - Overall Summary > - Event Details @@ -555,6 +553,7 @@ ssh -i ~/my-ec2.pem -L 3001:localhost:3001 -L 3002:localhost:3002 ubuntu@[PUBLIC > - Annotations > > **Key Elements:** +> > - **Timeline tracks**: Multiple engine/component activity visualization > - **Tensor(nc0)**: Dense orange activity showing tensor engine utilization > - **MFU/HFU columns**: Model/Hardware FLOPS Utilization percentages @@ -562,9 +561,7 @@ ssh -i ~/my-ec2.pem -L 3001:localhost:3001 -L 3002:localhost:3002 ubuntu@[PUBLIC > - **Time scale**: 0 to ~2.07 ms execution window > - **Two NeuronCores**: nc0 and nc1 tracks shown separately - -* If you hover over any engine instruction in the timeline with your mouse, you will see instruction details in a pop-up box. - +- If you hover over any engine instruction in the timeline with your mouse, you will see instruction details in a pop-up box. > **Figure: nki profiler 5** > @@ -573,6 +570,7 @@ ssh -i ~/my-ec2.pem -L 3001:localhost:3001 -L 3002:localhost:3002 ubuntu@[PUBLIC > This screenshot displays the Device Timeline view within the Neuron Profiler, showing execution traces across multiple NeuronCore components with a detailed popup for a selected MATMUL instruction. > > **Timeline Tracks (top to bottom):** +> > - Tensor(nc1): Sparse activity markers > - Tensor(nc0): Dense orange/colored activity bars showing tensor engine operations > - TensorMatrix (nc1): Activity pattern @@ -589,11 +587,12 @@ ssh -i ~/my-ec2.pem -L 3001:localhost:3001 -L 3002:localhost:3002 ubuntu@[PUBLIC > > **Detailed Instruction Popup (purple/lavender background):** > The popup shows information for a selected Tensor operation: -> - Name: S[S] (Tensor)+x@complete acc_flags=0 fp32_mode=LOW_HIGH src=p[32*x34f7*[1,0,0][S12,1,1]] dst=0x20018001[1,0,0][S12,1,1] 128*128 +> +> - Name: S[S] (Tensor)+x@complete acc_flags=0 fp32_mode=LOW_HIGH src=p[32*x34f7*[1,0,0][S12,1,1]] dst=0x20018001[1,0,0][S12,1,1] 128\*128 > - Time: 595,748 ns - 596,339 ns > - Duration: 591 ns > - Opcode: MATMUL -> - Hierarchy: xla__op_CallImpl_custom-call.3 +> - Hierarchy: xla\_\_op_CallImpl_custom-call.3 > - Instruction Type: REGULAR > - Compiler PC: 2704 > - NKI Source Location: /home/ethschan/mlp_with_mm_kernel.py:116 @@ -606,6 +605,7 @@ ssh -i ~/my-ec2.pem -L 3001:localhost:3001 -L 3002:localhost:3002 ubuntu@[PUBLIC > The State Buffer Usage tracks show grayscale area graphs indicating memory utilization over time, with varying levels throughout the execution. > > **Key Elements:** +> > - **MATMUL opcode**: Matrix multiplication operation highlighted > - **591 ns duration**: Time for this specific matrix operation > - **Source location**: Python file path and line 116 shown @@ -613,11 +613,9 @@ ssh -i ~/my-ec2.pem -L 3001:localhost:3001 -L 3002:localhost:3002 ubuntu@[PUBLIC > - **State Buffer Usage**: Memory utilization visualization > - **Custom-call.3**: XLA operator hierarchy reference > - **fp32_mode=LOW_HIGH**: Floating point precision mode -> - **128*128**: Matrix dimensions for the operation - - -* If you click on any engine instruction in the timeline with your mouse, you will see event details in a panel below the timeline. +> - **128\*128**: Matrix dimensions for the operation +- If you click on any engine instruction in the timeline with your mouse, you will see event details in a panel below the timeline. > **Figure: nki profiler 6** > @@ -627,6 +625,7 @@ ssh -i ~/my-ec2.pem -L 3001:localhost:3001 -L 3002:localhost:3002 ubuntu@[PUBLIC > > **Timeline Tracks (visible portion):** > The upper section shows multiple component tracks including: +> > - Sync(nc1), Sync(nc0): Synchronization events with colored markers > - Tensor(nc1), Tensor(nc0): Tensor engine activity with orange bars > - TensorMatrix tracks: Matrix operation indicators @@ -641,20 +640,22 @@ ssh -i ~/my-ec2.pem -L 3001:localhost:3001 -L 3002:localhost:3002 ubuntu@[PUBLIC > > **Field-Value Table:** > The table displays detailed information with Field and Value columns: +> > - ctv_pc: 29679 > - ctv_ind: 65535 > - duration_ns: 189 > - Engine: Tensor > - ete_wait_time_ns: (value not visible) > - fully_qualified_subgraph: nsp0 -> - hierarchyName: xla__op+locals+CallImpl_custom-call.3 -> - Nki_attrs: {"op_type":"xla__op_uls03+locals+u03eCallImpl","source_file":"/shared/ethan/nn/x_server/lib/python3.10/site-packages/torch_xla/core/xla_ops_registry.py","source_line":"44"} +> - hierarchyName: xla\_\_op+locals+CallImpl_custom-call.3 +> - Nki_attrs: {"op_type":"xla\_\_op_uls03+locals+u03eCallImpl","source_file":"/shared/ethan/nn/x_server/lib/python3.10/site-packages/torch_xla/core/xla_ops_registry.py","source_line":"44"} > - Nki_name: %custom-call.3 = custom-call(%transpose.14, %transpose.12, %constant.4) > - instructionId: 141955007395627936 > - instructionName: fp32_mode=LOW transpose_mode=DISABLED src=fp32@block32fhb 1,0,0|12*[1,1] 128*128 > - instructionType: REGULAR > > **Key Elements:** +> > - **Event Details tab**: Selected view showing instruction metadata > - **hierarchyName**: XLA operation reference (custom-call.3) > - **Nki_attrs**: JSON attributes including source file and line information @@ -664,9 +665,7 @@ ssh -i ~/my-ec2.pem -L 3001:localhost:3001 -L 3002:localhost:3002 ubuntu@[PUBLIC > - **Engine: Tensor**: Indicates this is a tensor engine operation > - **Blue selection bar**: Shows selected event region in timeline - -* To view hierarchy of this profile, click on Add Widget and select Hierarchy. - +- To view hierarchy of this profile, click on Add Widget and select Hierarchy. > **Figure: nki profiler 7** > @@ -675,6 +674,7 @@ ssh -i ~/my-ec2.pem -L 3001:localhost:3001 -L 3002:localhost:3002 ubuntu@[PUBLIC > This screenshot displays the Neuron Explorer view with the "Add Widget" dropdown menu expanded, revealing the various widget types that can be added to the profiler interface for analysis. > > **Header and Controls:** +> > - Title: "Neuron Explorer" in header bar > - Profile name: "mlp_with_mm_kernel" > - Search controls with category/field selectors and Submit/Clear result buttons @@ -683,6 +683,7 @@ ssh -i ~/my-ec2.pem -L 3001:localhost:3001 -L 3002:localhost:3002 ubuntu@[PUBLIC > > **Add Widget Dropdown Menu:** > The expanded dropdown shows the following widget options: +> > - **Search**: Widget for searching events > - **Hierarchy** (highlighted/selected): Shows operator hierarchy > - **Device Timeline**: Timeline visualization of device activity @@ -697,6 +698,7 @@ ssh -i ~/my-ec2.pem -L 3001:localhost:3001 -L 3002:localhost:3002 ubuntu@[PUBLIC > > **Device Timeline (background):** > Behind the dropdown, a partial view of the Device Timeline is visible showing tracks for: +> > - qSync100(nc1) and qSync100(nc0): Sync operations > - qGpSimdDynamic (nc1): GPSIMD dynamic operations > - qScalarDynamic (nc0): Scalar dynamic operations with colored activity markers @@ -705,11 +707,13 @@ ssh -i ~/my-ec2.pem -L 3001:localhost:3001 -L 3002:localhost:3002 ubuntu@[PUBLIC > - qScalarDynamic (nc0): Activity shown with markers > > **Left Sidebar:** +> > - Profile Manager > - Profile (selected) > - Summary > > **Key Elements:** +> > - **Add Widget dropdown**: Central feature showing all available widget types > - **Hierarchy option**: Currently highlighted/hovered option > - **AI Recommendation**: Notable feature for AI-assisted optimization @@ -718,9 +722,7 @@ ssh -i ~/my-ec2.pem -L 3001:localhost:3001 -L 3002:localhost:3002 ubuntu@[PUBLIC > - **Layout button**: For arranging widgets in the interface > - **Widget variety**: 11 different widget options available - -* Using the Profiler’s flexible layout support, you can drag and group every widget into any panel of your choice to customize the layout for your workflow. - +- Using the Profiler’s flexible layout support, you can drag and group every widget into any panel of your choice to customize the layout for your workflow. > **Figure: nki profiler 8** > @@ -730,10 +732,11 @@ ssh -i ~/my-ec2.pem -L 3001:localhost:3001 -L 3002:localhost:3002 ubuntu@[PUBLIC > > **Hierarchy View (Top Panel):** > The upper panel shows a hierarchical timeline visualization with a "Model" row on the left axis. The timeline displays operator execution blocks at different levels: -> - First level shows operators like "aten_mm..." and "xla__op+locals+CallImpl_custom-call.3" -> - A highlighted region shows "aten_view" -> - Later in the timeline: "xla__op+locals+CallImpl_custom-call.4" -> - At the end: "xla__op+locals+CallImpl..." (truncated) +> +> - First level shows operators like "aten_mm..." and "xla\_\_op+locals+CallImpl_custom-call.3" +> - A highlighted region shows "aten_view" +> - Later in the timeline: "xla\_\_op+locals+CallImpl_custom-call.4" +> - At the end: "xla\_\_op+locals+CallImpl..." (truncated) > - Colored blocks represent different operators: gray, cyan/teal, orange, and green blocks > > The hierarchy shows nested relationships between operations, with some operators containing sub-operations indicated by smaller blocks within larger ones. @@ -759,6 +762,7 @@ ssh -i ~/my-ec2.pem -L 3001:localhost:3001 -L 3002:localhost:3002 ubuntu@[PUBLIC > The two views are time-aligned, allowing users to correlate high-level operator execution (Hierarchy) with low-level engine activity (Device Timeline). > > **Key Elements:** +> > - **Hierarchy view**: Shows operator-level execution with nested relationships > - **Device Timeline**: Shows engine-level instruction traces > - **custom-call.3, custom-call.4**: XLA custom call operators visible in hierarchy @@ -767,9 +771,7 @@ ssh -i ~/my-ec2.pem -L 3001:localhost:3001 -L 3002:localhost:3002 ubuntu@[PUBLIC > - **Time alignment**: Both panels synchronized for correlation > - **~2.07 ms total**: Full kernel execution time span - -* If you right-click on an operator in the hierarchy timeline, it will highlight all related instructions in the instruction timeline. - +- If you right-click on an operator in the hierarchy timeline, it will highlight all related instructions in the instruction timeline. > **Figure: nki profiler 9** > @@ -779,13 +781,15 @@ ssh -i ~/my-ec2.pem -L 3001:localhost:3001 -L 3002:localhost:3002 ubuntu@[PUBLIC > > **Hierarchy View (Top Panel):** > The upper panel shows the operator hierarchy timeline with: +> > - "Model" label on the left axis -> - Operator blocks including "xla__op+locals+CallImpl_custom-call.3" highlighted/selected (shown in orange/golden color) +> - Operator blocks including "xla\_\_op+locals+CallImpl_custom-call.3" highlighted/selected (shown in orange/golden color) > - The selected operator spans a significant portion of the timeline > > **Operator Detail Popup (Golden/Yellow Background):** > A detailed popup appears for the selected operator showing: -> - Name: xla__op_CallImpl_custom-call.3 +> +> - Name: xla\_\_op_CallImpl_custom-call.3 > - Duration: 1.046ms > - Time: 225,365ns - 1,285,725ns > - Subgraph: 0 @@ -794,6 +798,7 @@ ssh -i ~/my-ec2.pem -L 3001:localhost:3001 -L 3002:localhost:3002 ubuntu@[PUBLIC > > **Device Timeline (Bottom Panel):** > The lower panel shows comprehensive engine-level traces: +> > - qScalarDynamic (nc0): Dense magenta activity > - qGpSimdDynamic: Activity markers > - qSync: Synchronization events @@ -815,6 +820,7 @@ ssh -i ~/my-ec2.pem -L 3001:localhost:3001 -L 3002:localhost:3002 ubuntu@[PUBLIC > The 64.3% MFU indicates the custom-call operation achieves relatively good utilization of the tensor engine's theoretical peak performance. The 1.046ms duration represents the bulk of the kernel's execution time. > > **Key Elements:** +> > - **MFU: 64.3%**: Model FLOPS Utilization for the selected operator > - **Duration: 1.046ms**: Operator execution time > - **custom-call.3**: XLA custom call operator selected @@ -823,7 +829,6 @@ ssh -i ~/my-ec2.pem -L 3001:localhost:3001 -L 3002:localhost:3002 ubuntu@[PUBLIC > - **Vector(nc0) activity**: Yellow bars showing vector operations > - **Time range**: 225,365ns to 1,285,725ns - ### View NKI Source Code in Neuron Profile You can optionally include your NKI source code files for display in Neuron Profile. When provided, Neuron Profile loads the source code into an integrated viewer, displayed side-by-side with the execution timeline in the web UI. This makes it easier to navigate between the instruction trace and the corresponding NKI source code, and to track the exact version of the code that generated the profile. @@ -831,12 +836,10 @@ You can optionally include your NKI source code files for display in Neuron Prof > **Note** > > Note -> -> +> > Even if you don’t upload the source code, the NKI source filename and line number remain available in the instruction detail view as noted in View Neuron Profile UI. -* If source code is uploaded with NEFF and NTFF file, you will be able to see the source code in the code editor. To open the code editor, click on **Add Widget** and select **Code Editor**. - +- If source code is uploaded with NEFF and NTFF file, you will be able to see the source code in the code editor. To open the code editor, click on **Add Widget** and select **Code Editor**. > **Figure: nki profiler 10** > @@ -845,6 +848,7 @@ You can optionally include your NKI source code files for display in Neuron Prof > This screenshot displays the Neuron Explorer interface with the Add Widget dropdown menu expanded, focusing on the Code Editor option for viewing NKI source code alongside profiling data. > > **Header and Layout:** +> > - Title: "Neuron Explorer" > - Profile: "mlp_with_mm_kernel" > - "+ Add Widget" button expanded (blue dropdown arrow) @@ -853,6 +857,7 @@ You can optionally include your NKI source code files for display in Neuron Prof > > **Add Widget Dropdown Menu:** > The expanded menu shows widget options with "Code Editor" highlighted (indicated by darker/selected background): +> > - Search > - Hierarchy > - Device Timeline @@ -867,31 +872,32 @@ You can optionally include your NKI source code files for display in Neuron Prof > > **Hierarchy View:** > Below the menu, the Hierarchy timeline is visible showing: +> > - "Model" row with operator blocks -> - "aten_zero_..." block at the start -> - "xla__op+locals+CallImpl_custom-call.3" (orange block) +> - "aten*zero*..." block at the start +> - "xla\_\_op+locals+CallImpl_custom-call.3" (orange block) > - "aten_p..." (cyan/teal block) -> - "xla__op+locals+CallImpl_custom-call" continuing on the right -> - "aten_v..." and "aten_..." blocks (colored in teal and green) +> - "xla\_\_op+locals+CallImpl_custom-call" continuing on the right +> - "aten*v..." and "aten*..." blocks (colored in teal and green) > > **Device Timeline (Partial View):** > At the bottom, a partial view of the Device Timeline shows: +> > - "qGpSimdDynamic (nc0)" track with blue/colored activity markers > > **Time Scale:** > The visible timeline spans from 0 to approximately 1,856,613 ns, with markers at 200,000, 400,000, 600,000, 800,000, 1,000,000, 1,200,000, and 1,400,000. > > **Key Elements:** +> > - **Code Editor option**: Highlighted widget selection for viewing source code > - **Add Widget menu**: Full list of 11 available widgets > - **Hierarchy view**: Shows operator-level execution timeline > - **custom-call.3**: Main NKI kernel operator visible -> - **aten_* operations**: Framework operations surrounding the custom call +> - **aten\_\* operations**: Framework operations surrounding the custom call > - **Integration purpose**: Code Editor enables viewing NKI source alongside traces - -* The code editor will be open on the right-hand side. - +- The code editor will be open on the right-hand side. > **Figure: nki profiler 11** > @@ -949,11 +955,13 @@ You can optionally include your NKI source code files for display in Neuron Prof > > **Hierarchy View (Top Left):** > Shows the operator hierarchy with: +> > - "Model" row containing operator blocks > - custom-call operations and aten operations visible > > **Device Timeline (Bottom Left):** > Shows multiple engine tracks: +> > - qScalarDynamic, qGpSimdDynamic: Activity markers > - Sync, Tensor, TensorMatrix, Vector, Scalar, GpSimd tracks > - Dense orange bars in Tensor(nc0) and yellow in Vector(nc0) @@ -962,17 +970,16 @@ You can optionally include your NKI source code files for display in Neuron Prof > Above the code editor: "EXPLORER" tab with "mlp_with_mm_kernel.py" file selected > > **Key Elements:** +> > - **Code Editor widget**: Displays NKI source code > - **@nki.jit decorator**: Platform target "trn2" visible -> - **TILES_IN_BLOCK_***: Tiling parameters for the matmul kernel +> - **TILES*IN_BLOCK*\***: Tiling parameters for the matmul kernel > - **Matrix dimensions**: M, K, N blocking strategy described > - **Source correlation**: Code visible alongside execution traces > - **mlp_with_mm_kernel.py**: The profiled NKI kernel file > - **Three-panel layout**: Hierarchy + Timeline + Code Editor - -* Hover on an instruction that has NKI source location and **Command + left click** on Mac (**Ctrl + right click** on Windows), and it will pop-up a window for showing file selection for stack trace. - +- Hover on an instruction that has NKI source location and **Command + left click** on Mac (**Ctrl + right click** on Windows), and it will pop-up a window for showing file selection for stack trace. > **Figure: nki profiler 12** > @@ -981,6 +988,7 @@ You can optionally include your NKI source code files for display in Neuron Prof > This screenshot shows a dark-themed modal dialog that appears when navigating between source code locations in the Neuron Profiler. The dialog enables jumping to specific lines in the NKI kernel source code that correspond to profiled instructions. > > **Dialog Header:** +> > - Title: "Select Source Location" > - Instructions: "Use up/down arrows to navigate, Enter to select, Esc to cancel" > @@ -989,16 +997,12 @@ You can optionally include your NKI source code files for display in Neuron Prof > > 1. **mlp_with_mm_kernel.py:116** (highlighted/selected with blue background) > - Subtitle: mlp_with_mm_kernel.py -> > 2. **mlp_with_mm_kernel.py:157** > - Subtitle: mlp_with_mm_kernel.py -> > 3. **mlp_with_mm_kernel.py:169** > - Subtitle: mlp_with_mm_kernel.py -> > 4. **mlp_with_mm_kernel.py:184** > - Subtitle: mlp_with_mm_kernel.py -> > 5. **mlp_with_mm_kernel.py:192** > - Subtitle: mlp_with_mm_kernel.py > @@ -1009,6 +1013,7 @@ You can optionally include your NKI source code files for display in Neuron Prof > The first option (line 116) is currently selected/highlighted with a blue background, indicating it will be the destination if Enter is pressed. > > **Key Elements:** +> > - **Select Source Location dialog**: Navigation modal for jumping to code lines > - **Line 116**: First and currently selected location (likely main kernel code) > - **Line 157, 169, 184, 192**: Additional source locations referenced by profile @@ -1016,9 +1021,7 @@ You can optionally include your NKI source code files for display in Neuron Prof > - **Keyboard navigation**: Arrow keys to select, Enter to confirm, Esc to cancel > - **Multiple references**: Shows that profiled instructions map to multiple source lines - -* Selecting any option from the list, it will jump to the line of the source code and highlight all of instructions related to this line. - +- Selecting any option from the list, it will jump to the line of the source code and highlight all of instructions related to this line. > **Figure: nki profiler 13** > @@ -1054,6 +1057,7 @@ You can optionally include your NKI source code files for display in Neuron Prof > ``` > > A yellow/gold highlighted line indicates the currently selected or referenced source location. The code shows NKI-specific constructs including: +> > - `range()` for loop iterations > - `nisa.dma_copy()` for DMA operations > - `nisa.nc_transpose()` for tensor transposition @@ -1067,11 +1071,13 @@ You can optionally include your NKI source code files for display in Neuron Prof > > **Device Timeline (Bottom Left):** > Shows engine-level traces with: +> > - Dense activity in Scalar and Tensor tracks > - Yellow/orange bars indicating tensor engine operations > - Multiple sync and dynamic operation tracks > > **Key Elements:** +> > - **range()**: NKI loop construct > - **nisa.dma_copy()**: DMA load operation > - **nisa.nc_transpose()**: Tensor transposition instruction @@ -1080,9 +1086,7 @@ You can optionally include your NKI source code files for display in Neuron Prof > - **Code minimap**: Overview of full source file > - **Loop nesting**: Shows M, K, N blocking structure for matmul - -* You can also enable different source code decorations in **Source Code Settings**. - +- You can also enable different source code decorations in **Source Code Settings**. > **Figure: nki profiler 14** > @@ -1092,12 +1096,14 @@ You can optionally include your NKI source code files for display in Neuron Prof > > **Tab Bar:** > The top shows two closeable tabs: +> > - "Instruction X" (closeable) > - "Settings X" (currently active, highlighted in blue) > - Window control icons (expand/close) on the right > > **Left Navigation Panel:** > Three settings categories listed vertically: +> > - Display Settings > - **Source Code Settings** (selected, highlighted with blue background and left border) > - Timeline Settings @@ -1105,34 +1111,39 @@ You can optionally include your NKI source code files for display in Neuron Prof > **Source Code Settings Options:** > > **Top-Level Options (with toggle switches):** +> > 1. **Source Code Time Range Decorations** (toggle: OFF) > - Likely shows time information inline with source code -> > 2. **Source Code Lowest Level Navigation** (toggle: OFF) > - Controls navigation granularity to lowest-level instructions > > **Source Code Navigation Section:** > Two toggle options for framework-specific navigation: +> > - **NKI** (toggle: ON, blue filled) > - **PyTorch** (toggle: ON, blue filled) > > **Source Code Decorations Section:** > Four toggle options controlling what information is shown in the code editor: +> > - **InstructionCount** (toggle: OFF) > - **FLOPS** (toggle: OFF) > - **Clicked** (toggle: ON, blue filled) > - **Dependencies** (toggle: ON, blue filled) > > **Toggle States:** +> > - Blue filled toggle = ON/enabled > - Gray/empty toggle = OFF/disabled > > The settings allow users to customize the profiler experience by: +> > - Enabling/disabling source code navigation for different frameworks > - Showing/hiding performance metrics inline with code > - Controlling visual decorations and dependency visualization > > **Key Elements:** +> > - **Source Code Settings tab**: Currently selected settings category > - **NKI toggle**: Enable NKI source code navigation (ON) > - **PyTorch toggle**: Enable PyTorch source code navigation (ON) @@ -1140,7 +1151,6 @@ You can optionally include your NKI source code files for display in Neuron Prof > - **Clicked/Dependencies toggles**: Enable click highlighting and dependency visualization (ON) > - **Time Range Decorations**: Show time information in code (OFF) - > **Figure: nki profiler 15** > > A Neuron Profiler Code Editor view showing NKI kernel source code with inline performance decorations displaying FLOPS counts, instruction counts, and engine breakdown information for specific code lines. @@ -1148,6 +1158,7 @@ You can optionally include your NKI source code files for display in Neuron Prof > This screenshot shows the Code Editor widget in the Neuron Profiler displaying detailed NKI kernel code with performance annotations (decorations) overlaid on specific source lines. > > **File Explorer Panel (Left):** +> > - "EXPLORER" header > - "mlp_with_mm_kernel.py" file selected > @@ -1178,14 +1189,12 @@ You can optionally include your NKI source code files for display in Neuron Prof > - InstructionCount: 4,096 > - Engine Breakdown: > - - Tensor: 4,096 instructions(100%) nc_range(TILES_IN_BLOCK_N): -> > 2. **For nisa.nc_matmul:** > - Shows "nisa.nc_matmul(" > - "dst=result_tile," -> - "stationary=lhs_f_tiles[bk_i][TILE_K, bn =" +> - "stationary=lhs_f_tiles[bk_i]TILE_K, bn =" > - "TILE_K*[bm * 1 + TILE_M," > - "moving=rhs_tiles[bk_i][TILE_K, bn * 1 + TILE_N]" -> > 3. **For nisa.tensor_tensor section:** > - "# Accumulate the result into the result_tmp tile." > - "nisa.tensor_tensor(data=result_tile[bm:bm+bm]," @@ -1197,6 +1206,7 @@ You can optionally include your NKI source code files for display in Neuron Prof > A code minimap shows the full file structure with highlighted regions indicating the current viewport position. > > **Key Elements:** +> > - **FLOPS: 17179869184**: Total floating-point operations for the matmul > - **InstructionCount: 4,096**: Number of tensor instructions > - **Engine Breakdown: Tensor 100%**: All instructions execute on tensor engine @@ -1206,13 +1216,12 @@ You can optionally include your NKI source code files for display in Neuron Prof > - **Inline decorations**: Yellow highlighted performance metrics > - **Loop structure**: M, K, N blocking visible in code - ## Next Steps Great! Now that you’ve learned how to profile an NKI kernel, it’s time to take this further: -* Dive into the NKI Performance Guide to discover techniques for making your kernels faster and more efficient. +- Dive into the NKI Performance Guide to discover techniques for making your kernels faster and more efficient. -* Explore the [NKI sample kernels](https://github.com/aws-neuron/nki-samples) to see real-world examples of high-performance kernel implementations — and get inspiration for your own NKI kernels. +- Explore the [NKI sample kernels](https://github.com/aws-neuron/nki-samples) to see real-world examples of high-performance kernel implementations — and get inspiration for your own NKI kernels. -By combining profiling insights with optimization strategies and practical examples, you’ll be well-equipped to write NKI kernels that leverage Neuron hardware in an efficient way. \ No newline at end of file +By combining profiling insights with optimization strategies and practical examples, you’ll be well-equipped to write NKI kernels that leverage Neuron hardware in an efficient way. diff --git a/skills/neuron-nki-docs/references/programming/api/api-nki-collectives.md b/skills/neuron-nki-docs/references/programming/api/api-nki-collectives.md index 4cbbf6d..4d3f69e 100644 --- a/skills/neuron-nki-docs/references/programming/api/api-nki-collectives.md +++ b/skills/neuron-nki-docs/references/programming/api/api-nki-collectives.md @@ -16,6 +16,7 @@ Collective Communication instructions. **Engine:** DMA Engine **Signature:** + ```python collectives.all_gather(srcs, dsts, replica_group, collective_dim, priority=None, name=None) ``` @@ -34,11 +35,11 @@ communication (multiple tensors) is only supported when tensors are on HBM. - **dsts** — List of output tensors to store results - **replica_group** — ReplicaGroup defining rank groups for the collective - **collective_dim** — Dimension along which output tensors are concatenated. - Currently only 0 is supported for HBM tensors. For SBUF tensors, 0 or 1 is - supported as SBUF collectives currently only operate on 2D tensors with a - single free dimension. + Currently only 0 is supported for HBM tensors. For SBUF tensors, 0 or 1 is + supported as SBUF collectives currently only operate on 2D tensors with a + single free dimension. - **priority** — DMA quality-of-service priority level 0-3 where lower is higher - priority (NeuronCore-v4+ only) + priority (NeuronCore-v4+ only) - **name** — (optional) name for the instruction. --- @@ -50,6 +51,7 @@ communication (multiple tensors) is only supported when tensors are on HBM. **Engine:** DMA Engine **Signature:** + ```python collectives.all_reduce(srcs, dsts, replica_group, op, priority=None, name=None) ``` @@ -69,7 +71,7 @@ communication (multiple tensors) is only supported when tensors are on HBM. - **replica_group** — ReplicaGroup defining rank groups for the collective - **op** — The reduction operation to perform (`nl.add`, `nl.minimum`, or `nl.maximum`) - **priority** — DMA quality-of-service priority level 0-3 where lower is higher - priority (NeuronCore-v4+ only) + priority (NeuronCore-v4+ only) - **name** — (optional) name for the instruction. --- @@ -81,6 +83,7 @@ communication (multiple tensors) is only supported when tensors are on HBM. **Engine:** DMA Engine **Signature:** + ```python collectives.all_to_all(srcs, dsts, replica_group, collective_dim, priority=None, name=None) ``` @@ -97,9 +100,9 @@ Tensors must reside on HBM. SBUF is not currently supported for all-to-all. - **dsts** — List of output tensors to store results - **replica_group** — ReplicaGroup defining rank groups for the collective - **collective_dim** — Dimension along which input tensors are split and output tensors are concatenated. - Currently only 0 is supported. + Currently only 0 is supported. - **priority** — DMA quality-of-service priority level 0-3 where lower is higher - priority (NeuronCore-v4+ only) + priority (NeuronCore-v4+ only) - **name** — (optional) name for the instruction. --- @@ -111,6 +114,7 @@ Tensors must reside on HBM. SBUF is not currently supported for all-to-all. **Engine:** DMA Engine **Signature:** + ```python collectives.all_to_all_v(srcs, dsts, replica_group, metadata_tensor, recv_counts_known=False, has_rdispls=False, priority=None, name=None) ``` @@ -127,7 +131,7 @@ metadata tensor, making per-rank payload sizes dynamic. **Current restrictions:** On instances with a NeuronSwitch fabric (see `Trn3 architecture -`_), +`\_), `all_to_all_v` requires LNC=2 and more than one participating device. Multiple ranks per device are supported, but for every replica-group rank-list, every device participating in that @@ -142,54 +146,51 @@ ranks from different nodes (a node refers to a different Trn EC2 instance). - **srcs** — Input tensor list. Currently supports exactly one tensor. - Must be HBM-backed. + Must be HBM-backed. - **dsts** — Output tensor list. Currently supports exactly one tensor. - Must be HBM-backed. `src` and `dst` element counts can be - different; sizes are validated against the metadata at execution - time. + Must be HBM-backed. `src` and `dst` element counts can be + different; sizes are validated against the metadata at execution + time. - **replica_group** — ReplicaGroup defining which ranks participate. - **metadata_tensor** — `uint32` tensor laid out contiguously in - memory. Shape depends on backing buffer, where `rows` is 3 when - `has_rdispls=False` and 4 when `has_rdispls=True`: - - - HBM: `(rows, replica_group_size)`. - - SBUF: `(1, rows, replica_group_size)` — the whole buffer must - live on a single partition, so a trivial partition dim is - prepended. - - For each other rank `r` in the replica group, the rows are: - - - Row 0 `send_counts[r]`: number of elements sent to rank `r`. - Always an input. - - Row 1 `send_displs[r]`: offset in elements within `src` where - the chunk destined for rank `r` begins. Always an input. - - Row 2 `recv_counts[r]`: number of elements received from rank - `r`. Controlled by `recv_counts_known` — see that flag. - - Row 3 `recv_displs[r]`: offset in elements within `dst` where - the chunk from rank `r` is written. Only present when - `has_rdispls=True`. - -- **recv_counts_known** — - Controls whether row 2 is populated by the collective during - execution. Row 2 is never read as input. - - - `True`: row 2 is left untouched, avoiding a small per-rank - writeback. - - `False` (default): row 2 is an **output** — per-rank received - counts are written during execution, and can be read after the - op to learn received sizes. - -- **has_rdispls** — - - `True`: row 3 is an **input**; recv_displs must be populated. - The chunk from sender rank `r` is written at - `dst[recv_displs[r] : recv_displs[r] + recv_counts[r]]`. - - `False` (default): row 3 may be omitted from `metadata_tensor` (pass a - 3-row tensor). Incoming chunks are laid out equally-spaced at - `recv_displs[r] = (dst.total_elements / replica_group_size) * r`, - regardless of the actual recv_count per rank. + memory. Shape depends on backing buffer, where `rows` is 3 when + `has_rdispls=False` and 4 when `has_rdispls=True`: + - HBM: `(rows, replica_group_size)`. + - SBUF: `(1, rows, replica_group_size)` — the whole buffer must + live on a single partition, so a trivial partition dim is + prepended. + + For each other rank `r` in the replica group, the rows are: + - Row 0 `send_counts[r]`: number of elements sent to rank `r`. + Always an input. + - Row 1 `send_displs[r]`: offset in elements within `src` where + the chunk destined for rank `r` begins. Always an input. + - Row 2 `recv_counts[r]`: number of elements received from rank + `r`. Controlled by `recv_counts_known` — see that flag. + - Row 3 `recv_displs[r]`: offset in elements within `dst` where + the chunk from rank `r` is written. Only present when + `has_rdispls=True`. + +- **recv_counts_known** — + Controls whether row 2 is populated by the collective during + execution. Row 2 is never read as input. + - `True`: row 2 is left untouched, avoiding a small per-rank + writeback. + - `False` (default): row 2 is an **output** — per-rank received + counts are written during execution, and can be read after the + op to learn received sizes. + +- **has_rdispls** — + - `True`: row 3 is an **input**; recv_displs must be populated. + The chunk from sender rank `r` is written at + `dst[recv_displs[r] : recv_displs[r] + recv_counts[r]]`. + - `False` (default): row 3 may be omitted from `metadata_tensor` (pass a + 3-row tensor). Incoming chunks are laid out equally-spaced at + `recv_displs[r] = (dst.total_elements / replica_group_size) * r`, + regardless of the actual recv_count per rank. - **priority** — DMA QoS priority level 0-3 where lower is higher - priority (NeuronCore-v4+ only). + priority (NeuronCore-v4+ only). - **name** — (optional) name for the instruction. --- @@ -201,6 +202,7 @@ instance). **Engine:** DMA Engine **Signature:** + ```python collectives.collective_permute(srcs, dsts, source_target_pairs, priority=None, name=None) ``` @@ -223,7 +225,7 @@ each list parameter must contain exactly one tensor. - **dsts** — List of destination tensors to receive into - **source_target_pairs** — List of (source, target) rank ID pairs - **priority** — DMA quality-of-service priority level 0-3 where lower is higher - priority (NeuronCore-v4+ only) + priority (NeuronCore-v4+ only) - **name** — (optional) name for the instruction. --- @@ -235,6 +237,7 @@ each list parameter must contain exactly one tensor. **Engine:** DMA Engine **Signature:** + ```python collectives.collective_permute_implicit(srcs_by_channel, dsts_by_channel, replica_group, channel_ids=[0], priority=None, name=None) ``` @@ -264,9 +267,9 @@ and 2 for other supported replica groups. - **dsts_by_channel** — List of destination tensor lists, one per channel. Each inner list must contain exactly one tensor. - **replica_group** — ReplicaGroup defining rank groups for the collective - **channel_ids** — List of channel IDs to use for communication (default [0] for single channel). - Currently must be consecutive integers starting from 0. + Currently must be consecutive integers starting from 0. - **priority** — DMA quality-of-service priority level 0-3 where lower is higher - priority (NeuronCore-v4+ only) + priority (NeuronCore-v4+ only) - **name** — (optional) name for the instruction. --- @@ -276,6 +279,7 @@ and 2 for other supported replica groups. `nki.collectives.collective_permute_implicit_current_processing_rank_id(iteration_id, replica_group, channel_id, name)` **Signature:** + ```python collectives.collective_permute_implicit_current_processing_rank_id(iteration_id, replica_group, channel_id=0, name=None) ``` @@ -312,7 +316,7 @@ and 2 for other supported replica groups. - **replica_group** — ReplicaGroup defining the ring topology - **channel_id** — Channel ID for the communication (0 to num_channels-1) - **name** — (optional) name for the instruction. -**Returns:** Scalar register containing the rank ID of the data to be processed + **Returns:** Scalar register containing the rank ID of the data to be processed --- @@ -323,6 +327,7 @@ and 2 for other supported replica groups. **Engine:** DMA Engine **Signature:** + ```python collectives.collective_permute_implicit_reduce(srcs0_by_channel, srcs1_by_channel, dsts_by_channel, replica_group, op, channel_ids=[0], priority=None, name=None) ``` @@ -356,9 +361,9 @@ and 2 for other supported replica groups. - **replica_group** — ReplicaGroup defining rank groups for the collective - **op** — The reduction operation to perform (`nl.add`, `nl.minimum`, or `nl.maximum`) - **channel_ids** — List of channel IDs to use for communication (default [0] for single channel). - Currently must be consecutive integers starting from 0. + Currently must be consecutive integers starting from 0. - **priority** — DMA quality-of-service priority level 0-3 where lower is higher - priority (NeuronCore-v4+ only) + priority (NeuronCore-v4+ only) - **name** — (optional) name for the instruction. --- @@ -368,6 +373,7 @@ and 2 for other supported replica groups. `nki.collectives.rank_id(name)` **Signature:** + ```python collectives.rank_id(name=None) ``` @@ -375,7 +381,7 @@ collectives.rank_id(name=None) Get the rank ID of the current rank. - **name** — (optional) name for the instruction. -**Returns:** The rank ID of the current rank within the collective group + **Returns:** The rank ID of the current rank within the collective group --- @@ -386,6 +392,7 @@ Get the rank ID of the current rank. **Engine:** DMA Engine **Signature:** + ```python collectives.reduce_scatter(srcs, dsts, replica_group, collective_dim, op, priority=None, name=None) ``` @@ -404,10 +411,10 @@ communication (multiple tensors) is only supported when tensors are on HBM. - **dsts** — List of output tensors to store results - **replica_group** — ReplicaGroup defining rank groups for the collective - **collective_dim** — Dimension along which input tensors are split. - Currently only 0 is supported for both HBM and SBUF tensors. + Currently only 0 is supported for both HBM and SBUF tensors. - **op** — The reduction operation to perform (`nl.add`, `nl.minimum`, or `nl.maximum`) - **priority** — DMA quality-of-service priority level 0-3 where lower is higher - priority (NeuronCore-v4+ only) + priority (NeuronCore-v4+ only) - **name** — (optional) name for the instruction. --- @@ -419,6 +426,7 @@ communication (multiple tensors) is only supported when tensors are on HBM. **Engine:** DMA Engine **Signature:** + ```python collectives.all_gather_v(srcs, dsts, replica_group, metadata_tensor, recv_counts_known=False, has_rdispls=False, priority=None, name=None) ``` @@ -450,59 +458,57 @@ rows 2/3. - Each replica subgroup must have exactly 4 ranks (intra-chip). - **srcs** — Input tensor list. Currently supports exactly one tensor. - Must be HBM-backed. + Must be HBM-backed. - **dsts** — Output tensor list. Currently supports exactly one tensor. - Must be HBM-backed. `src` and `dst` element counts are free to - differ; sizes are validated against the metadata at execution time. + Must be HBM-backed. `src` and `dst` element counts are free to + differ; sizes are validated against the metadata at execution time. - **replica_group** — ReplicaGroup defining which ranks participate. - **metadata_tensor** — `uint32` tensor laid out contiguously in - memory. Shape depends on backing buffer, where `rows` is 3 when - `has_rdispls=False` and 4 when `has_rdispls=True`: - - - HBM: `(rows, replica_group_size)`. - - SBUF: `(1, rows, replica_group_size)` — the whole buffer must - live on a single partition, so a trivial partition dim is - prepended. - - Rows 0/1 are single-valued for all-gather: only their first - column is read. - - The rows are: - - - Row 0 `send_count`: number of elements in the chunk broadcast - to every rank. Only the first column is read; the same count - applies to all destinations. Always an input. - - Row 1 `send_displ`: offset in elements within `src` where the - broadcast chunk begins. Only the first column is read; the same - displacement applies to all destinations. Always an input. - - Row 2 `recv_counts[r]`: number of elements received from rank - `r`. Per-src-rank. Controlled by `recv_counts_known` — see - that flag. - - Row 3 `recv_displs[r]`: offset in elements within `dst` where - the chunk from rank `r` is written. Per-src-rank. Only present - when `has_rdispls=True`. - -- **recv_counts_known** — - Controls whether row 2 is populated by the collective during - execution. Row 2 is never read as input. - - - `True`: row 2 is left untouched, avoiding a small per-rank - writeback. - - `False` (default): row 2 is an **output** — per-rank received - counts are written during execution, and can be read after the - op to learn received sizes. -- **has_rdispls** — - - `True`: row 3 is an **input**; recv_displs must be populated. - The chunk from sender rank `r` is written at - `dst[recv_displs[r] : recv_displs[r] + recv_counts[r]]`. - - `False` (default): row 3 may be omitted from `metadata_tensor` - (pass a 3-row tensor). Incoming chunks are laid out - equally-spaced at - `block_offset(r) = dst.total_elements / replica_group_size * r`, - regardless of the actual recv_count per rank. + memory. Shape depends on backing buffer, where `rows` is 3 when + `has_rdispls=False` and 4 when `has_rdispls=True`: + - HBM: `(rows, replica_group_size)`. + - SBUF: `(1, rows, replica_group_size)` — the whole buffer must + live on a single partition, so a trivial partition dim is + prepended. + + Rows 0/1 are single-valued for all-gather: only their first + column is read. + + The rows are: + - Row 0 `send_count`: number of elements in the chunk broadcast + to every rank. Only the first column is read; the same count + applies to all destinations. Always an input. + - Row 1 `send_displ`: offset in elements within `src` where the + broadcast chunk begins. Only the first column is read; the same + displacement applies to all destinations. Always an input. + - Row 2 `recv_counts[r]`: number of elements received from rank + `r`. Per-src-rank. Controlled by `recv_counts_known` — see + that flag. + - Row 3 `recv_displs[r]`: offset in elements within `dst` where + the chunk from rank `r` is written. Per-src-rank. Only present + when `has_rdispls=True`. + +- **recv_counts_known** — + Controls whether row 2 is populated by the collective during + execution. Row 2 is never read as input. + - `True`: row 2 is left untouched, avoiding a small per-rank + writeback. + - `False` (default): row 2 is an **output** — per-rank received + counts are written during execution, and can be read after the + op to learn received sizes. + +- **has_rdispls** — + - `True`: row 3 is an **input**; recv_displs must be populated. + The chunk from sender rank `r` is written at + `dst[recv_displs[r] : recv_displs[r] + recv_counts[r]]`. + - `False` (default): row 3 may be omitted from `metadata_tensor` + (pass a 3-row tensor). Incoming chunks are laid out + equally-spaced at + `block_offset(r) = dst.total_elements / replica_group_size * r`, + regardless of the actual recv_count per rank. - **priority** — DMA QoS priority level 0-3 where lower is higher - priority (NeuronCore-v4+ only). + priority (NeuronCore-v4+ only). - **name** — (optional) name for the instruction. --- diff --git a/skills/neuron-nki-docs/references/programming/api/api-nki-isa-local-collective.md b/skills/neuron-nki-docs/references/programming/api/api-nki-isa-local-collective.md index 2910a01..1c28d6d 100644 --- a/skills/neuron-nki-docs/references/programming/api/api-nki-isa-local-collective.md +++ b/skills/neuron-nki-docs/references/programming/api/api-nki-isa-local-collective.md @@ -16,6 +16,7 @@ Logical NeuronCore (LNC) instructions. **Engine:** GpSimd Engine **Signature:** + ```python isa.core_barrier(data, cores, engine=engine_enum.gpsimd, name=None) ``` @@ -24,7 +25,6 @@ Synchronize execution across multiple NeuronCores by implementing a barrier mech > **Note:** > Available only on NeuronCore-v3 or newer. -> This instruction creates a synchronization point where all specified NeuronCores must reach before any can proceed. The barrier is implemented using a semaphore-based protocol @@ -83,6 +83,7 @@ nisa.core_barrier(data=shared_tensor, cores=(0, 1)) **Engine:** DMA Engine **Signature:** + ```python isa.sendrecv(src, dst, send_to_rank, recv_from_rank, pipe_id, dma_engine=dma_engine_enum.dma, name=None) ``` @@ -92,7 +93,6 @@ simultaneously using DMA engines. > **Note:** > Available only on NeuronCore-v3 or newer. -> This instruction enables bidirectional data exchange between two NeuronCores within a Logical NeuronCore (LNC) configuration. @@ -120,10 +120,9 @@ The `dma_engine` parameter specifies which DMA transfer mechanism to use: - `nisa.dma_engine.gpsimd_dma`: Uses the GPSIMD's internal DMA engine for low-latency SB-to-SB swaps in LNC=2. Implies GPSIMD as the trigger engine. This mode restricts the data size per partition to not exceed: - - - 1024 bytes for 32-bit types - - 512 bytes for 16-bit types - - 256 bytes for 8-bit types + - 1024 bytes for 32-bit types + - 512 bytes for 16-bit types + - 256 bytes for 8-bit types **Constraints.** diff --git a/skills/neuron-nki-docs/references/programming/api/api-nki-isa-memory.md b/skills/neuron-nki-docs/references/programming/api/api-nki-isa-memory.md index fe01703..0c7670d 100644 --- a/skills/neuron-nki-docs/references/programming/api/api-nki-isa-memory.md +++ b/skills/neuron-nki-docs/references/programming/api/api-nki-isa-memory.md @@ -16,6 +16,7 @@ DMA and memory management instructions. **Engine:** DMA Engine **Signature:** + ```python isa.dma_compute(dst, srcs, reduce_op, scales=None, unique_indices=True, oob_mode=oob_mode_enum.error, name=None) ``` @@ -41,19 +42,22 @@ When one of the source tensors has a `vector_offset` (indirect indexing), `dma_compute` performs read-modify-write with two modes: **Scatter RMW**: `dst(HBM)[indices] = dst(HBM)[indices] + src(SB)` - - `dst` is in HBM with indirect indexing - - One source matches `dst` and has `vector_offset` - - The other source is data in SBUF + +- `dst` is in HBM with indirect indexing +- One source matches `dst` and has `vector_offset` +- The other source is data in SBUF **Gather RMW**: `dst(SB) = dst(SB) + src(HBM)[indices]` - - `dst` is in SBUF - - One source is data in HBM with `vector_offset` - - The other source matches `dst` + +- `dst` is in SBUF +- One source is data in HBM with `vector_offset` +- The other source matches `dst` Both modes require: - - Exactly 2 source tensors - - All `scales` must be `1.0` (or `None`) - - `unique_indices` must be `True` (non-unique indices not yet supported) + +- Exactly 2 source tensors +- All `scales` must be `1.0` (or `None`) +- `unique_indices` must be `True` (non-unique indices not yet supported) The only supported DGE mode for read-modify-write (scatter/gather) is SW DGE. For `dma_compute` without `vector_offset`, the only supported DGE mode is None (static DMA). @@ -85,27 +89,25 @@ The max number of source tensors in `srcs` is 16. - **srcs** — a list of input tensors to be scaled and reduced - **reduce_op** — the reduction operation to apply (currently only `nl.add` is supported) - **scales** — (optional) a list of scale factors corresponding to each - tensor in `srcs`. Must be all 1.0 if provided. - Defaults to None (equivalent to [1.0, 1.0, ...]). + tensor in `srcs`. Must be all 1.0 if provided. + Defaults to None (equivalent to [1.0, 1.0, ...]). - **unique_indices** — (optional) Whether scatter indices are unique. - Must be True when using vector_offset (non-unique - not yet supported). Default: True. + Must be True when using vector_offset (non-unique + not yet supported). Default: True. - **oob_mode** — (optional) Specifies how to handle out-of-bounds (oob) - array indices during indirect access operations. Valid - modes are: + array indices during indirect access operations. Valid + modes are: + - `oob_mode.error`: (Default) Raises an error when encountering + out-of-bounds indices. + - `oob_mode.skip`: Silently skips any operations involving + out-of-bounds indices. - - `oob_mode.error`: (Default) Raises an error when encountering - out-of-bounds indices. - - `oob_mode.skip`: Silently skips any operations involving - out-of-bounds indices. - - For example, when using indirect gather/scatter operations with - `vector_offset`, out-of-bounds indices can occur if the index - array contains values that exceed the dimensions of the target array. + For example, when using indirect gather/scatter operations with + `vector_offset`, out-of-bounds indices can occur if the index + array contains values that exceed the dimensions of the target array. --- - ### nki.isa.dma_copy {#nki-isa-dma_copy} `nki.isa.dma_copy(dst, src, priority, oob_mode, dge_mode, engine, name)` @@ -113,6 +115,7 @@ The max number of source tensors in `srcs` is 16. **Engine:** Scalar Engine **Signature:** + ```python isa.dma_copy(dst, src, priority=None, oob_mode=oob_mode_enum.error, dge_mode=dge_mode_enum.unknown, engine=engine_enum.unknown, name=None) ``` @@ -145,9 +148,9 @@ address or index when it is out of bound using `oob_mode=oob_mode.skip`. Both `src` and `dst` tiles can be in HBM or SBUF. However, if both tiles are in SBUF, consider using an alternative for better performance: -- nisa.tensor_copy for direct copies -- nisa.nc_n_gather to gather elements within each partition independently -- nisa.local_gather to gather elements within groups of partitions +- nisa.tensor_copy for direct copies +- nisa.nc_n_gather to gather elements within each partition independently +- nisa.local_gather to gather elements within groups of partitions **Data types.** @@ -170,13 +173,13 @@ parameter. There are two types of indirect addressing: -*Vector indirection* provides per-partition dynamic offsets. Each of the hardware partitions +_Vector indirection_ provides per-partition dynamic offsets. Each of the hardware partitions gets its own index, enabling gather/scatter where different partitions access different rows. Use `.ap(pattern=..., vector_offset=idx_tensor, indirect_dim=0)` where `idx_tensor` is an SBUF tensor of shape `(P, 1)` containing one row index per partition. The tensor being indexed (the one `.ap()` is called on) must be in HBM. -*Scalar indirection* provides a single dynamic offset applied uniformly to all partitions. +_Scalar indirection_ provides a single dynamic offset applied uniformly to all partitions. Use `.ap(pattern=..., scalar_offset=reg_or_tensor, indirect_dim=N)` where the offset is either a 1x1 SBUF tensor or a `VirtualRegister` from `nisa.register_alloc()`. @@ -192,7 +195,7 @@ The hardware reads the index tensor in column-major order DMA Batching has the following hardware-imposed restrictions: #. The 2D vector_offset must be on `src` (gather); a 2D vector_offset on `dst` - (multi-column scatter) is not supported. +(multi-column scatter) is not supported. #. When `M > 1`, `P` must be exactly 128 #. Both `src` and `dst` tensors must be contiguous in memory. #. `src` and `dst` must have the same dtype. @@ -259,16 +262,15 @@ def indirect_scatter_kernel(src_data, indices, output): - **dst** — the destination tensor to copy data into - **src** — the source tensor to copy data from - **priority** — (optional): DMA quality-of-service priority level 0-3 where lower is higher priority (NeuronCore-v4+ only) -- **dge_mode** — (optional) specify which Descriptor Generation Engine (DGE) mode to use for DMA descriptor generation: `nki.isa.dge_mode.none` (turn off DGE) or `nki.isa.dge_mode.swdge` (software DGE) or `nki.isa.dge_mode.hwdge` (hardware DGE) or `nki.isa.dge_mode.unknown` (by default, let compiler select the best DGE mode). Hardware based DGE is only supported for NeuronCore-v3 or newer. See [Trainium2 arch guide](../../architecture/trainium2_arch.md) for more information. +- **dge_mode** — (optional) specify which Descriptor Generation Engine (DGE) mode to use for DMA descriptor generation: `nki.isa.dge_mode.none` (turn off DGE) or `nki.isa.dge_mode.swdge` (software DGE) or `nki.isa.dge_mode.hwdge` (hardware DGE) or `nki.isa.dge_mode.unknown` (by default, let compiler select the best DGE mode). Hardware based DGE is only supported for NeuronCore-v3 or newer. See [Trainium2 arch guide](../../architecture/trainium2_arch.md) for more information. - **oob_mode** — (optional) Specifies how to handle out-of-bounds (oob) array indices during indirect access operations. Valid modes are: + - `oob_mode.error`: (Default) Raises an error when encountering out-of-bounds indices. + - `oob_mode.skip`: Silently skips any operations involving out-of-bounds indices. - - `oob_mode.error`: (Default) Raises an error when encountering out-of-bounds indices. - - `oob_mode.skip`: Silently skips any operations involving out-of-bounds indices. - - For example, when using indirect gather/scatter operations, out-of-bounds indices can occur if the index array contains values that exceed the dimensions of the target array. + For example, when using indirect gather/scatter operations, out-of-bounds indices can occur if the index array contains values that exceed the dimensions of the target array. - **engine** — (optional) the engine to use for HWDGE descriptor generation: `nki.isa.engine.sync` or `nki.isa.engine.scalar`. - Only valid when `dge_mode=nisa.dge_mode.hwdge`. `nki.isa.engine.unknown` by default. + Only valid when `dge_mode=nisa.dge_mode.hwdge`. `nki.isa.engine.unknown` by default. --- @@ -279,6 +281,7 @@ def indirect_scatter_kernel(src_data, indices, output): **Engine:** DMA Engine **Signature:** + ```python isa.dma_transpose(dst, src, axes=None, priority=None, dge_mode=dge_mode_enum.unknown, oob_mode=oob_mode_enum.error, name=None) ``` @@ -425,14 +428,13 @@ def gather_transpose_4d_kernel(src_hbm, idx_hbm): - **dst** — the destination of transpose, must be a tile in SBUF. - **src** — the source of transpose, must be a tile in HBM or SBUF. `src.dtype == dst.dtype` - **axes** — transpose axes where the i-th axis of the transposed tile will correspond to the axes[i] of the source. - Supported axes are `(1, 0)`, `(2, 1, 0)`, and `(3, 1, 2, 0)`. + Supported axes are `(1, 0)`, `(2, 1, 0)`, and `(3, 1, 2, 0)`. - **priority** — (optional): DMA quality-of-service priority level 0-3 where lower is higher priority (NeuronCore-v4+ only) -- **dge_mode** — (optional) specify which Descriptor Generation Engine (DGE) mode to use for DMA descriptor generation: `nki.isa.dge_mode.none` (turn off DGE) or `nki.isa.dge_mode.swdge` (software DGE) or `nki.isa.dge_mode.hwdge` (hardware DGE) or `nki.isa.dge_mode.unknown` (by default, let compiler select the best DGE mode). Hardware based DGE is only supported for NeuronCore-v3 or newer. See [Trainium2 arch guide](../../architecture/trainium2_arch.md) for more information. +- **dge_mode** — (optional) specify which Descriptor Generation Engine (DGE) mode to use for DMA descriptor generation: `nki.isa.dge_mode.none` (turn off DGE) or `nki.isa.dge_mode.swdge` (software DGE) or `nki.isa.dge_mode.hwdge` (hardware DGE) or `nki.isa.dge_mode.unknown` (by default, let compiler select the best DGE mode). Hardware based DGE is only supported for NeuronCore-v3 or newer. See [Trainium2 arch guide](../../architecture/trainium2_arch.md) for more information. - **oob_mode** — (optional) Specifies how to handle runtime out-of-bounds (oob) array indices during indirect access operations. Valid modes are: + - `oob_mode.error`: (Default) Raises an error when encountering runtime out-of-bounds indices. - - `oob_mode.error`: (Default) Raises an error when encountering runtime out-of-bounds indices. - - - `oob_mode.skip`: Silently skips any operations involving out-of-bounds indices. Only valid when `src` uses indirect indexing. + - `oob_mode.skip`: Silently skips any operations involving out-of-bounds indices. Only valid when `src` uses indirect indexing. --- @@ -442,12 +444,12 @@ def gather_transpose_4d_kernel(src_hbm, idx_hbm): nki.isa.local_gather -nki.isa.local_gather(*dst*, *src_buffer*, *index*, *num_elem_per_idx=1*, *num_valid_indices=None*, *name=None*)[[source]](../../../_modules/nki/isa.html#local_gather) +nki.isa.local*gather(\_dst*, _src_buffer_, _index_, _num_elem_per_idx=1_, _num_valid_indices=None_, _name=None_)[[source]](../../../\_modules/nki/isa.html#local_gather) Gather SBUF data in `src_buffer` using `index` on GpSimd Engine. Each of the eight GpSimd cores in GpSimd Engine connects to 16 contiguous SBUF partitions (e.g., core[0] connected to partition[0:16]) and performs gather from the connected 16 -SBUF partitions *independently* in parallel. The indices used for gather on each core should also +SBUF partitions _independently_ in parallel. The indices used for gather on each core should also come from the same 16 connected SBUF partitions. During execution of the instruction, each GpSimd core reads a 16-partition slice from `index`, flattens @@ -459,7 +461,7 @@ is not a multiple of 16, users can explicitly specify the valid index count per Note, `num_valid_indices` must not exceed the total element count in each 16-partition `index` slice (i.e., `num_valid_indices <= index.size / (index.shape[0] / 16)`). -Next, each GpSimd core uses the flattened `indices_1d` indices as *partition offsets* to gather from +Next, each GpSimd core uses the flattened `indices_1d` indices as _partition offsets_ to gather from the connected 16-partition slice of `src_buffer`. Optionally, this API also allows gathering of multiple contiguous elements starting at each index to improve gather throughput, as indicated by `num_elem_per_idx`. Behavior of out-of-bound index access is undefined. @@ -472,7 +474,6 @@ users can generate indices into 16 partitions, replicate them eight times to 128 As an example, if `src_buffer` is (128, 512) in shape and `index` is (128, 4) in shape, where the partition dimension size is 128, `local_gather` effectively performs the following operation: - ```python num_gpsimd_cores = 8 num_partitions_per_core = 16 @@ -505,32 +506,31 @@ for i_core in range(num_gpsimd_cores): output_np = output_np.reshape(output_shape) ``` - `local_gather` preserves the input data types from `src_buffer` in the gather output. Therefore, no data type casting is allowed in this API. The indices in `index` tile must be uint16 types. This API has three tile size constraints [subject to future relaxation]: -* The partition axis size of `src_buffer` must match that of `index` and must -be a multiple of 16. In other words, `src_buffer.shape[0] == index.shape[0] and src_buffer.shape[0] % 16 == 0`. +- The partition axis size of `src_buffer` must match that of `index` and must + be a multiple of 16. In other words, `src_buffer.shape[0] == index.shape[0] and src_buffer.shape[0] % 16 == 0`. -* The number of contiguous elements to gather per index per partition `num_elem_per_idx` -must be one of the following values: `[1, 2, 4, 8, 16, 32]`. +- The number of contiguous elements to gather per index per partition `num_elem_per_idx` + must be one of the following values: `[1, 2, 4, 8, 16, 32]`. -* The number of indices for gather per core must be less than or equal to 4096. +- The number of indices for gather per core must be less than or equal to 4096. Parameters: -* **dst** – an output tile of the gathered data +- **dst** – an output tile of the gathered data -* **src_buffer** – an input tile for gathering. +- **src_buffer** – an input tile for gathering. -* **index** – an input tile with indices used for gathering. +- **index** – an input tile with indices used for gathering. -* **num_elem_per_idx** – an optional integer value to read multiple contiguous elements per index per partition; default is 1. +- **num_elem_per_idx** – an optional integer value to read multiple contiguous elements per index per partition; default is 1. -* **num_valid_indices** – an optional integer value to specify the number of valid indices per GpSimd core; default is -`index.size / (index.shape[0] / 16)`. +- **num_valid_indices** – an optional integer value to specify the number of valid indices per GpSimd core; default is + `index.size / (index.shape[0] / 16)`. Click [`here`](../../downloads/test_nki_isa_local_gather.py) to download the full NKI code example with equivalent numpy implementation. @@ -544,6 +544,7 @@ full NKI code example with equivalent numpy implementation. **Engine:** GpSimd Engine **Signature:** + ```python isa.memset(dst, value, engine=engine_enum.unknown, name=None) ``` @@ -554,8 +555,8 @@ The memset instruction supports all valid NKI dtypes (see [Supported Data Types] - **dst** — destination tile to initialize. - **value** — the constant value to initialize with - **engine** — specify which engine to use for memset: `nki.isa.engine.vector` or `nki.isa.engine.gpsimd` ; - `nki.isa.engine.unknown` by default, lets compiler select the best engine for the given - input tile shape + `nki.isa.engine.unknown` by default, lets compiler select the best engine for the given + input tile shape > **Note:** > For x4 packed types (`float8_e4m3fn_x4`, `float8_e5m2_x4`, diff --git a/skills/neuron-nki-docs/references/programming/api/api-nki-isa-misc.md b/skills/neuron-nki-docs/references/programming/api/api-nki-isa-misc.md index 3160c42..b74d2c1 100644 --- a/skills/neuron-nki-docs/references/programming/api/api-nki-isa-misc.md +++ b/skills/neuron-nki-docs/references/programming/api/api-nki-isa-misc.md @@ -9,24 +9,22 @@ Other ISA functions. ## Functions - ### nki.isa.dge_mode {#nki-isa-dge_mode} # nki.isa.dge_mode nki.isa.dge_mode -*class *nki.isa.dge_mode(*value*)[[source]](../../../_modules/nki/isa.html#dge_mode) +*class *nki.isa.dge*mode(\_value*)[[source]](../../../\_modules/nki/isa.html#dge_mode) Neuron Descriptor Generation Engine Mode Attributes - | unknown | Unknown DGE mode, i.e., let compiler decide the DGE mode | -| --- | --- | -| swdge | Software DGE | -| hwdge | Hardware DGE | -| none | Not using DGE | +| ------- | -------------------------------------------------------- | +| swdge | Software DGE | +| hwdge | Hardware DGE | +| none | Not using DGE | --- @@ -36,19 +34,18 @@ Attributes nki.isa.engine -*class *nki.isa.engine(*value*)[[source]](../../../_modules/nki/isa.html#engine) +*class *nki.isa.engine(_value_)[[source]](../../../\_modules/nki/isa.html#engine) Neuron Device engines Attributes - -| tensor | Tensor Engine | -| --- | --- | -| vector | Vector Engine | -| scalar | Scalar Engine | -| gpsimd | GpSIMD Engine | -| dma | DMA Engine | -| sync | Sync Engine | +| tensor | Tensor Engine | +| ------- | -------------- | +| vector | Vector Engine | +| scalar | Scalar Engine | +| gpsimd | GpSIMD Engine | +| dma | DMA Engine | +| sync | Sync Engine | | unknown | Unknown Engine | --- @@ -60,6 +57,7 @@ Attributes **Engine:** Tensor Engine **Signature:** + ```python isa.quantize_mx(dst, src, dst_scale, name=None) ``` @@ -69,7 +67,6 @@ Quantize FP16/BF16 data to MXFP8 tensors (both data and scales) using Vector Eng > **Note:** > > Available only on NeuronCore-v4 and newer. -> The resulting `dst` and `dst_scale` tensors use the MXFP8 element and scale data types as defined in the [OCP Microscaling standard](https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf). @@ -119,7 +116,7 @@ For a group :math:`V` of 32 values, let .. math:: - a_{\max} = \max_{V_i \in V} |V_i| +a*{\max} = \max*{V_i \in V} |V_i| Let :math:`E_{\max}` be the maximum unbiased exponent of the destination element data type: 8 for `float8_e4m3fn` and 15 for `float8_e5m2`. @@ -128,11 +125,10 @@ The block scale :math:`X` is calculated as .. math:: - X = - 2^{ - \left\lfloor \log_2(a_{\max}) \right\rfloor - - (E_{\max} - 1) - } +X = +2^{ +\left\lfloor \log*2(a*{\max}) \right\rfloor - (E\_{\max} - 1) +} For an all-zero group, where :math:`a_{\max} = 0`, the block scale is set to :math:`X = 2^{-127}`, the minimum value representable by `float8_e8m0fnu`. @@ -150,6 +146,7 @@ For an all-zero group, where :math:`a_{\max} = 0`, the block scale is set to **Engine:** Vector Engine **Signature:** + ```python isa.rand2(dst, min, max, name=None) ``` @@ -159,7 +156,6 @@ Generate pseudo random numbers with uniform distribution using Vector Engine. > **Note:** > > Available only on NeuronCore-v4 and newer. -> This instruction generates pseudo random numbers and stores them into SBUF/PSUM. The generated values follow a uniform distribution within the specified [min, max] range. @@ -203,6 +199,7 @@ elements per partition of `dst` must not exceed the physical size of each SBUF/P **Engine:** GpSimd Engine **Signature:** + ```python isa.rand_get_state(dst, engine=engine_enum.gpsimd, name=None) ``` @@ -233,7 +230,7 @@ The output `dst` tile must be uint32. - **dst** — the destination tensor to store PRNG state values; must be a 2D uint32 tensor - **engine** — specify which engine to use: `nki.isa.engine.gpsimd` (default) - or `nki.isa.engine.vector` (NeuronCore-v4+) + or `nki.isa.engine.vector` (NeuronCore-v4+) --- @@ -244,6 +241,7 @@ The output `dst` tile must be uint32. **Engine:** GpSimd Engine **Signature:** + ```python isa.rand_set_state(src_seeds, engine=engine_enum.gpsimd, name=None) ``` @@ -276,10 +274,10 @@ The input `src_seeds` tile must be uint32. - `src_seeds` must be in SBUF. - **src_seeds** — the source tensor containing seed values for the PRNG; must be a 2D uint32 tensor - with the partition dimension representing the compute lanes and the free dimension - containing the seed values + with the partition dimension representing the compute lanes and the free dimension + containing the seed values - **engine** — specify which engine to use: `nki.isa.engine.gpsimd` (default) - or `nki.isa.engine.vector` (NeuronCore-v4+) + or `nki.isa.engine.vector` (NeuronCore-v4+) --- @@ -289,18 +287,17 @@ The input `src_seeds` tile must be uint32. nki.isa.reduce_cmd -*class *nki.isa.reduce_cmd(*value*)[[source]](../../../_modules/nki/isa.html#reduce_cmd) +*class *nki.isa.reduce*cmd(\_value*)[[source]](../../../\_modules/nki/isa.html#reduce_cmd) Engine Register Reduce commands Attributes - -| idle | Not using the accumulator registers | -| --- | --- | -| reset | Resets the accumulator registers to its initial state | -| reduce | Keeps accumulating over the current value of the accumulator registers | -| reset_reduce | Resets the accumulator registers then immediately accumulate the results of the current instruction into the accumulators | -| load_reduce | Loads a value into the accumulator registers, then accumulate the results of the current instruction into the accumulators | +| idle | Not using the accumulator registers | +| ------------ | -------------------------------------------------------------------------------------------------------------------------- | +| reset | Resets the accumulator registers to its initial state | +| reduce | Keeps accumulating over the current value of the accumulator registers | +| reset_reduce | Resets the accumulator registers then immediately accumulate the results of the current instruction into the accumulators | +| load_reduce | Loads a value into the accumulator registers, then accumulate the results of the current instruction into the accumulators | --- @@ -311,6 +308,7 @@ Attributes **Engine:** GpSimd Engine **Signature:** + ```python isa.register_alloc(x=None) ``` @@ -362,6 +360,7 @@ nisa.register_load(reg2, zero_tensor) `nki.isa.register_load(dst, src)` **Signature:** + ```python isa.register_load(dst, src) ``` @@ -396,6 +395,7 @@ nisa.register_load(loop_reg, computed_bound) `nki.isa.register_move(dst, src)` **Signature:** + ```python isa.register_move(dst, src) ``` @@ -433,6 +433,7 @@ nisa.register_move(reg2, loop_count) # Copy value from loop_count `nki.isa.register_store(dst, src)` **Signature:** + ```python isa.register_store(dst, src) ``` @@ -469,6 +470,7 @@ nisa.register_store(result_tensor, counter_reg) **Engine:** GpSimd Engine **Signature:** + ```python isa.rng(dst, engine=engine_enum.vector, name=None) ``` @@ -505,11 +507,10 @@ elements per partition of `dst` must not exceed the physical size of each SBUF/P - **dst** — the destination tensor to write random values to - **engine** — specify which engine to use: `nki.isa.engine.vector` (default) - or `nki.isa.engine.gpsimd` (NeuronCore-v3+) + or `nki.isa.engine.gpsimd` (NeuronCore-v3+) --- - ### nki.isa.set_rng_seed {#nki-isa-set_rng_seed} `nki.isa.set_rng_seed(src_seeds, name)` @@ -517,6 +518,7 @@ elements per partition of `dst` must not exceed the physical size of each SBUF/P **Engine:** Vector Engine **Signature:** + ```python isa.set_rng_seed(src_seeds, name=None) ``` @@ -552,6 +554,7 @@ The input `src_seeds` must be a [1,1] tensor. **Engine:** Vector Engine **Signature:** + ```python isa.exponential(dst, src, max_value=0.0, reduce_res=None, reduce_cmd=reduce_cmd_enum.idle, reduce_init=0.0, name=None) ``` @@ -560,7 +563,6 @@ Apply exponential function to each element after subtracting a max_value using V > **Note:** > Available only on NeuronCore-v4 and newer. -> This instruction computes `exp(src - max_value)` for each element. The instruction can optionally maintain a running sum of the exponential values using shared internal reduction @@ -615,10 +617,11 @@ When operands are manually allocated, their base partitions must satisfy: - **max_value** — The maximum value to subtract from each element before applying exponential (for numerical stability). Can be a scalar or vector of shape `(src.shape[0], 1)`. Supported dtypes: float32. - **reduce_res** — Optional tile to store reduction results (sum of exponentials). Must have shape `(src.shape[0], 1)`. Supported buffers: SBUF, PSUM. Supported dtypes: float8_e4m3, float8_e5m2, float16, bfloat16, float32, tfloat32. - Pass `None` to keep the reduction result in the Vector Engine's internal - accumulator without writing it out. This is useful when chaining multiple - calls that reduce into the same accumulator — only the final call needs to - pass a tile to retrieve the accumulated result. + Pass `None` to keep the reduction result in the Vector Engine's internal + accumulator without writing it out. This is useful when chaining multiple + calls that reduce into the same accumulator — only the final call needs to + pass a tile to retrieve the accumulated result. + - **reduce_cmd** — Control the state of reduction registers for accumulating exponential results. Supported: `idle`, `reset_reduce`, `reduce`, `load_reduce`. - **reduce_init** — Initial value for reduction when using `reduce_cmd.load_reduce`. Supported dtypes: float32. @@ -641,7 +644,6 @@ The Vector Engine maintains internal accumulator registers that can be controlle > nki.isa.range_select , nki.isa.select_reduce , > nki.isa.tensor_scalar_reduce , and > nki.isa.tensor_scalar_cumulative . -> **Behavior** @@ -672,6 +674,7 @@ for i in range(num_elements): **Engine:** GpSimd Engine **Signature:** + ```python isa.nonzero_with_count(dst, src, index_offset=0, padding_val=-1, name=None) ``` @@ -681,7 +684,6 @@ Find indices of nonzero elements in an input tensor and their total count using > **Note:** > > Available only on NeuronCore-v3 and newer. -> NOTE: this instruction only operates on partitions [0, 16, 32, ..., 112] of the input tile and writes to partitions [0, 16, 32, ..., 112] of the destination tile. The data in other diff --git a/skills/neuron-nki-docs/references/programming/api/api-nki-isa-scalar.md b/skills/neuron-nki-docs/references/programming/api/api-nki-isa-scalar.md index 9c5561a..8fd1c0b 100644 --- a/skills/neuron-nki-docs/references/programming/api/api-nki-isa-scalar.md +++ b/skills/neuron-nki-docs/references/programming/api/api-nki-isa-scalar.md @@ -16,6 +16,7 @@ Scalar Engine instructions. **Engine:** Scalar Engine **Signature:** + ```python isa.activation(dst, op, data, bias=None, scale=1.0, reduce_op=None, reduce_res=None, reduce_cmd=reduce_cmd_enum.idle, name=None) ``` @@ -129,7 +130,6 @@ In addition, on the Scalar engine a scattered `dst` cannot be in PSUM. > `sin`, `arctan`, `log`, `sqrt`, `rsqrt`, and `reciprocal` > have limited valid input ranges. See [Supported Activation Functions for NKI ISA](nki.api.shared.md#nki-act-func) for their > ranges and out-of-range behavior. -> - **dst** — the activation output - **op** — an activation function (see [Supported Activation Functions for NKI ISA](nki.api.shared.md#nki-act-func) for supported functions). @@ -139,10 +139,11 @@ In addition, on the Scalar engine a scattered `dst` cannot be in PSUM. - **reduce_op** — the reduce operation to perform on the free dimension of the activated data - **reduce_res** — a tile of shape `(data.shape[0], 1)` to hold the final state of `reduce_regs`. - Pass `None` to keep the reduction result in the Scalar Engine's internal - accumulator without writing it out. This is useful when chaining multiple - calls that reduce into the same accumulator — only the final call needs to - pass a tile to retrieve the accumulated result. + Pass `None` to keep the reduction result in the Scalar Engine's internal + accumulator without writing it out. This is useful when chaining multiple + calls that reduce into the same accumulator — only the final call needs to + pass a tile to retrieve the accumulated result. + - **reduce_cmd** — an enum member from `nisa.reduce_cmd` to control the state of `reduce_regs`. --- @@ -154,6 +155,7 @@ In addition, on the Scalar engine a scattered `dst` cannot be in PSUM. **Engine:** Scalar Engine **Signature:** + ```python isa.activate2(dst, op, data, imm0, imm1, op0, op1, relu_param=0.0, reverse0=False, reverse1=False, reduce_op=None, reduce_res=None, reduce_cmd=reduce_cmd_enum.idle, name=None) ``` @@ -163,7 +165,6 @@ using Scalar Engine. > **Note:** > Available only on NeuronCore-v4 and newer. -> This instruction provides a three-stage pipeline per partition: @@ -221,36 +222,36 @@ In addition, on the Scalar engine a scattered `dst` cannot be in PSUM. > `sin`, `arctan`, `log`, `sqrt`, `rsqrt`, and `reciprocal` > have limited valid input ranges. See [Supported Activation Functions for NKI ISA](nki.api.shared.md#nki-act-func) for their > ranges and out-of-range behavior. -> - **dst** — the activation output tile. Supported buffers: SBUF, PSUM. - **op** — an activation function (see [Supported Activation Functions for NKI ISA](nki.api.shared.md#nki-act-func) for supported functions). - **data** — the input tile; layout: (partition axis <= 128, free axis). Supported buffers: SBUF, PSUM. - **imm0** — scalar or `[N, 1]` vector value for the first tensor-scalar operation. - `N` must match the partition dimension size of `data`. + `N` must match the partition dimension size of `data`. - **imm1** — scalar or `[N, 1]` vector value for the second tensor-scalar operation. - `N` must match the partition dimension size of `data`. + `N` must match the partition dimension size of `data`. - **op0** — first ALU operation in tensor-scalar pipeline. Must be an arithmetic operator - (e.g., `nl.multiply`, `nl.add`, `nl.subtract`) or `nl.bypass` for no operation. + (e.g., `nl.multiply`, `nl.add`, `nl.subtract`) or `nl.bypass` for no operation. - **op1** — second ALU operation in tensor-scalar pipeline. Must be an arithmetic operator - (e.g., `nl.add`, `nl.subtract`) or `nl.bypass` for no operation. + (e.g., `nl.add`, `nl.subtract`) or `nl.bypass` for no operation. - **relu_param** — scalar or vector parameter for parameterized activation functions (e.g., PReLU). - Defaults to `0.0`. + Defaults to `0.0`. - **reverse0** — reverse operand order for `op0`. When `True`, computes - `imm0 data` instead of `data imm0`. Requires `op0` to be set. + `imm0 data` instead of `data imm0`. Requires `op0` to be set. - **reverse1** — reverse operand order for `op1`. When `True`, computes - `imm1 result` instead of `result imm1`. Requires `op1` to be set. + `imm1 result` instead of `result imm1`. Requires `op1` to be set. - **reduce_op** — the reduce operation to perform on the free dimension of the activated data. - Supported: `nl.add`, `nl.maximum`, `nl.minimum`, `nl.abs_max`, `nl.abs_min`. + Supported: `nl.add`, `nl.maximum`, `nl.minimum`, `nl.abs_max`, `nl.abs_min`. - **reduce_res** — a tile of shape `(data.shape[0], 1)` to hold the final state of the - reduction registers. Supported buffers: SBUF, PSUM. + reduction registers. Supported buffers: SBUF, PSUM. + + Pass `None` to keep the reduction result in the Scalar Engine's internal + accumulator without writing it out. This is useful when chaining multiple + calls that reduce into the same accumulator — only the final call needs to + pass a tile to retrieve the accumulated result. - Pass `None` to keep the reduction result in the Scalar Engine's internal - accumulator without writing it out. This is useful when chaining multiple - calls that reduce into the same accumulator — only the final call needs to - pass a tile to retrieve the accumulated result. - **reduce_cmd** — an enum member from `nisa.reduce_cmd` to control the state of the - reduction registers. + reduction registers. **Accumulator behavior:** @@ -268,8 +269,7 @@ values into the output tile. > **Note:** > The accumulator registers are shared across Scalar Engine accumulation instructions including -> nki.isa.activation and `nki.isa.activate2`. -> +> nki.isa.activation and `nki.isa.activate2`. **Example** @@ -334,6 +334,7 @@ for i in range(num_elements_per_partition): **Engine:** Scalar Engine **Signature:** + ```python isa.activation_reduce(dst, op, data, reduce_op, reduce_res, bias=None, scale=1.0, name=None) ``` @@ -346,10 +347,10 @@ This API is equivalent to calling `nisa.activation` with `reduce_cmd=nisa.reduce_cmd.reset_reduce` and passing in reduce_res. This API is kept for backward compatibility, we recommend using `nisa.activation` moving forward. -Refer to nisa.activation for semantics of `op/data/bias/scale`. +Refer to nisa.activation for semantics of `op/data/bias/scale`. -In addition to nisa.activation computation, this API also performs a reduction -along the free dimension(s) of the nisa.activation result, at a small additional +In addition to nisa.activation computation, this API also performs a reduction +along the free dimension(s) of the nisa.activation result, at a small additional performance cost. The reduction result is written into `reduce_res`, which must be a SBUF/PSUM tile with the same partition axis size as the input tile `data` and one element per partition. On NeuronCore-v2, the `reduce_op` must be `nl.add`. @@ -401,24 +402,24 @@ In addition, on the Scalar engine a scattered `dst` cannot be in PSUM. > `sin`, `arctan`, `log`, `sqrt`, `rsqrt`, and `reciprocal` > have limited valid input ranges. See [Supported Activation Functions for NKI ISA](nki.api.shared.md#nki-act-func) for their > ranges and out-of-range behavior. -> - **dst** — output tile of the activation instruction; layout: same as input `data` tile - **op** — an activation function (see [Supported Activation Functions for NKI ISA](nki.api.shared.md#nki-act-func) for supported functions). - **data** — the input tile; layout: (partition axis <= 128, free axis) - **reduce_op** — the reduce operation to perform on the free dimension of the activation result - **reduce_res** — a tile of shape `(data.shape[0], 1)`, where data.shape[0] - is the partition axis size of the input `data` tile. The result of `sum(ReductionResult)` - is written into the tensor. + is the partition axis size of the input `data` tile. The result of `sum(ReductionResult)` + is written into the tensor. Pass `None` to keep the reduction result in the Scalar Engine's internal accumulator without writing it out. This is useful when chaining multiple calls that reduce into the same accumulator — only the final call needs to pass a tile to retrieve the accumulated result. + - **bias** — a vector with the same partition axis size as `data` - for broadcast add (after broadcast multiply with `scale`) + for broadcast add (after broadcast multiply with `scale`) - **scale** — a scalar or a vector with the same partition axis size as `data` - for broadcast multiply + for broadcast multiply --- @@ -428,11 +429,12 @@ In addition, on the Scalar engine a scattered `dst` cannot be in PSUM. nki.isa.dropout -nki.isa.dropout(*dst*, *data*, *prob*, *name=None*)[[source]](../../../_modules/nki/isa.html#dropout) +nki.isa.dropout(_dst_, _data_, _prob_, _name=None_)[[source]](../../../\_modules/nki/isa.html#dropout) Randomly replace some elements of the input tile `data` with zeros based on input probabilities using Vector Engine. The probability of replacing input elements with zeros (i.e., drop probability) is specified using the `prob` field: + - If the probability is 1.0, all elements are replaced with zeros. - If the probability is 0.0, all elements are kept with their original values. @@ -445,22 +447,22 @@ Data type of the input `data` tile can be any valid NKI data types (see [Supported Data Types](nki.api.shared.md#nki-dtype) for more information). However, data type of `prob` has restrictions based on the data type of `data`: -* If data type of `data` is any of the integer types (e.g., int32, int16), -`prob` data type must be float32 +- If data type of `data` is any of the integer types (e.g., int32, int16), + `prob` data type must be float32 -* If data type of data is any of the float types (e.g., float32, bfloat16), -`prob` data can be any valid float type +- If data type of data is any of the float types (e.g., float32, bfloat16), + `prob` data can be any valid float type The output data type `dst.dtype` must match the input data type `data.dtype`. Parameters: -* **dst** – an output tile of the dropout result +- **dst** – an output tile of the dropout result -* **data** – the input tile +- **data** – the input tile -* **prob** – a scalar or a tile of shape `(data.shape[0], 1)` to indicate the -probability of replacing elements with zeros +- **prob** – a scalar or a tile of shape `(data.shape[0], 1)` to indicate the + probability of replacing elements with zeros --- @@ -470,7 +472,7 @@ probability of replacing elements with zeros nki.isa.reciprocal -nki.isa.reciprocal(*dst*, *data*, *name=None*)[[source]](../../../_modules/nki/isa.html#reciprocal) +nki.isa.reciprocal(_dst_, _data_, _name=None_)[[source]](../../../\_modules/nki/isa.html#reciprocal) Compute element-wise reciprocal (1.0/x) of the input `data` tile using Vector Engine. **Memory types.** @@ -495,8 +497,8 @@ that of `data` and must not exceed the physical size of each SBUF partition. Parameters: -* **dst** – the output tile +- **dst** – the output tile -* **data** – the input tile +- **data** – the input tile --- diff --git a/skills/neuron-nki-docs/references/programming/api/api-nki-isa-tensor.md b/skills/neuron-nki-docs/references/programming/api/api-nki-isa-tensor.md index ed94991..da2c4ea 100644 --- a/skills/neuron-nki-docs/references/programming/api/api-nki-isa-tensor.md +++ b/skills/neuron-nki-docs/references/programming/api/api-nki-isa-tensor.md @@ -15,7 +15,7 @@ Tensor Engine instructions for matrix operations. nki.isa.get_nc_version -nki.isa.get_nc_version()[[source]](../../../_modules/nki/isa.html#get_nc_version) +nki.isa.get_nc_version()[[source]](../../../\_modules/nki/isa.html#get_nc_version) Returns the `nc_version` of the current target context. --- @@ -26,7 +26,7 @@ Returns the `nc_version` of the current target context. nki.isa.nc_find_index8 -nki.isa.nc_find_index8(*dst*, *data*, *vals*, *name=None*)[[source]](../../../_modules/nki/isa.html#nc_find_index8) +nki.isa.nc*find_index8(\_dst*, _data_, _vals_, _name=None_)[[source]](../../../\_modules/nki/isa.html#nc_find_index8) Find indices of the 8 given vals in each partition of the data tensor. This instruction first loads the 8 values, @@ -46,11 +46,11 @@ If provided, a mask is applied only to the data tensor. Parameters: -* **dst** – a 2D tile containing indices (uint16 or uint32) of the 8 values in each partition with shape [par_dim, 8] +- **dst** – a 2D tile containing indices (uint16 or uint32) of the 8 values in each partition with shape [par_dim, 8] -* **data** – the data tensor to find indices from +- **data** – the data tensor to find indices from -* **vals** – tensor containing the 8 values per partition whose indices will be found +- **vals** – tensor containing the 8 values per partition whose indices will be found --- @@ -60,7 +60,7 @@ Parameters: nki.isa.nc_match_replace8 -nki.isa.nc_match_replace8(*dst*, *data*, *vals*, *imm*, *dst_idx=None*, *name=None*)[[source]](../../../_modules/nki/isa.html#nc_match_replace8) +nki.isa.nc*match_replace8(\_dst*, _data_, _vals_, _imm_, _dst_idx=None_, _name=None_)[[source]](../../../\_modules/nki/isa.html#nc_match_replace8) Replace first occurrence of each value in `vals` with `imm` in `data` using the Vector engine and return the replaced tensor. If `dst_idx` tile is provided, the indices of the matched values are written to `dst_idx`. @@ -87,7 +87,6 @@ If provided, a mask is applied to the data tensor. **NumPy equivalent:** - ```python # Let's assume we work with NumPy, and ``data``, ``vals`` are 2-dimensional arrays # (with shape[0] being the partition axis) and imm is a constant float32 value. @@ -118,18 +117,17 @@ output = data_2d.reshape(data.shape) indices = indices.reshape(vals.shape) # Computed only if ``dst_idx`` is specified ``` - Parameters: -* **dst** – the modified data tensor +- **dst** – the modified data tensor -* **data** – the data tensor to modify +- **data** – the data tensor to modify -* **dst_idx** – (optional) the destination tile to write flattened indices of matched values +- **dst_idx** – (optional) the destination tile to write flattened indices of matched values -* **vals** – tensor containing the 8 values per partition to replace +- **vals** – tensor containing the 8 values per partition to replace -* **imm** – float32 constant to replace matched values with +- **imm** – float32 constant to replace matched values with --- @@ -140,6 +138,7 @@ Parameters: **Engine:** Tensor Engine **Signature:** + ```python isa.nc_matmul(dst, stationary, moving, is_stationary_onezero=False, is_moving_onezero=False, is_transpose=False, accumulate=None, tile_position=(), tile_size=(), perf_mode=matmul_perf_mode.none, name=None) ``` @@ -197,14 +196,13 @@ The `accumulate` parameter controls whether the matmul result overwrites or accu - `accumulate=None` (default): The compiler automatically infers the correct flag — first write overwrites, subsequent writes accumulate. -*Hint*: Use `accumulate=(i > 0)` in a loop to explicitly overwrite on the first iteration +_Hint_: Use `accumulate=(i > 0)` in a loop to explicitly overwrite on the first iteration and accumulate on subsequent iterations, or simply omit the parameter to let the compiler handle it. > **Note:** > On NeuronCore-v2 and NeuronCore-v3, matmul accumulating into a PSUM location that was value-initialized > by a non-matmul instruction (e.g., `memset`, `tensor_copy`) is not supported and is > undefined behavior on hardware. -> **Transpose mode.** @@ -220,7 +218,7 @@ ensures neuron-profile identifies this instruction as a transpose for performanc **Memory types.** -The `nc_matmul` instruction *must* read inputs from SBUF and +The `nc_matmul` instruction _must_ read inputs from SBUF and write outputs to PSUM. Therefore, the `stationary` and `moving` must be SBUF tiles, and `dst` tile must be a PSUM tile. @@ -316,16 +314,16 @@ When operands are manually allocated, their base partitions must satisfy: - **stationary** — the stationary operand - **moving** — the moving operand - **is_stationary_onezero** — hints to the compiler whether the `stationary` operand is a tile with ones/zeros only; - setting this field explicitly could lead to 2x better performance - if `stationary` tile is in float32; the field has no impact for non-float32 `stationary` + setting this field explicitly could lead to 2x better performance + if `stationary` tile is in float32; the field has no impact for non-float32 `stationary` - **is_moving_onezero** — hints to the compiler whether the `moving` operand is a tile with ones/zeros only; - setting this field explicitly could lead to 2x better performance - if `moving` tile is in float32; the field has no impact for non-float32 `moving` + setting this field explicitly could lead to 2x better performance + if `moving` tile is in float32; the field has no impact for non-float32 `moving` - **is_transpose** — controls Tensor Engine transpose mode on/off starting NeuronCore-v3 - **accumulate** — if True, accumulate the matmul result into the existing `dst` PSUM tile content; - if False, overwrite the existing content; - if None (default), auto-detect based on whether this PSUM location was previously written. - Not exposed for `nc_transpose`. + if False, overwrite the existing content; + if None (default), auto-detect based on whether this PSUM location was previously written. + Not exposed for `nc_transpose`. - **tile_position** — a 2D tuple (start_row, start_column) to control starting row in Tensor Engine tiling mode; start_column must be 0 - **tile_size** — a 2D tuple (row_size, column_size) to control row tile size in Tensor Engine tiling mode; column_size must be 128 - **perf_mode** — controls Tensor Engine FP8 double performance mode on/off starting NeuronCore-v3: `matmul_perf_mode.none` (default) disables double FP8 mode; `matmul_perf_mode.double_row` enables double FP8 mode which achieves 2x matmul throughput by packing two FP8 weight/ifmap element pairs and computing two multiplications in parallel per cycle; cannot be combined with column tiling mode. See the [Trainium2 arch guide](../../architecture/trainium2_arch.md) for more information. @@ -372,6 +370,7 @@ def nc_matmul_accumulate_kernel(lhsT, rhs): **Engine:** Tensor Engine **Signature:** + ```python isa.nc_matmul_mx(dst, stationary, moving, stationary_scale, moving_scale, tile_position=None, tile_size=None, accumulate=None, name=None) ``` @@ -381,7 +380,6 @@ Compute matrix multiplication of MXFP8/MXFP4 quantized matrices with integrated > **Note:** > > Available only on NeuronCore-v4 and newer. -> The NeuronCore-v4 Tensor Engine supports matrix multiplication of MXFP8/MXFP4 quantized matrices as defined in the OCP Microscaling standard. @@ -439,7 +437,7 @@ matches the free dimension of the `dst` tile in size. The scale tensors follow a special layout requirement. See more details in `nisa.quantize_mx` API doc. -*Tile size* +_Tile size_ - The partition dimension size of `stationary` and `moving` must be identical and be a multiple of 32, not exceeding 128. @@ -479,20 +477,19 @@ When operands are manually allocated, their base partitions must satisfy: - **stationary** — the stationary quantized matrix (SBUF tile) - **moving** — the moving quantized matrix (SBUF tile) - **stationary_scale** — the dequantization scales for stationary matrix - (SBUF tile) + (SBUF tile) - **moving_scale** — the dequantization scales for moving matrix (SBUF tile) - **tile_position** — a 2D tuple (start_row, start_column) to control - starting row and column in Tensor Engine tiling mode + starting row and column in Tensor Engine tiling mode - **tile_size** — a 2D tuple (row_size, column_size) to control row and - column tile sizes in Tensor Engine tiling mode + column tile sizes in Tensor Engine tiling mode - **accumulate** — if True, accumulate the matmul result into the existing - `dst` PSUM tile content; if False, overwrite the - existing content; if None (default), auto-detect based on - whether this PSUM location was previously written + `dst` PSUM tile content; if False, overwrite the + existing content; if None (default), auto-detect based on + whether this PSUM location was previously written --- - ### nki.isa.nc_n_gather {#nki-isa-nc_n_gather} `nki.isa.nc_n_gather(dst, data, indices, name)` @@ -500,6 +497,7 @@ When operands are manually allocated, their base partitions must satisfy: **Engine:** DMA Engine **Signature:** + ```python isa.nc_n_gather(dst, data, indices, name=None) ``` @@ -519,7 +517,7 @@ flattened indices from the same partition in `indices`. If you need to gather el The `n` in `nc_n_gather` indicates that this instruction corresponds to `n` groups of instructions in the underlying ISA, where `n = ceil(elems_per_partition / 512)`. -Alternatively, we could gather elements by calling nisa.dma_copy with an +Alternatively, we could gather elements by calling nisa.dma_copy with an indirect access pattern derived from `indices`. However, this is less efficient than `nc_n_gather`, which uses GpSimd Engine to perform local data movement within SBUF, without using DMA engines. @@ -552,14 +550,13 @@ The indices' values must be within the range `[0, data.size / data.shape[0])`. --- - ### nki.isa.nc_stream_shuffle {#nki-isa-nc_stream_shuffle} # nki.isa.nc_stream_shuffle nki.isa.nc_stream_shuffle -nki.isa.nc_stream_shuffle(*dst*, *src*, *shuffle_mask*, *name=None*)[[source]](../../../_modules/nki/isa.html#nc_stream_shuffle) +nki.isa.nc*stream_shuffle(\_dst*, _src_, _shuffle_mask_, _name=None_)[[source]](../../../\_modules/nki/isa.html#nc_stream_shuffle) Apply cross-partition data movement within a quadrant of 32 partitions from source tile `src` to destination tile `dst` using Vector Engine. @@ -574,30 +571,30 @@ see [Cross-partition Data Movement](../../architecture/trainium_inferentia2_arch This API has 3 constraints on `src` and `dst`: -* `dst` must have same data type as `src`. +- `dst` must have same data type as `src`. -* `dst` must have the same number of elements per partition as `src`. +- `dst` must have the same number of elements per partition as `src`. -* The access start partition of `src` (`src_start_partition`), does not have to match or be in the same quadrant -as that of `dst` (`dst_start_partition`). However, `src_start_partition`/`dst_start_partition` needs to follow -some special hardware rules with the number of active partitions `num_active_partitions`. -`num_active_partitions = ceil(max(src_num_partitions, dst_num_partitions)/32) * 32`, where `src_num_partitions` and -`dst_num_partitions` refer to the number of partitions the `src` and `dst` tensors access respectively. -`src_start_partition`/`dst_start_partition` is constrained based on the value of `num_active_partitions`: +- The access start partition of `src` (`src_start_partition`), does not have to match or be in the same quadrant + as that of `dst` (`dst_start_partition`). However, `src_start_partition`/`dst_start_partition` needs to follow + some special hardware rules with the number of active partitions `num_active_partitions`. + `num_active_partitions = ceil(max(src_num_partitions, dst_num_partitions)/32) * 32`, where `src_num_partitions` and + `dst_num_partitions` refer to the number of partitions the `src` and `dst` tensors access respectively. + `src_start_partition`/`dst_start_partition` is constrained based on the value of `num_active_partitions`: -* If `num_active_partitions` is 96/128, `src_start_partition`/`dst_start_partition` must be 0. +- If `num_active_partitions` is 96/128, `src_start_partition`/`dst_start_partition` must be 0. -* If `num_active_partitions` is 64, `src_start_partition`/`dst_start_partition` must be 0/64. +- If `num_active_partitions` is 64, `src_start_partition`/`dst_start_partition` must be 0/64. -* If `num_active_partitions` is 32, `src_start_partition`/`dst_start_partition` must be 0/32/64/96. +- If `num_active_partitions` is 32, `src_start_partition`/`dst_start_partition` must be 0/32/64/96. Parameters: -* **dst** – the destination tile +- **dst** – the destination tile -* **src** – the source tile +- **src** – the source tile -* **shuffle_mask** – a 32-element list that specifies the shuffle source and destination partition +- **shuffle_mask** – a 32-element list that specifies the shuffle source and destination partition --- @@ -607,7 +604,7 @@ Parameters: nki.isa.nc_transpose -nki.isa.nc_transpose(*dst*, *data*, *engine=engine.unknown*, *name=None*)[[source]](../../../_modules/nki/isa.html#nc_transpose) +nki.isa.nc*transpose(\_dst*, _data_, _engine=engine.unknown_, _name=None_)[[source]](../../../\_modules/nki/isa.html#nc_transpose) Perform a 2D transpose between the partition axis and the free axis of input `data` using Tensor or Vector Engine. If the `data` tile has more than one free axis, this API implicitly flattens all free axes into one axis @@ -643,12 +640,12 @@ based on the input shape. Parameters: -* **dst** – the transpose output +- **dst** – the transpose output -* **data** – the input tile to be transposed +- **data** – the input tile to be transposed -* **engine** – specify which engine to use for transpose: `nki.isa.tensor_engine` or `nki.isa.vector_engine`; -by default, the best engine will be selected for the given input tile shape +- **engine** – specify which engine to use for transpose: `nki.isa.tensor_engine` or `nki.isa.vector_engine`; + by default, the best engine will be selected for the given input tile shape --- @@ -658,18 +655,17 @@ by default, the best engine will be selected for the given input tile shape nki.isa.nc_version -*class *nki.isa.nc_version(*value*)[[source]](../../../_modules/nki/isa.html#nc_version) +*class *nki.isa.nc*version(\_value*)[[source]](../../../\_modules/nki/isa.html#nc_version) NeuronCore version -__init__() +**init**() Attributes - | gen2 | Trn1/Inf2 target | -| --- | --- | -| gen3 | Trn2 target | -| gen4 | Trn3 target | +| ---- | ---------------- | +| gen3 | Trn2 target | +| gen4 | Trn3 target | --- @@ -679,10 +675,11 @@ Attributes nki.isa.scalar_tensor_tensor -nki.isa.scalar_tensor_tensor(*dst*, *data*, *op0*, *operand0*, *op1*, *operand1*, *reverse0=False*, *reverse1=False*, *name=None*)[[source]](../../../_modules/nki/isa.html#scalar_tensor_tensor) +nki.isa.scalar*tensor_tensor(\_dst*, _data_, _op0_, _operand0_, _op1_, _operand1_, _reverse0=False_, _reverse1=False_, _name=None_)[[source]](../../../\_modules/nki/isa.html#scalar_tensor_tensor) Apply two math operators in sequence using Vector Engine: `(data operand0) operand1`. This instruction is equivalent to running two operations back-to-back: + 1. `temp_result = tensor_scalar(data, op0, operand0)` - broadcast `operand0` and apply `op0` 2. `dst = tensor_tensor(temp_result, op1, operand1)` - element-wise operation with `operand1` @@ -726,24 +723,24 @@ and the number of elements per partition of `operand0` must be 1. Parameters: -* **dst** – the output tile +- **dst** – the output tile -* **data** – the input tile +- **data** – the input tile -* **op0** – the first math operator used with operand0 (see [Supported Math Operators for NKI ISA](nki.api.shared.md#nki-aluop) for supported operators) +- **op0** – the first math operator used with operand0 (see [Supported Math Operators for NKI ISA](nki.api.shared.md#nki-aluop) for supported operators) -* **operand0** – a scalar constant or a tile of shape `(data.shape[0], 1)`, where data.shape[0] -is the partition axis size of the input `data` tile +- **operand0** – a scalar constant or a tile of shape `(data.shape[0], 1)`, where data.shape[0] + is the partition axis size of the input `data` tile -* **reverse0** – reverse ordering of inputs to `op0`; if false, `operand0` is the rhs of `op0`; -if true, `operand0` is the lhs of `op0` +- **reverse0** – reverse ordering of inputs to `op0`; if false, `operand0` is the rhs of `op0`; + if true, `operand0` is the lhs of `op0` -* **op1** – the second math operator used with operand1 (see [Supported Math Operators for NKI ISA](nki.api.shared.md#nki-aluop) for supported operators) +- **op1** – the second math operator used with operand1 (see [Supported Math Operators for NKI ISA](nki.api.shared.md#nki-aluop) for supported operators) -* **operand1** – a tile with the same size as `data` for element-wise operation +- **operand1** – a tile with the same size as `data` for element-wise operation -* **reverse1** – reverse ordering of inputs to `op1`; if false, `operand1` is the rhs of `op1`; -if true, `operand1` is the lhs of `op1` +- **reverse1** – reverse ordering of inputs to `op1`; if false, `operand1` is the rhs of `op1`; + if true, `operand1` is the lhs of `op1` --- @@ -754,6 +751,7 @@ if true, `operand1` is the lhs of `op1` **Engine:** Vector Engine **Signature:** + ```python isa.tensor_copy(dst, src, engine=engine_enum.unknown, name=None) ``` @@ -776,7 +774,7 @@ In addition, since GpSimd Engine cannot access PSUM in NeuronCore, Scalar or Vec output tile is in PSUM (see [NeuronCore-v2 Compute Engines](../../architecture/trainium_inferentia2_arch.md#arch-sec-neuron-core-engines) for details). By default, this API returns a tile in SBUF, unless the returned value is assigned to a pre-declared PSUM tile. -On NeuronCore v2, `tensor_copy` is not supported on the Scalar Engine. Instead, use nisa.activation with `op=nl.copy`. +On NeuronCore v2, `tensor_copy` is not supported on the Scalar Engine. Instead, use nisa.activation with `op=nl.copy`. **Tensor indirection.** @@ -801,18 +799,17 @@ In addition, on the Scalar engine a scattered `dst` cannot be in PSUM. - **dst** — a tile with the same content and partition axis size as the `src` tile. - **src** — the source of copy, must be a tile in SBUF or PSUM. - **engine** — (optional) the engine to use for the operation: `nki.isa.engine.vector`, `nki.isa.engine.scalar`, - `nki.isa.engine.gpsimd` or `nki.isa.engine.unknown` (default, compiler selects best engine based on engine workload). + `nki.isa.engine.gpsimd` or `nki.isa.engine.unknown` (default, compiler selects best engine based on engine workload). --- - ### nki.isa.tensor_copy_dynamic_dst {#nki-isa-tensor_copy_dynamic_dst} # nki.isa.tensor_copy_dynamic_dst nki.isa.tensor_copy_dynamic_dst -nki.isa.tensor_copy_dynamic_dst(*dst*, *src*, *engine=engine.unknown*, *name=None*)[[source]](../../../_modules/nki/isa.html#tensor_copy_dynamic_dst) +nki.isa.tensor*copy_dynamic_dst(\_dst*, _src_, _engine=engine.unknown_, _name=None_)[[source]](../../../\_modules/nki/isa.html#tensor_copy_dynamic_dst) Create a copy of `src` tile within NeuronCore on-chip SRAMs using Vector or Scalar or GpSimd Engine, with `dst` located at a dynamic offset within each partition. @@ -827,12 +824,12 @@ once per offset. Parameters: -* **dst** – the destination of copy, must be a tile in SBUF of PSUM that is dynamically indexed within each dimension. +- **dst** – the destination of copy, must be a tile in SBUF of PSUM that is dynamically indexed within each dimension. -* **src** – the source of copy, must be a tile in SBUF or PSUM. +- **src** – the source of copy, must be a tile in SBUF or PSUM. -* **engine** – (optional) the engine to use for the operation: nki.isa.vector_engine, nki.isa.gpsimd_engine, -nki.isa.scalar_engine or nki.isa.unknown_engine (default, let compiler select best engine). +- **engine** – (optional) the engine to use for the operation: nki.isa.vector_engine, nki.isa.gpsimd_engine, + nki.isa.scalar_engine or nki.isa.unknown_engine (default, let compiler select best engine). --- @@ -842,7 +839,7 @@ nki.isa.scalar_engine or nki.isa.unknown_engine (default, let compiler select be nki.isa.tensor_copy_dynamic_src -nki.isa.tensor_copy_dynamic_src(*dst*, *src*, *engine=engine.unknown*, *name=None*)[[source]](../../../_modules/nki/isa.html#tensor_copy_dynamic_src) +nki.isa.tensor*copy_dynamic_src(\_dst*, _src_, _engine=engine.unknown_, _name=None_)[[source]](../../../\_modules/nki/isa.html#tensor_copy_dynamic_src) Create a copy of `src` tile within NeuronCore on-chip SRAMs using Vector or Scalar or GpSimd Engine, with `src` located at a dynamic offset within each partition. @@ -858,12 +855,12 @@ once per offset. Parameters: -* **src** – the source of copy, must be a tile in SBUF or PSUM that is dynamically indexed within each partition. +- **src** – the source of copy, must be a tile in SBUF or PSUM that is dynamically indexed within each partition. -* **engine** – (optional) the engine to use for the operation: nki.isa.vector_engine, nki.isa.gpsimd_engine, -nki.isa.scalar_engine or nki.isa.unknown_engine (default, let compiler select best engine). +- **engine** – (optional) the engine to use for the operation: nki.isa.vector_engine, nki.isa.gpsimd_engine, + nki.isa.scalar_engine or nki.isa.unknown_engine (default, let compiler select best engine). -* **return** – the modified destination of copy. +- **return** – the modified destination of copy. --- @@ -874,6 +871,7 @@ nki.isa.scalar_engine or nki.isa.unknown_engine (default, let compiler select be **Engine:** Vector Engine **Signature:** + ```python isa.tensor_copy_predicated(dst, src, predicate, reverse_pred=False, name=None) ``` @@ -917,27 +915,27 @@ When operands are manually allocated, their base partitions must satisfy: :param `src`: The source tile or number to copy elements from when `predicate` is True :param `dst`: The destination tile to copy elements to :param `predicate`: A tile that determines which elements to copy + - **reverse_pred** — A boolean that reverses the effect of `predicate`. --- - ### nki.isa.tensor_partition_reduce {#nki-isa-tensor_partition_reduce} # nki.isa.tensor_partition_reduce nki.isa.tensor_partition_reduce -nki.isa.tensor_partition_reduce(*dst*, *op*, *data*, *name=None*)[[source]](../../../_modules/nki/isa.html#tensor_partition_reduce) +nki.isa.tensor*partition_reduce(\_dst*, _op_, _data_, _name=None_)[[source]](../../../\_modules/nki/isa.html#tensor_partition_reduce) Apply a reduction operation across partitions of an input `data` tile using GpSimd Engine. Parameters: -* **dst** – output tile with reduced result +- **dst** – output tile with reduced result -* **op** – the reduction operator (add, max, bitwise_or, bitwise_and) +- **op** – the reduction operator (add, max, bitwise_or, bitwise_and) -* **data** – the input tile to be reduced +- **data** – the input tile to be reduced --- @@ -948,6 +946,7 @@ Parameters: **Engine:** Tensor Engine **Signature:** + ```python isa.tensor_reduce(dst, op, data, axis, negate=False, keepdims=False, name=None) ``` @@ -1027,13 +1026,13 @@ When operands are manually allocated, their base partitions must satisfy: - **op** — the reduction operator (see [Supported Math Operators for NKI ISA](nki.api.shared.md#nki-aluop) for supported reduction operators) - **data** — the input tile to be reduced - **axis** — int or tuple/list of ints. The axis (or axes) along which to reduce; - must be the last contiguous free dimension(s) ending at the final dim. - For example, for a 4D tile `(P, D1, D2, D3)`: valid values are - `(3,)`, `(2, 3)`, or `(1, 2, 3)`. Axis 0 (partition dim) cannot be reduced. + must be the last contiguous free dimension(s) ending at the final dim. + For example, for a 4D tile `(P, D1, D2, D3)`: valid values are + `(3,)`, `(2, 3)`, or `(1, 2, 3)`. Axis 0 (partition dim) cannot be reduced. - **negate** — if True, reduction result is multiplied by `-1.0`; - only applicable when op is an arithmetic operator + only applicable when op is an arithmetic operator - **keepdims** — If this is set to True, the axes which are reduced are left in the result as dimensions with size one. - With this option, the result will broadcast correctly against the input array. + With this option, the result will broadcast correctly against the input array. --- @@ -1044,6 +1043,7 @@ When operands are manually allocated, their base partitions must satisfy: **Engine:** Vector Engine **Signature:** + ```python isa.tensor_scalar(dst, data, op0, operand0, reverse0=False, op1=None, operand1=None, reverse1=False, engine=_engine_enum.unknown, name=None) ``` @@ -1060,8 +1060,8 @@ Note, performing one operator has the same performance cost as performing two op When the operators are non-commutative (e.g., subtract), we can reverse ordering of the inputs for each operator through: - - `reverse0 = True`: `tmp_res = operand0 data` - - `reverse1 = True`: `operand1 tmp_res` +- `reverse0 = True`: `tmp_res = operand0 data` +- `reverse1 = True`: `operand1 tmp_res` The `tensor_scalar` instruction supports two types of operators: 1) bitvec operators (e.g., bitwise_and) and 2) arithmetic operators (e.g., add). @@ -1075,9 +1075,9 @@ If arithmetic operators are used, the `tensor_scalar` instruction can run on Vec However, each engine supports limited arithmetic operators (see tbl-aluop). The Scalar Engine on trn2 only supports some operator combinations: - - `op0=nl.multiply` and `op1=nl.add` - - `op0=nl.multiply` and `op1=None` - - `op0=nl.add` and `op1=None` +- `op0=nl.multiply` and `op1=nl.add` +- `op0=nl.multiply` and `op1=None` +- `op0=nl.add` and `op1=None` Also, arithmetic operators impose no restriction on the data types of input tensor `data` and output tensor `dst`, but the operand0 and operand1 (if used) must be float32. @@ -1112,19 +1112,19 @@ In addition, on the Scalar engine a scattered `dst` cannot be in PSUM. - **data** — the input tile - **op0** — the first math operator used with operand0 (see [Supported Math Operators for NKI ISA](nki.api.shared.md#nki-aluop) for supported operators). - **operand0** — a scalar constant or a tile of shape `(data.shape[0], 1)`, where data.shape[0] - is the partition axis size of the input `data` tile. - Must be `None` or `0` when `op0` is a unary operator (e.g., `nl.abs`). + is the partition axis size of the input `data` tile. + Must be `None` or `0` when `op0` is a unary operator (e.g., `nl.abs`). - **reverse0** — reverse ordering of inputs to `op0`; if false, `operand0` is the rhs of `op0`; - if true, `operand0` is the lhs of `op0` + if true, `operand0` is the lhs of `op0` - **op1** — the second math operator used with operand1 (see [Supported Math Operators for NKI ISA](nki.api.shared.md#nki-aluop) for supported operators); - this operator is optional + this operator is optional - **operand1** — a scalar constant or a tile of shape `(data.shape[0], 1)`, where data.shape[0] - is the partition axis size of the input `data` tile + is the partition axis size of the input `data` tile - **reverse1** — reverse ordering of inputs to `op1`; if false, `operand1` is the rhs of `op1`; - if true, `operand1` is the lhs of `op1` + if true, `operand1` is the lhs of `op1` - **engine** — (optional) the engine to use for the operation: `nki.isa.engine.vector`, `nki.isa.engine.scalar`, - `nki.isa.engine.gpsimd` (only allowed for rsqrt) or `nki.isa.engine.unknown` (default, let - compiler select best engine based on the input tile shape). + `nki.isa.engine.gpsimd` (only allowed for rsqrt) or `nki.isa.engine.unknown` (default, let + compiler select best engine based on the input tile shape). --- @@ -1135,6 +1135,7 @@ In addition, on the Scalar engine a scattered `dst` cannot be in PSUM. **Engine:** Vector Engine **Signature:** + ```python isa.tensor_scalar_cumulative(dst, src, op0, op1, imm0, imm1=None, reduce_cmd=reduce_cmd_enum.reset_reduce, name=None) ``` @@ -1176,7 +1177,7 @@ for i in len(in_tensor): - Scalar operation (`op0`) must be an arithmetic op (e.g., add, mult, max) - Reduction operation (`op1`) is limited to add, subtract, mult, max, min - Input / output dtypes are restricted to BF16, FP16, FP32, FP8, UINT8, UINT16, INT8, INT16 - - INT32/UINT32 are not supported as input/output dtypes (ISA limitation) + - INT32/UINT32 are not supported as input/output dtypes (ISA limitation) **Accumulator behavior:** @@ -1191,7 +1192,6 @@ The Vector Engine maintains internal accumulator registers controlled via `reduc > nki.isa.exponential , nki.isa.range_select , > nki.isa.select_reduce , and > nki.isa.tensor_scalar_reduce . -> **Tensor indirection.** @@ -1217,13 +1217,12 @@ When operands are manually allocated, their base partitions must satisfy: - **op1** — Cumulative arithmetic operation for cumulative computation - **imm0** — Scalar or vector value for tensor-scalar operation. Must be FP32 datatype - **imm1** — (optional) Initial scalar or vector value for the accumulator when `load_reduce` - is specified as the `reduce_cmd`. Must be FP32 datatype + is specified as the `reduce_cmd`. Must be FP32 datatype - **reduce_cmd** — (optional) Control accumulator behavior using `nisa.reduce_cmd` values, - defaults to `reset_reduce` + defaults to `reset_reduce` --- - ### nki.isa.tensor_scalar_reduce {#nki-isa-tensor_scalar_reduce} `nki.isa.tensor_scalar_reduce(dst, data, op0, operand0, reduce_op, reduce_res, reverse0, reduce_cmd, reduce_init, name)` @@ -1231,6 +1230,7 @@ When operands are manually allocated, their base partitions must satisfy: **Engine:** Vector Engine **Signature:** + ```python isa.tensor_scalar_reduce(dst, data, op0, operand0, reduce_op, reduce_res, reverse0=False, reduce_cmd=reduce_cmd_enum.reset_reduce, reduce_init=None, name=None) ``` @@ -1238,13 +1238,13 @@ isa.tensor_scalar_reduce(dst, data, op0, operand0, reduce_op, reduce_res, revers Perform the same computation as `nisa.tensor_scalar` with one math operator and also a reduction along the free dimension of the `nisa.tensor_scalar` result using Vector Engine. -Refer to nisa.tensor_scalar for semantics of `data/op0/operand0`. +Refer to nisa.tensor_scalar for semantics of `data/op0/operand0`. Unlike regular `nisa.tensor_scalar` where two operators are supported, only one operator is supported in this API. Also, `op0` can only be arithmetic operation in nki-aluop. Bitvec operators are not supported in this API. -In addition to nisa.tensor_scalar computation, this API also performs a reduction -along the free dimension(s) of the nisa.tensor_scalar result, at a small additional +In addition to nisa.tensor_scalar computation, this API also performs a reduction +along the free dimension(s) of the nisa.tensor_scalar result, at a small additional performance cost. The reduction result is returned in `reduce_res` in-place, which must be a SBUF/PSUM tile with the same partition axis size as the input tile `data` and one element per partition. The `reduce_op` can be any of `nl.add`, `nl.multiply`, `nl.max` or `nl.min`. @@ -1253,8 +1253,8 @@ Reduction axis is not configurable in this API. If the input tile has multiple f reduce across all of them. .. math:: - result = data operand0 \\ - reduce\_res = reduce\_op(dst, axis=) +result = data operand0 \\ +reduce_res = reduce_op(dst, axis=) **Accumulator behavior:** @@ -1273,7 +1273,6 @@ The Vector Engine maintains internal accumulator registers that can be controlle > nki.isa.exponential , nki.isa.range_select , > nki.isa.select_reduce , and > nki.isa.tensor_scalar_cumulative . -> **Tensor indirection.** @@ -1298,22 +1297,21 @@ When operands are manually allocated, their base partitions must satisfy: - **data** — the input tile - **op0** — the math operator used with operand0 (any arithmetic operator in nki-aluop is allowed). - **operand0** — a scalar constant or a tile of shape `(data.shape[0], 1)`, where data.shape[0] - is the partition axis size of the input `data` tile. - Must be `None` or `0` when `op0` is a unary operator (e.g., `nl.abs`). + is the partition axis size of the input `data` tile. + Must be `None` or `0` when `op0` is a unary operator (e.g., `nl.abs`). - **reverse0** — `(not supported yet)` reverse ordering of inputs to `op0`; if false, `operand0` is the rhs of `op0`; - if true, `operand0` is the lhs of `op0`. `<-- currently not supported yet.` + if true, `operand0` is the lhs of `op0`. `<-- currently not supported yet.` - **reduce_op** — the reduce operation to perform on the free dimension of `data operand0` - **reduce_res** — a tile of shape `(data.shape[0], 1)`, where data.shape[0] - is the partition axis size of the input `data` tile. The result of `reduce_op(data operand0)` - is written in-place into the tile. + is the partition axis size of the input `data` tile. The result of `reduce_op(data operand0)` + is written in-place into the tile. - **reduce_cmd** — Control the state of reduction registers for accumulating reduction results. - Supported: `reset_reduce` (default), `reduce`, `load_reduce`. + Supported: `reset_reduce` (default), `reduce`, `load_reduce`. - **reduce_init** — Initial value for reduction when using `reduce_cmd.load_reduce`. - Must be provided when `reduce_cmd` is `load_reduce`. Supported dtypes: float32. + Must be provided when `reduce_cmd` is `load_reduce`. Supported dtypes: float32. --- - ### nki.isa.tensor_tensor {#nki-isa-tensor_tensor} `nki.isa.tensor_tensor(dst, data1, data2, op, engine, name)` @@ -1321,6 +1319,7 @@ When operands are manually allocated, their base partitions must satisfy: **Engine:** Vector Engine **Signature:** + ```python isa.tensor_tensor(dst, data1, data2, op, engine=_engine_enum.unknown, name=None) ``` @@ -1330,7 +1329,7 @@ The two tiles must have the same partition axis size and the same number of elem The element-wise operator is specified using the `op` field. Valid choices for `op`: -1. Any supported *binary* operator that runs on the Vector Engine. (See [Supported Math Operators for NKI ISA](nki.api.shared.md#nki-aluop) for details.) +1. Any supported _binary_ operator that runs on the Vector Engine. (See [Supported Math Operators for NKI ISA](nki.api.shared.md#nki-aluop) for details.) 2. `nl.power`. (Which runs on the GpSimd engine.) For bitvec operators, the input/output data types must be integer types and Vector Engine treats @@ -1356,7 +1355,7 @@ The three legal cases are: 3. `data1` is in PSUM, while `data2` is in SBUF. Note, if you need broadcasting capability in the free dimension for either input tile, you should consider -using nki.isa.tensor_scalar API instead, +using nki.isa.tensor_scalar API instead, which has better performance than `nki.isa.tensor_tensor` in general. **Tensor indirection.** @@ -1383,7 +1382,7 @@ When operands are manually allocated, their base partitions must satisfy: - **data2** — rhs input operand of the element-wise operation - **op** — a binary math operator (see [Supported Math Operators for NKI ISA](nki.api.shared.md#nki-aluop) for supported operators) - **engine** — (optional) the engine to use for the operation: `nki.isa.engine.vector`, `nki.isa.engine.gpsimd` - or `nki.isa.engine.unknown` (default, let compiler select best engine based on the input tile shape). + or `nki.isa.engine.unknown` (default, let compiler select best engine based on the input tile shape). --- @@ -1394,6 +1393,7 @@ When operands are manually allocated, their base partitions must satisfy: **Engine:** Vector Engine **Signature:** + ```python isa.tensor_tensor_scan(dst, data0, data1, initial, op0, op1, reverse0=False, reverse1=False, name=None) ``` @@ -1444,13 +1444,13 @@ cast to `dst.dtype` at no additional performance cost. - **data0** — lhs input operand of the scan operation - **data1** — rhs input operand of the scan operation - **initial** — starting state of the scan; can be a SBUF/PSUM tile with 1 element/partition or a scalar - compile-time constant + compile-time constant - **op0** — a binary arithmetic math operator (see [Supported Math Operators for NKI ISA](nki.api.shared.md#nki-aluop) for supported operators) - **op1** — a binary arithmetic math operator (see [Supported Math Operators for NKI ISA](nki.api.shared.md#nki-aluop) for supported operators) - **reverse0** — reverse ordering of inputs to `op0`; if false, `data0` is the lhs of `op0`; - if true, `data0` is the rhs of `op0` + if true, `data0` is the rhs of `op0` - **reverse1** — reverse ordering of inputs to `op1`; if false, `data1` is the rhs of `op1`; - if true, `data1` is the lhs of `op1` + if true, `data1` is the lhs of `op1` --- @@ -1461,6 +1461,7 @@ cast to `dst.dtype` at no additional performance cost. **Engine:** GpSimd Engine **Signature:** + ```python isa.topk(val_dst, idx_dst, src, n, name=None) ``` diff --git a/skills/neuron-nki-docs/references/programming/api/api-nki-isa-utility.md b/skills/neuron-nki-docs/references/programming/api/api-nki-isa-utility.md index d2d5a21..0bcb160 100644 --- a/skills/neuron-nki-docs/references/programming/api/api-nki-isa-utility.md +++ b/skills/neuron-nki-docs/references/programming/api/api-nki-isa-utility.md @@ -16,6 +16,7 @@ Utility and helper ISA functions. **Engine:** GpSimd Engine **Signature:** + ```python isa.affine_select(dst, pattern, channel_multiplier, on_true_tile, on_false_value, cmp_op=equal, offset=0, name=None) ``` @@ -102,6 +103,7 @@ must have the same partition dimension size and. **Engine:** GpSimd Engine **Signature:** + ```python isa.iota(dst, pattern, offset=0, channel_multiplier=0, name=None) ``` @@ -169,7 +171,7 @@ The total number of elements in `pattern` must match the number of elements per nki.isa.max8 -nki.isa.max8(*dst*, *src*, *name=None*)[[source]](../../../_modules/nki/isa.html#max8) +nki.isa.max8(_dst_, _src_, _name=None_)[[source]](../../../\_modules/nki/isa.html#max8) Find the 8 largest values in each partition of the source tile. This instruction reads the input elements, converts them to fp32 internally, and outputs @@ -181,15 +183,15 @@ The number of elements read per partition must be between 8 and 16,384 inclusive The output will always contain exactly 8 elements per partition. The source and output must have the same partition dimension size: -* source: [par_dim, …] +- source: [par_dim, …] -* output: [par_dim, 8] +- output: [par_dim, 8] Parameters: -* **dst** – a 2D tile containing the 8 largest values per partition in descending order with shape [par_dim, 8] +- **dst** – a 2D tile containing the 8 largest values per partition in descending order with shape [par_dim, 8] -* **src** – the source tile to find maximum values from +- **src** – the source tile to find maximum values from --- @@ -199,14 +201,13 @@ Parameters: nki.isa.range_select -nki.isa.range_select(*dst*, *on_true_tile*, *comp_op0*, *comp_op1*, *bound0*, *bound1*, *reduce_cmd=reduce_cmd.idle*, *reduce_res=None*, *reduce_op=*, *range_start=0.0*, *on_false_value=0.0*, *name=None*)[[source]](../../../_modules/nki/isa.html#range_select) +nki.isa.range*select(\_dst*, _on_true_tile_, _comp_op0_, _comp_op1_, _bound0_, _bound1_, _reduce_cmd=reduce_cmd.idle_, _reduce_res=None_, _reduce_op=_, _range_start=0.0_, _on_false_value=0.0_, _name=None_)[[source]](../../../\_modules/nki/isa.html#range_select) Select elements from `on_true_tile` based on comparison with bounds using Vector Engine. > **Note** > > Note -> -> +> > Available only on NeuronCore-v3 and newer. For each element in `on_true_tile`, compares its free dimension index + `range_start` against `bound0` and `bound1` @@ -222,41 +223,40 @@ In self-attention, we often have this instruction sequence: `range_select` (Vect When `range_select` outputs a full row of `fill_value`, caution is needed to avoid NaN in the activation instruction that subtracts the output of `range_select` by `reduce_res` (max value): -* If `dst.dtype` and `reduce_res.dtype` are both FP32, we should not hit any NaN issue -since `FP32_MIN - FP32_MIN = 0`. Exponentiation on 0 is stable (1.0 exactly). +- If `dst.dtype` and `reduce_res.dtype` are both FP32, we should not hit any NaN issue + since `FP32_MIN - FP32_MIN = 0`. Exponentiation on 0 is stable (1.0 exactly). -* If `dst.dtype` is FP16/BF16/FP8, the fill_value in the output tile will become `-INF` -since HW performs a downcast from FP32_MIN to a smaller dtype. -In this case, you must make sure `reduce_res.dtype` is FP32 to avoid NaN in `activation`. -NaN can be avoided because `activation` always upcasts input tiles to FP32 to perform math operations: `-INF - FP32_MIN = -INF`. -Exponentiation on `-INF` is stable (0.0 exactly). +- If `dst.dtype` is FP16/BF16/FP8, the fill_value in the output tile will become `-INF` + since HW performs a downcast from FP32_MIN to a smaller dtype. + In this case, you must make sure `reduce_res.dtype` is FP32 to avoid NaN in `activation`. + NaN can be avoided because `activation` always upcasts input tiles to FP32 to perform math operations: `-INF - FP32_MIN = -INF`. + Exponentiation on `-INF` is stable (0.0 exactly). **Constraints:** The comparison operators must be one of: -* nl.equal +- nl.equal -* nl.less +- nl.less -* nl.less_equal +- nl.less_equal -* nl.greater +- nl.greater -* nl.greater_equal +- nl.greater_equal Partition dim sizes must match across `on_true_tile`, `bound0`, and `bound1`: -* `bound0` and `bound1` must have one element per partition +- `bound0` and `bound1` must have one element per partition -* `on_true_tile` must be one of the FP dtypes, and `bound0/bound1` must be FP32 types. +- `on_true_tile` must be one of the FP dtypes, and `bound0/bound1` must be FP32 types. The comparison with `bound0`, `bound1`, and free dimension index is done in FP32. Make sure `range_start` + free dimension index is within 2^24 range. **Numpy equivalent:** - ```python indices = np.zeros_like(on_true_tile, dtype=np.float32) indices[:] = range_start + np.arange(on_true_tile[0].size) @@ -266,30 +266,29 @@ select_out_tile = np.where(mask, on_true_tile, on_false_value) reduce_tile = reduce_op(select_out_tile, axis=1, keepdims=True) ``` - Parameters: -* **dst** – output tile with selected elements +- **dst** – output tile with selected elements -* **on_true_tile** – input tile containing elements to select from +- **on_true_tile** – input tile containing elements to select from -* **on_false_value** – constant value to use when selection condition is False. -Due to HW constraints, this must be FP32_MIN FP32 bit pattern +- **on_false_value** – constant value to use when selection condition is False. + Due to HW constraints, this must be FP32_MIN FP32 bit pattern -* **comp_op0** – first comparison operator +- **comp_op0** – first comparison operator -* **comp_op1** – second comparison operator +- **comp_op1** – second comparison operator -* **bound0** – tile with one element per partition for first comparison +- **bound0** – tile with one element per partition for first comparison -* **bound1** – tile with one element per partition for second comparison +- **bound1** – tile with one element per partition for second comparison -* **reduce_op** – reduction operator to apply on across the selected output. Currently only `nl.maximum` is supported. +- **reduce_op** – reduction operator to apply on across the selected output. Currently only `nl.maximum` is supported. -* **reduce_res** – optional tile to store reduction results. +- **reduce_res** – optional tile to store reduction results. -* **range_start** – starting base offset for index array for the free dimension of `on_true_tile`. -Defaults to 0, and must be a compile-time integer. +- **range_start** – starting base offset for index array for the free dimension of `on_true_tile`. + Defaults to 0, and must be a compile-time integer. --- @@ -299,13 +298,12 @@ Defaults to 0, and must be a compile-time integer. nki.isa.select_reduce -nki.isa.select_reduce(*dst*, *predicate*, *on_true*, *on_false*, *reduce_res=None*, *reduce_cmd=reduce_cmd.idle*, *reduce_op=*, *reverse_pred=False*, *name=None*)[[source]](../../../_modules/nki/isa.html#select_reduce) +nki.isa.select*reduce(\_dst*, _predicate_, _on_true_, _on_false_, _reduce_res=None_, _reduce_cmd=reduce_cmd.idle_, _reduce_op=_, _reverse_pred=False_, _name=None_)[[source]](../../../\_modules/nki/isa.html#select_reduce) Selectively copy elements from either `on_true` or `on_false` to the destination tile based on a `predicate` using Vector Engine, with optional reduction (max). The operation can be expressed in NumPy as: - ```python # Select: predicate = ~predicate if reverse_pred else predicate @@ -315,56 +313,54 @@ result = np.where(predicate, on_true, on_false) reduction_result = np.max(result, axis=1, keepdims=True) ``` - **Memory constraints:** -* Both `on_true` and `predicate` are permitted to be in SBUF +- Both `on_true` and `predicate` are permitted to be in SBUF -* Either `on_true` or `predicate` may be in PSUM, but not both simultaneously +- Either `on_true` or `predicate` may be in PSUM, but not both simultaneously -* The destination `dst` can be in either SBUF or PSUM +- The destination `dst` can be in either SBUF or PSUM **Shape and data type constraints:** -* `on_true`, `dst`, and `predicate` must have identical shapes (same number of partitions and elements per partition) +- `on_true`, `dst`, and `predicate` must have identical shapes (same number of partitions and elements per partition) -* `on_true` can be any supported dtype except `tfloat32`, `int32`, `uint32` +- `on_true` can be any supported dtype except `tfloat32`, `int32`, `uint32` -* `on_false` dtype must be `float32` if `on_false` is a scalar. +- `on_false` dtype must be `float32` if `on_false` is a scalar. -* `on_false` has to be either scalar or vector of shape `(on_true.shape[0], 1)` +- `on_false` has to be either scalar or vector of shape `(on_true.shape[0], 1)` -* `predicate` dtype can be any supported integer type `int8`, `uint8`, `int16`, `uint16` +- `predicate` dtype can be any supported integer type `int8`, `uint8`, `int16`, `uint16` -* `reduce_res` must be a vector of shape `(on_true.shape[0], 1)` +- `reduce_res` must be a vector of shape `(on_true.shape[0], 1)` -* `reduce_res` dtype must of float type +- `reduce_res` dtype must of float type -* `reduce_op` only supports `max` +- `reduce_op` only supports `max` **Behavior:** -* Where predicate is True: The corresponding elements from `on_true` are copied to `dst` +- Where predicate is True: The corresponding elements from `on_true` are copied to `dst` -* Where predicate is False: The corresponding elements from `on_false` are copied to `dst` +- Where predicate is False: The corresponding elements from `on_false` are copied to `dst` -* When reduction is enabled, the max value from each partition of the `result` is computed and stored in `reduce_res` +- When reduction is enabled, the max value from each partition of the `result` is computed and stored in `reduce_res` **Accumulator behavior:** The Vector Engine maintains internal accumulator registers that can be controlled via the `reduce_cmd` parameter: -* `nisa.reduce_cmd.reset_reduce`: Reset accumulators to -inf, then accumulate the current results +- `nisa.reduce_cmd.reset_reduce`: Reset accumulators to -inf, then accumulate the current results -* `nisa.reduce_cmd.reduce`: Continue accumulating without resetting (useful for multi-step reductions) +- `nisa.reduce_cmd.reduce`: Continue accumulating without resetting (useful for multi-step reductions) -* `nisa.reduce_cmd.idle`: No accumulation performed (default) +- `nisa.reduce_cmd.idle`: No accumulation performed (default) > **Note** > > Note -> -> +> > Even when `reduce_cmd` is set to `idle`, the accumulator state may still be modified. > Always use `reset_reduce` after any operations that ran with `idle` mode to ensure > consistent behavior. @@ -372,27 +368,26 @@ The Vector Engine maintains internal accumulator registers that can be controlle > **Note** > > Note -> -> +> > The accumulator registers are shared for other Vector Engine accumulation instructions such [nki.isa.range_select](nki.isa.range_select.md) Parameters: -* **dst** – The destination tile to write the selected values to +- **dst** – The destination tile to write the selected values to -* **predicate** – Tile that determines which value to select (on_true or on_false) +- **predicate** – Tile that determines which value to select (on_true or on_false) -* **on_true** – Tile to select from when predicate is True +- **on_true** – Tile to select from when predicate is True -* **on_false** – Value to use when predicate is False, can be a scalar value or a vector tile of `(on_true.shape[0], 1)` +- **on_false** – Value to use when predicate is False, can be a scalar value or a vector tile of `(on_true.shape[0], 1)` -* **reduce_res** – (optional) Tile to store reduction results, must have shape `(on_true.shape[0], 1)` +- **reduce_res** – (optional) Tile to store reduction results, must have shape `(on_true.shape[0], 1)` -* **reduce_cmd** – (optional) Control accumulator behavior using `nisa.reduce_cmd` values, defaults to idle +- **reduce_cmd** – (optional) Control accumulator behavior using `nisa.reduce_cmd` values, defaults to idle -* **reduce_op** – (optional) Reduction operator to apply (only `nl.maximum` is supported) +- **reduce_op** – (optional) Reduction operator to apply (only `nl.maximum` is supported) -* **reverse_pred** – (optional) Reverse the meaning of the predicate condition, defaults to False +- **reverse_pred** – (optional) Reverse the meaning of the predicate condition, defaults to False --- @@ -402,7 +397,7 @@ Parameters: nki.isa.sequence_bounds -nki.isa.sequence_bounds(*dst*, *segment_ids*, *name=None*)[[source]](../../../_modules/nki/isa.html#sequence_bounds) +nki.isa.sequence*bounds(\_dst*, _segment_ids_, _name=None_)[[source]](../../../\_modules/nki/isa.html#sequence_bounds) Compute the sequence bounds for a given set of segment IDs using GpSIMD Engine. Given a tile of segment IDs, this function identifies where each segment begins and ends. @@ -421,7 +416,6 @@ Both the input tile (`segment_ids`) and output tile (`dst`) must have data type **NumPy equivalent:** - ```python def compute_sequence_bounds(sequence): n = len(sequence) @@ -461,11 +455,10 @@ b = ( ) ``` - Parameters: -* **dst** – tile containing the sequence bounds. +- **dst** – tile containing the sequence bounds. -* **segment_ids** – tile containing the segment IDs. Elements with ID=0 are treated as padding. +- **segment_ids** – tile containing the segment IDs. Elements with ID=0 are treated as padding. --- diff --git a/skills/neuron-nki-docs/references/programming/api/api-nki-isa-vector.md b/skills/neuron-nki-docs/references/programming/api/api-nki-isa-vector.md index 3031e2e..3b99b1a 100644 --- a/skills/neuron-nki-docs/references/programming/api/api-nki-isa-vector.md +++ b/skills/neuron-nki-docs/references/programming/api/api-nki-isa-vector.md @@ -15,7 +15,7 @@ Vector Engine instructions. nki.isa.bn_aggr -nki.isa.bn_aggr(*dst*, *data*, *name=None*)[[source]](../../../_modules/nki/isa.html#bn_aggr) +nki.isa.bn*aggr(\_dst*, _data_, _name=None_)[[source]](../../../\_modules/nki/isa.html#bn_aggr) Aggregate one or multiple `bn_stats` outputs to generate a mean and variance per partition using Vector Engine. @@ -35,9 +35,9 @@ The float32 computation results are cast to `dst.dtype` at no additional perform Parameters: -* **dst** – an output tile with two elements per partition: a mean followed by a variance +- **dst** – an output tile with two elements per partition: a mean followed by a variance -* **data** – an input tile with results of one or more [bn_stats](nki.isa.bn_stats.md) +- **data** – an input tile with results of one or more [bn_stats](nki.isa.bn_stats.md) --- @@ -47,32 +47,32 @@ Parameters: nki.isa.bn_stats -nki.isa.bn_stats(*dst*, *data*, *name=None*)[[source]](../../../_modules/nki/isa.html#bn_stats) +nki.isa.bn*stats(\_dst*, _data_, _name=None_)[[source]](../../../\_modules/nki/isa.html#bn_stats) Compute mean- and variance-related statistics for each partition of an input tile `data` in parallel using Vector Engine. The output tile of the instruction has 6 elements per partition: -* the `count` of the even elements (of the input tile elements from the same partition) +- the `count` of the even elements (of the input tile elements from the same partition) -* the `mean` of the even elements +- the `mean` of the even elements -* `variance * count` of the even elements +- `variance * count` of the even elements -* the `count` of the odd elements +- the `count` of the odd elements -* the `mean` of the odd elements +- the `mean` of the odd elements -* `variance * count` of the odd elements +- `variance * count` of the odd elements To get the final mean and variance of the input tile, we need to pass the above `bn_stats` instruction output into the [bn_aggr](nki.isa.bn_aggr.md) instruction, which will output two elements per partition: -* mean (of the original input tile elements from the same partition) +- mean (of the original input tile elements from the same partition) -* variance +- variance Due to hardware limitation, the number of elements per partition (i.e., free dimension size) of the input `data` must not exceed 512 (nl.tile_size.bn_stats_fmax). @@ -87,8 +87,8 @@ The float32 computation results are cast to `dst.dtype` at no additional perform Parameters: -* **dst** – an output tile with 6-element statistics per partition +- **dst** – an output tile with 6-element statistics per partition -* **data** – the input tile (up to 512 elements per partition) +- **data** – the input tile (up to 512 elements per partition) --- diff --git a/skills/neuron-nki-docs/references/programming/api/api-nki-language-creation.md b/skills/neuron-nki-docs/references/programming/api/api-nki-language-creation.md index 9ca86a7..a1f9745 100644 --- a/skills/neuron-nki-docs/references/programming/api/api-nki-language-creation.md +++ b/skills/neuron-nki-docs/references/programming/api/api-nki-language-creation.md @@ -14,6 +14,7 @@ Functions for creating and initializing arrays and tensors. `nki.language.ndarray(shape, dtype, buffer, name, address)` **Signature:** + ```python language.ndarray(shape, dtype, buffer=sbuf, name='', address=None) ``` @@ -25,7 +26,7 @@ Create a new tensor of given shape and dtype on the specified buffer. - **buffer** — the specific buffer (ie, sbuf, psum, hbm), defaults to sbuf. - **name** — the name of the tensor, used in scheduling. - **address** — optional memory address `(partition_offset, free_offset)`. -**Returns:** a new `NkiTensor` allocated on the buffer. + **Returns:** a new `NkiTensor` allocated on the buffer. --- @@ -34,6 +35,7 @@ Create a new tensor of given shape and dtype on the specified buffer. `nki.language.zeros(shape, dtype, buffer, name)` **Signature:** + ```python language.zeros(shape, dtype, buffer=sbuf, name='') ``` @@ -45,13 +47,12 @@ Create a new tensor of given shape and dtype on the specified buffer, filled wit > **Warning:** > > This API is experimental and may change in future releases. -> - **shape** — the shape of the tensor. - **dtype** — the data type of the tensor. - **buffer** — the specific buffer (ie, sbuf, psum, hbm), defaults to sbuf. - **name** — the name of the tensor, used in scheduling. -**Returns:** a new `NkiTensor` allocated on the buffer. + **Returns:** a new `NkiTensor` allocated on the buffer. --- @@ -60,6 +61,7 @@ Create a new tensor of given shape and dtype on the specified buffer, filled wit `nki.language.empty_like(x, dtype, buffer, name)` **Signature:** + ```python language.empty_like(x, dtype=None, buffer=None, name='') ``` @@ -71,13 +73,12 @@ Create a new tensor with the same shape and type as a given tensor. > **Warning:** > > This API is experimental and may change in future releases. -> - **x** — the tensor. - **dtype** — the data type of the tensor (default: same as `x`). - **buffer** — the specific buffer (ie, sbuf, psum, hbm), (default: same as `x`). - **name** — the name of the tensor, used in scheduling. -**Returns:** a new `NkiTensor` with the same shape and type as `x`. + **Returns:** a new `NkiTensor` with the same shape and type as `x`. --- @@ -86,6 +87,7 @@ Create a new tensor with the same shape and type as a given tensor. `nki.language.gather_flattened(data, indices, axis, dtype)` **Signature:** + ```python language.gather_flattened(data, indices, axis=0, dtype=None) ``` @@ -100,13 +102,12 @@ to select from the free dimension of data. > **Warning:** > > This API is experimental and may change in future releases. -> - **data** — input tensor to gather from. - **indices** — indices to gather. - **axis** — axis along which to gather. - **dtype** — (optional) data type to cast the output type to (see [Supported Data Types](nki.api.shared.md#nki-dtype) for more information); if not specified, it will default to be the same as the data type of the input tile. -**Returns:** gathered tensor. + **Returns:** gathered tensor. Examples: @@ -127,6 +128,7 @@ nl.store(actual_tensor[0:128, 0:512], result) `nki.language.load(src, dtype)` **Signature:** + ```python language.load(src, dtype=None) ``` @@ -136,11 +138,10 @@ Load a tensor from device memory (HBM) into on-chip memory (SBUF). > **Warning:** > > This API is experimental and may change in future releases. -> - **src** — HBM tensor to load the data from. - **dtype** — (optional) data type to cast the output type to (see [Supported Data Types](nki.api.shared.md#nki-dtype) for more information); if not specified, it will default to be the same as the data type of the input tile. -**Returns:** a new tile on SBUF with values from `src`. + **Returns:** a new tile on SBUF with values from `src`. --- @@ -149,6 +150,7 @@ Load a tensor from device memory (HBM) into on-chip memory (SBUF). `nki.language.load_transpose2d(src, dtype)` **Signature:** + ```python language.load_transpose2d(src, dtype=None) ``` @@ -158,11 +160,10 @@ Load a tensor from device memory (HBM) and 2D-transpose the data before storing > **Warning:** > > This API is experimental and may change in future releases. -> - **src** — HBM tensor to load the data from. - **dtype** — (optional) data type to cast the output type to (see [Supported Data Types](nki.api.shared.md#nki-dtype) for more information); if not specified, it will default to be the same as the data type of the input tile. -**Returns:** a new tile on SBUF with values from `src` 2D-transposed. + **Returns:** a new tile on SBUF with values from `src` 2D-transposed. --- @@ -171,6 +172,7 @@ Load a tensor from device memory (HBM) and 2D-transpose the data before storing `nki.language.ones(shape, dtype, buffer, name)` **Signature:** + ```python language.ones(shape, dtype, buffer=sbuf, name='') ``` @@ -182,13 +184,12 @@ Create a new tensor of given shape and dtype on the specified buffer, filled wit > **Warning:** > > This API is experimental and may change in future releases. -> - **shape** — the shape of the tensor. - **dtype** — the data type of the tensor. - **buffer** — the specific buffer (ie, sbuf, psum, hbm), defaults to sbuf. - **name** — the name of the tensor, used in scheduling. -**Returns:** a new `NkiTensor` allocated on the buffer. + **Returns:** a new `NkiTensor` allocated on the buffer. --- @@ -197,6 +198,7 @@ Create a new tensor of given shape and dtype on the specified buffer, filled wit `nki.language.rand(shape, dtype, buffer, name)` **Signature:** + ```python language.rand(shape, dtype, buffer=sbuf, name='') ``` @@ -208,13 +210,12 @@ Values are sampled from a uniform distribution between 0 and 1. > **Warning:** > > This API is experimental and may change in future releases. -> - **shape** — the shape of the tensor. - **dtype** — the data type of the tensor (see [Supported Data Types](nki.api.shared.md#nki-dtype) for more information). - **buffer** — the specific buffer (ie, sbuf, psum, hbm), defaults to sbuf. - **name** — the name of the tensor, used in scheduling. -**Returns:** a new `NkiTensor` allocated on the buffer with random values. + **Returns:** a new `NkiTensor` allocated on the buffer with random values. Examples: @@ -232,6 +233,7 @@ a = nl.rand((128, 512), dtype=nl.float32) `nki.language.shared_constant(constant)` **Signature:** + ```python language.shared_constant(constant) ``` @@ -250,8 +252,8 @@ and tfloat32 are supported at the MLIR level but not yet tested end-to-end on hardware. - **constant** — the constant data. Can be a numpy array or a file path - to a `.npy` file. -**Returns:** an NkiTensor in shared_hbm containing the constant data. + to a `.npy` file. + **Returns:** an NkiTensor in shared_hbm containing the constant data. --- @@ -260,6 +262,7 @@ end-to-end on hardware. `nki.language.store(dst, value)` **Signature:** + ```python language.store(dst, value) ``` @@ -269,7 +272,6 @@ Store into a tensor on device memory (HBM) from on-chip memory (SBUF). > **Warning:** > > This API is experimental and may change in future releases. -> - **dst** — HBM tensor to store the data into. - **value** — an SBUF tile that contains the values to store. diff --git a/skills/neuron-nki-docs/references/programming/api/api-nki-language-dims.md b/skills/neuron-nki-docs/references/programming/api/api-nki-language-dims.md index 928b0ce..955f6ce 100644 --- a/skills/neuron-nki-docs/references/programming/api/api-nki-language-dims.md +++ b/skills/neuron-nki-docs/references/programming/api/api-nki-language-dims.md @@ -14,6 +14,7 @@ Dimension and range management functions. `nki.language.affine_range(start, stop, step)` **Signature:** + ```python language.affine_range(start, stop=None, step=1) ``` @@ -27,12 +28,11 @@ function. Prefer using `range()` directly instead. > **Warning:** > > This API is deprecated and will be removed in future releases. -> - **start** — start value (or stop if `stop` is None). - **stop** — stop value (exclusive). - **step** — step size. -**Returns:** an iterator yielding integer values from start to stop. + **Returns:** an iterator yielding integer values from start to stop. Examples: @@ -54,6 +54,7 @@ for i in nl.affine_range(input_tensor.shape[1] // 512): `nki.language.sequential_range(start, stop, step)` **Signature:** + ```python language.sequential_range(start, stop=None, step=1) ``` @@ -67,12 +68,11 @@ function. Prefer using `range()` directly instead. > **Warning:** > > This API is deprecated and will be removed in future releases. -> - **start** — start value (or stop if `stop` is None). - **stop** — stop value (exclusive). - **step** — step size. -**Returns:** an iterator yielding integer values from start to stop. + **Returns:** an iterator yielding integer values from start to stop. Examples: @@ -94,6 +94,7 @@ for i in nl.sequential_range(input_tensor.shape[1] // 512): `nki.language.static_range(start, stop, step)` **Signature:** + ```python language.static_range(start, stop=None, step=1) ``` @@ -107,12 +108,11 @@ function. Prefer using `range()` directly instead. > **Warning:** > > This API is deprecated and will be removed in future releases. -> - **start** — start value (or stop if `stop` is None). - **stop** — stop value (exclusive). - **step** — step size. -**Returns:** an iterator yielding integer values from start to stop. + **Returns:** an iterator yielding integer values from start to stop. Examples: @@ -135,7 +135,7 @@ for i in nl.static_range(input_tensor.shape[1] // 512): nki.language.num_programs -nki.language.num_programs(*axes=None*)[[source]](../../../_modules/nki/language.html#num_programs) +nki.language.num*programs(\_axes=None*)[[source]](../../../\_modules/nki/language.html#num_programs) Number of SPMD programs along the given axes in the launch grid. If `axes` is not provided, returns the total number of programs. @@ -153,7 +153,7 @@ The number of SPMD(single process multiple data) programs along `axes` in the la nki.language.program_id -nki.language.program_id(*axis*)[[source]](../../../_modules/nki/language.html#program_id) +nki.language.program*id(\_axis*)[[source]](../../../\_modules/nki/language.html#program_id) Index of the current SPMD program along the given axis in the launch grid. Parameters: @@ -170,7 +170,7 @@ The program id along `axis` in the launch grid nki.language.program_ndim -nki.language.program_ndim()[[source]](../../../_modules/nki/language.html#program_ndim) +nki.language.program_ndim()[[source]](../../../\_modules/nki/language.html#program_ndim) Number of dimensions in the SPMD launch grid. Returns: @@ -184,26 +184,25 @@ The number of dimensions in the launch grid, i.e. the number of axes nki.language.tile_size -*class *nki.language.tile_size[[source]](../../../_modules/nki/language.html#tile_size) +*class *nki.language.tile_size[[source]](../../../\_modules/nki/language.html#tile_size) Tile size constants. Attributes - -| bn_stats_fmax | Maximum free dimension of BN_STATS | -| --- | --- | -| gemm_moving_fmax | Maximum free dimension of the moving operand of General Matrix Multiplication on Tensor Engine | -| gemm_stationary_fmax | Maximum free dimension of the stationary operand of General Matrix Multiplication on Tensor Engine | -| pmax | Maximum partition dimension of a tile | -| psum_fmax | Maximum free dimension of a tile on PSUM buffer, in FP32 elements | -| psum_fmax_bytes | Maximum free dimension of a tile on PSUM buffer, in bytes | -| psum_num_banks | Number of usable PSUM banks per partition | -| sbuf_size_bytes | Total SBUF capacity in bytes (all partitions combined) | -| sbuf_fmax | Maximum free dimension of a tile on SBUF buffer, in FP32 elements | -| sbuf_fmax_bytes | Maximum free dimension of a tile on SBUF buffer, in bytes | -| psum_min_align | Minimum byte alignment requirement for PSUM free dimension address | -| sbuf_min_align | Minimum byte alignment requirement for SBUF free dimension address | -| total_available_sbuf_size | **Deprecated.** Use `sbuf_fmax_bytes` (per-partition) or `sbuf_size_bytes` (total) | +| bn_stats_fmax | Maximum free dimension of BN_STATS | +| ------------------------- | -------------------------------------------------------------------------------------------------- | +| gemm_moving_fmax | Maximum free dimension of the moving operand of General Matrix Multiplication on Tensor Engine | +| gemm_stationary_fmax | Maximum free dimension of the stationary operand of General Matrix Multiplication on Tensor Engine | +| pmax | Maximum partition dimension of a tile | +| psum_fmax | Maximum free dimension of a tile on PSUM buffer, in FP32 elements | +| psum_fmax_bytes | Maximum free dimension of a tile on PSUM buffer, in bytes | +| psum_num_banks | Number of usable PSUM banks per partition | +| sbuf_size_bytes | Total SBUF capacity in bytes (all partitions combined) | +| sbuf_fmax | Maximum free dimension of a tile on SBUF buffer, in FP32 elements | +| sbuf_fmax_bytes | Maximum free dimension of a tile on SBUF buffer, in bytes | +| psum_min_align | Minimum byte alignment requirement for PSUM free dimension address | +| sbuf_min_align | Minimum byte alignment requirement for SBUF free dimension address | +| total_available_sbuf_size | **Deprecated.** Use `sbuf_fmax_bytes` (per-partition) or `sbuf_size_bytes` (total) | --- @@ -212,6 +211,7 @@ Attributes `nki.language.dynamic_range(start, stop, step)` **Signature:** + ```python language.dynamic_range(start, stop=None, step=1) ``` @@ -224,7 +224,7 @@ The loop runs on device with dynamic bounds. - **start** — start value (or stop if `stop` is None), can be VirtualRegister. - **stop** — stop value (exclusive), can be VirtualRegister. - **step** — step size, must be a compile-time positive integer (not VirtualRegister). -**Returns:** an iterator yielding integer values from start to stop. + **Returns:** an iterator yielding integer values from start to stop. Examples: @@ -245,6 +245,7 @@ for _ in nl.dynamic_range(1): `nki.language.fori_loop(lower, upper, body_fun, step)` **Signature:** + ```python language.fori_loop(lower, upper, body_fun, step=1) ``` @@ -272,9 +273,9 @@ through SBUF/HBM. - **lower** — start value (int or VirtualRegister). - **upper** — end value (exclusive) (int or VirtualRegister). - **body_fun** — function `(i: VirtualRegister) -> None` called each - iteration with the current iteration value. + iteration with the current iteration value. - **step** — step size, must be a compile-time positive integer. -**Returns:** None. Side effects in `body_fun` persist after the loop. + **Returns:** None. Side effects in `body_fun` persist after the loop. Examples: @@ -313,6 +314,7 @@ nl.fori_loop(0, ub_reg, body) `nki.language.while_loop(init, body_fun)` **Signature:** + ```python language.while_loop(init, body_fun) ``` @@ -340,9 +342,9 @@ through SBUF/HBM. - **init** — Initial condition register (VirtualRegister). - **body_fun** — function `(r: VirtualRegister) -> VirtualRegister` called - each iteration with the current condition value; returns the - next condition register. -**Returns:** None. Side effects in `body_fun` persist after the loop. + each iteration with the current condition value; returns the + next condition register. + **Returns:** None. Side effects in `body_fun` persist after the loop. Examples: diff --git a/skills/neuron-nki-docs/references/programming/api/api-nki-language-memory.md b/skills/neuron-nki-docs/references/programming/api/api-nki-language-memory.md index 408766a..99d8136 100644 --- a/skills/neuron-nki-docs/references/programming/api/api-nki-language-memory.md +++ b/skills/neuron-nki-docs/references/programming/api/api-nki-language-memory.md @@ -14,6 +14,7 @@ Memory management and data movement functions for loading/storing tensors. `nki.language.hbm()` **Signature:** + ```python language.hbm ``` @@ -27,6 +28,7 @@ HBM - Alias of private_hbm `nki.language.private_hbm()` **Signature:** + ```python language.private_hbm ``` @@ -40,6 +42,7 @@ HBM - Only visible to each individual kernel instance in the SPMD grid `nki.language.psum()` **Signature:** + ```python language.psum ``` @@ -53,6 +56,7 @@ PSUM - Only visible to each individual kernel instance in the SPMD grid `nki.language.sbuf()` **Signature:** + ```python language.sbuf ``` @@ -66,6 +70,7 @@ State Buffer - Only visible to each individual kernel instance in the SPMD grid `nki.language.shared_hbm()` **Signature:** + ```python language.shared_hbm ``` diff --git a/skills/neuron-nki-docs/references/programming/api/api-nki-language-misc.md b/skills/neuron-nki-docs/references/programming/api/api-nki-language-misc.md index ba2b32f..e80a43a 100644 --- a/skills/neuron-nki-docs/references/programming/api/api-nki-language-misc.md +++ b/skills/neuron-nki-docs/references/programming/api/api-nki-language-misc.md @@ -15,7 +15,7 @@ Other language functions. nki.language.device_print -nki.language.device_print(*print_prefix*, *tensor*)[[source]](../../../_modules/nki/language.html#device_print) +nki.language.device*print(\_print_prefix*, _tensor_)[[source]](../../../\_modules/nki/language.html#device_print) Print a message with a string `print_prefix` followed by the value of a tile `tensor`. By default, using this function will not result in your tensors being printed out. When running your kernel, @@ -38,19 +38,17 @@ def my_nki_kernel(input_tensor): ... ``` - > **Note** > > Warning -> -> +> > This feature is only available when using the NxD Inference library. Parameters: -* **print_prefix** ([*str*](https://docs.python.org/3/library/stdtypes.html#str)) – prefix of the print message. This string is evaluated at trace time and must be a constant expression. +- **print_prefix** ([_str_](https://docs.python.org/3/library/stdtypes.html#str)) – prefix of the print message. This string is evaluated at trace time and must be a constant expression. -* **tensor** – tensor to print out. Can be in SBUF or HBM. +- **tensor** – tensor to print out. Can be in SBUF or HBM. Returns: None @@ -63,10 +61,9 @@ None nki.language.ds -nki.language.ds(*start*, *size*)[[source]](../../../_modules/nki/language.html#ds) +nki.language.ds(_start_, _size_)[[source]](../../../\_modules/nki/language.html#ds) Construct a dynamic slice for simple tensor indexing. - ```python import nki.language as nl import nki.isa as nisa @@ -93,6 +90,7 @@ def example_kernel(in_tensor): `nki.language.all(x, axis, dtype)` **Signature:** + ```python language.all(x, axis, dtype=None) ``` @@ -104,14 +102,13 @@ Whether all elements along the specified axis (or axes) evaluate to True. > **Warning:** > > This API is experimental and may change in future releases. -> - **x** — a tile. - **axis** — int or tuple/list of ints. The axis (or axes) along which to operate; - must be free dimensions, not partition dimension (0); can only be the - last contiguous dim(s) of the tile: [1], [1,2], [1,2,3], [1,2,3,4]. + must be free dimensions, not partition dimension (0); can only be the + last contiguous dim(s) of the tile: [1], [1,2], [1,2,3], [1,2,3,4]. - **dtype** — (optional) data type to cast the output type to (see [Supported Data Types](nki.api.shared.md#nki-dtype) for more information); if not specified, it will default to be the same as the data type of the input tile. -**Returns:** a tile with the logical AND reduction along the provided axis. + **Returns:** a tile with the logical AND reduction along the provided axis. --- @@ -120,6 +117,7 @@ Whether all elements along the specified axis (or axes) evaluate to True. `nki.language.broadcast_to(x, shape, dtype)` **Signature:** + ```python language.broadcast_to(x, shape, dtype=None) ``` @@ -131,16 +129,15 @@ Broadcast a tile to a new shape following numpy broadcasting rules. > **Warning:** > > This API is experimental and may change in future releases. -> If `x.shape` is already the same as `shape`, returns `x` unchanged (or a dtype-cast copy if `dtype` differs). - **x** — the source tile in SBUF or PSUM. - **shape** — the target shape. Must have the same rank as `x`. - Each dimension must either match or be broadcast from size 1. + Each dimension must either match or be broadcast from size 1. - **dtype** — (optional) data type to cast the output type to (see [Supported Data Types](nki.api.shared.md#nki-dtype) for more information); if not specified, it will default to be the same as the data type of the input tile. -**Returns:** a tile with the target shape containing broadcast values from `x`. + **Returns:** a tile with the target shape containing broadcast values from `x`. --- @@ -149,6 +146,7 @@ If `x.shape` is already the same as `shape`, returns `x` unchanged `nki.language.dropout(x, rate, dtype)` **Signature:** + ```python language.dropout(x, rate, dtype=None) ``` @@ -158,13 +156,12 @@ Randomly zeroes some of the elements of the input tile given a probability rate. > **Warning:** > > This API is experimental and may change in future releases. -> - **x** — a tile. - **rate** — the probability of zeroing each element. Can be a scalar constant - or a tile of shape `(x.shape[0], 1)` for per-partition drop probabilities. + or a tile of shape `(x.shape[0], 1)` for per-partition drop probabilities. - **dtype** — (optional) data type to cast the output type to (see [Supported Data Types](nki.api.shared.md#nki-dtype) for more information); if not specified, it will default to be the same as the data type of the input tile. -**Returns:** a tile with randomly zeroed elements of `x`. + **Returns:** a tile with randomly zeroed elements of `x`. --- @@ -173,6 +170,7 @@ Randomly zeroes some of the elements of the input tile given a probability rate. `nki.language.expand_dims(x, axis)` **Signature:** + ```python language.expand_dims(x, axis) ``` @@ -184,13 +182,12 @@ Expand the shape of a tile. > **Warning:** > > This API is experimental and may change in future releases. -> Insert a new axis that will appear at the axis position in the expanded tile shape. - **x** — a tile. - **axis** — position in the expanded axes where the new axis is placed. -**Returns:** a tile with view of input data with the number of dimensions increased. + **Returns:** a tile with view of input data with the number of dimensions increased. --- @@ -201,6 +198,7 @@ Insert a new axis that will appear at the axis position in the expanded tile sha **Engine:** Tensor Engine **Signature:** + ```python language.matmul(x, y, transpose_x=False) ``` @@ -210,15 +208,14 @@ x @ y matrix multiplication of x and y. > **Warning:** > > This API is experimental and may change in future releases. -> - **x** — a tile on SBUF (partition dimension <= 128, free dimension <= 128), - x's free dimension must match y's partition dimension. + x's free dimension must match y's partition dimension. - **y** — a tile on SBUF (partition dimension <= 128, free dimension <= 512). - **transpose_x** — defaults to False. If True, x is treated as already transposed. - If False, an additional transpose will be inserted to make x's partition - dimension the contract dimension of the matmul to align with the Tensor Engine. -**Returns:** x @ y or x.T @ y if transpose_x=True. + If False, an additional transpose will be inserted to make x's partition + dimension the contract dimension of the matmul to align with the Tensor Engine. + **Returns:** x @ y or x.T @ y if transpose_x=True. Examples: @@ -244,6 +241,7 @@ assert nl.equal(result, expected) `nki.language.max(x, axis, dtype, keepdims)` **Signature:** + ```python language.max(x, axis, dtype=None, keepdims=False) ``` @@ -255,15 +253,14 @@ Maximum of elements along the specified axis (or axes) of the input. > **Warning:** > > This API is experimental and may change in future releases. -> - **x** — a tile. - **axis** — int or tuple/list of ints. The axis (or axes) along which to operate; - must be free dimensions, not partition dimension (0); can only be the - last contiguous dim(s) of the tile: [1], [1,2], [1,2,3], [1,2,3,4]. + must be free dimensions, not partition dimension (0); can only be the + last contiguous dim(s) of the tile: [1], [1,2], [1,2,3], [1,2,3,4]. - **dtype** — (optional) data type to cast the output type to (see [Supported Data Types](nki.api.shared.md#nki-dtype) for more information); if not specified, it will default to be the same as the data type of the input tile. - **keepdims** — if True, the reduced axes are kept as size-one dimensions. -**Returns:** a tile with the maximum along the provided axis. + **Returns:** a tile with the maximum along the provided axis. --- @@ -272,6 +269,7 @@ Maximum of elements along the specified axis (or axes) of the input. `nki.language.mean(x, axis, dtype, keepdims)` **Signature:** + ```python language.mean(x, axis, dtype=None, keepdims=False) ``` @@ -283,16 +281,15 @@ Arithmetic mean along the specified axis (or axes) of the input. > **Warning:** > > This API is experimental and may change in future releases. -> - **x** — a tile. - **axis** — int or tuple/list of ints. The axis (or axes) along which to operate; - must be free dimensions, not partition dimension (0); can only be the - last contiguous dim(s) of the tile: [1], [1,2], [1,2,3], [1,2,3,4]. + must be free dimensions, not partition dimension (0); can only be the + last contiguous dim(s) of the tile: [1], [1,2], [1,2,3], [1,2,3,4]. - **dtype** — (optional) data type to cast the output type to (see [Supported Data Types](nki.api.shared.md#nki-dtype) for more information); if not specified, it will default to be the same as the data type of the input tile. - **keepdims** — if True, the reduced axes are kept as size-one dimensions. -**Returns:** a tile with the average of elements along the provided axis. Float32 - intermediate values are used for the computation. + **Returns:** a tile with the average of elements along the provided axis. Float32 + intermediate values are used for the computation. --- @@ -301,6 +298,7 @@ Arithmetic mean along the specified axis (or axes) of the input. `nki.language.min(x, axis, dtype, keepdims)` **Signature:** + ```python language.min(x, axis, dtype=None, keepdims=False) ``` @@ -312,15 +310,14 @@ Minimum of elements along the specified axis (or axes) of the input. > **Warning:** > > This API is experimental and may change in future releases. -> - **x** — a tile. - **axis** — int or tuple/list of ints. The axis (or axes) along which to operate; - must be free dimensions, not partition dimension (0); can only be the - last contiguous dim(s) of the tile: [1], [1,2], [1,2,3], [1,2,3,4]. + must be free dimensions, not partition dimension (0); can only be the + last contiguous dim(s) of the tile: [1], [1,2], [1,2,3], [1,2,3,4]. - **dtype** — (optional) data type to cast the output type to (see [Supported Data Types](nki.api.shared.md#nki-dtype) for more information); if not specified, it will default to be the same as the data type of the input tile. - **keepdims** — if True, the reduced axes are kept as size-one dimensions. -**Returns:** a tile with the minimum along the provided axis. + **Returns:** a tile with the minimum along the provided axis. --- @@ -329,6 +326,7 @@ Minimum of elements along the specified axis (or axes) of the input. `nki.language.prod(x, axis, dtype, keepdims)` **Signature:** + ```python language.prod(x, axis, dtype=None, keepdims=False) ``` @@ -340,15 +338,14 @@ Product of elements along the specified axis (or axes) of the input. > **Warning:** > > This API is experimental and may change in future releases. -> - **x** — a tile. - **axis** — int or tuple/list of ints. The axis (or axes) along which to operate; - must be free dimensions, not partition dimension (0); can only be the - last contiguous dim(s) of the tile: [1], [1,2], [1,2,3], [1,2,3,4]. + must be free dimensions, not partition dimension (0); can only be the + last contiguous dim(s) of the tile: [1], [1,2], [1,2,3], [1,2,3,4]. - **dtype** — (optional) data type to cast the output type to (see [Supported Data Types](nki.api.shared.md#nki-dtype) for more information); if not specified, it will default to be the same as the data type of the input tile. - **keepdims** — if True, the reduced axes are kept as size-one dimensions. -**Returns:** a tile with the product along the provided axis. + **Returns:** a tile with the product along the provided axis. --- @@ -357,6 +354,7 @@ Product of elements along the specified axis (or axes) of the input. `nki.language.rms_norm(x, w, axis, n, epsilon, dtype, compute_dtype)` **Signature:** + ```python language.rms_norm(x, w, axis, n, epsilon=1e-06, dtype=None, compute_dtype=None) ``` @@ -366,7 +364,6 @@ Apply Root Mean Square Layer Normalization. > **Warning:** > > This API is experimental and may change in future releases. -> - **x** — input tile. - **w** — weight tile. @@ -375,7 +372,7 @@ Apply Root Mean Square Layer Normalization. - **epsilon** — epsilon value used by rms calculation to avoid divide-by-zero. - **dtype** — (optional) data type to cast the output type to (see [Supported Data Types](nki.api.shared.md#nki-dtype) for more information); if not specified, it will default to be the same as the data type of the input tile. - **compute_dtype** — (optional) dtype for the internal computation. -**Returns:** `x / RMS(x) * w` + **Returns:** `x / RMS(x) * w` Examples: @@ -395,6 +392,7 @@ result = nl.rms_norm(x, w, axis=1, n=512) `nki.language.softmax(x, axis, dtype)` **Signature:** + ```python language.softmax(x, axis=-1, dtype=None) ``` @@ -406,12 +404,11 @@ Softmax activation function on the input, element-wise. > **Warning:** > > This API is experimental and may change in future releases. -> - **x** — a tile. - **axis** — int or tuple/list of ints. The axis (or axes) along which to operate; must be free dimensions, not partition dimension (0); can only be the last contiguous dim(s) of the tile: `[1], [1,2], [1,2,3], [1,2,3,4]` - **dtype** — (optional) data type to cast the output type to (see [Supported Data Types](nki.api.shared.md#nki-dtype) for more information); if not specified, it will default to be the same as the data type of the input tile. -**Returns:** a tile that has softmax of `x`. + **Returns:** a tile that has softmax of `x`. Examples: @@ -430,6 +427,7 @@ result = nl.softmax(a, axis=1) `nki.language.sum(x, axis, dtype, keepdims)` **Signature:** + ```python language.sum(x, axis, dtype=None, keepdims=False) ``` @@ -441,15 +439,14 @@ Sum of elements along the specified axis (or axes) of the input. > **Warning:** > > This API is experimental and may change in future releases. -> - **x** — a tile. - **axis** — int or tuple/list of ints. The axis (or axes) along which to operate; - must be free dimensions, not partition dimension (0); can only be the - last contiguous dim(s) of the tile: [1], [1,2], [1,2,3], [1,2,3,4]. + must be free dimensions, not partition dimension (0); can only be the + last contiguous dim(s) of the tile: [1], [1,2], [1,2,3], [1,2,3,4]. - **dtype** — (optional) data type to cast the output type to (see [Supported Data Types](nki.api.shared.md#nki-dtype) for more information); if not specified, it will default to be the same as the data type of the input tile. - **keepdims** — if True, the reduced axes are kept as size-one dimensions. -**Returns:** a tile with the sum along the provided axis. + **Returns:** a tile with the sum along the provided axis. --- @@ -458,6 +455,7 @@ Sum of elements along the specified axis (or axes) of the input. `nki.language.transpose(x, dtype)` **Signature:** + ```python language.transpose(x, dtype=None) ``` @@ -467,12 +465,11 @@ Transposes a 2D tile between its partition and free dimension. > **Warning:** > > This API is experimental and may change in future releases. -> - **x** — 2D input tile. - **dtype** — (optional) data type to cast the output type to (see [Supported Data Types](nki.api.shared.md#nki-dtype) for more information); if not specified, it will default to be the same as the data type of the input tile. -**Returns:** a tile that has the values of the input tile with its partition and free - dimensions swapped. + **Returns:** a tile that has the values of the input tile with its partition and free + dimensions swapped. Examples: @@ -495,6 +492,7 @@ assert nl.equal(result, x) `nki.language.var(x, axis, dtype, keepdims)` **Signature:** + ```python language.var(x, axis, dtype=None, keepdims=False) ``` @@ -506,15 +504,14 @@ Variance along the specified axis (or axes) of the input. > **Warning:** > > This API is experimental and may change in future releases. -> - **x** — a tile. - **axis** — int or tuple/list of ints. The axis (or axes) along which to operate; - must be free dimensions, not partition dimension (0); can only be the - last contiguous dim(s) of the tile: [1], [1,2], [1,2,3], [1,2,3,4]. + must be free dimensions, not partition dimension (0); can only be the + last contiguous dim(s) of the tile: [1], [1,2], [1,2,3], [1,2,3,4]. - **dtype** — (optional) data type to cast the output type to (see [Supported Data Types](nki.api.shared.md#nki-dtype) for more information); if not specified, it will default to be the same as the data type of the input tile. - **keepdims** — currently ignored; result always has keepdims=True shape. -**Returns:** a tile with the variance of the elements along the provided axis. + **Returns:** a tile with the variance of the elements along the provided axis. --- @@ -523,6 +520,7 @@ Variance along the specified axis (or axes) of the input. `nki.language.where(condition, x, y, dtype)` **Signature:** + ```python language.where(condition, x, y, dtype=None) ``` @@ -534,13 +532,12 @@ Return elements chosen from x or y depending on condition. > **Warning:** > > This API is experimental and may change in future releases. -> - **condition** — condition tile with float values (1.0 for True, 0.0 for False). - **x** — tensor from which to take elements where condition is True. - **y** — tensor from which to take elements where condition is False. - **dtype** — (optional) data type to cast the output type to (see [Supported Data Types](nki.api.shared.md#nki-dtype) for more information); if not specified, it will default to be the same as the data type of the input tile. -**Returns:** tensor with elements from x or y based on condition. + **Returns:** tensor with elements from x or y based on condition. Examples: @@ -579,6 +576,7 @@ assert nl.equal(result, expected) `nki.language.zeros_like(x, dtype, buffer, name)` **Signature:** + ```python language.zeros_like(x, dtype=None, buffer=None, name='') ``` @@ -590,12 +588,11 @@ Create a new tensor of zeros with the same shape and type as a given tensor. > **Warning:** > > This API is experimental and may change in future releases. -> - **x** — the tensor. - **dtype** — the data type of the tensor. - **buffer** — the specific buffer (ie, sbuf, psum, hbm), defaults to sbuf. - **name** — the name of the tensor, used in scheduling. -**Returns:** a new `NkiTensor` of zeros with the same shape as `x`. + **Returns:** a new `NkiTensor` of zeros with the same shape as `x`. --- diff --git a/skills/neuron-nki-docs/references/programming/api/api-nki-language-operators.md b/skills/neuron-nki-docs/references/programming/api/api-nki-language-operators.md index 0098f32..f5ef8e9 100644 --- a/skills/neuron-nki-docs/references/programming/api/api-nki-language-operators.md +++ b/skills/neuron-nki-docs/references/programming/api/api-nki-language-operators.md @@ -14,6 +14,7 @@ Operator specifiers passed as the `op` argument to `nisa.activation()`, `nisa.te `nki.language.abs()` **Signature:** + ```python language.abs ``` @@ -27,6 +28,7 @@ Op specifier for `abs` operation. Pass as `op` argument to `nisa.activation()`, `nki.language.abs_max()` **Signature:** + ```python language.abs_max ``` @@ -40,6 +42,7 @@ Op specifier for `abs_max` operation. Pass as `op` argument to `nisa.activation( `nki.language.abs_min()` **Signature:** + ```python language.abs_min ``` @@ -53,6 +56,7 @@ Op specifier for `abs_min` operation. Pass as `op` argument to `nisa.activation( `nki.language.add()` **Signature:** + ```python language.add ``` @@ -66,6 +70,7 @@ Op specifier for `add` operation. Pass as `op` argument to `nisa.activation()`, `nki.language.arctan()` **Signature:** + ```python language.arctan ``` @@ -79,6 +84,7 @@ Op specifier for `arctan` operation. Pass as `op` argument to `nisa.activation() `nki.language.average()` **Signature:** + ```python language.average ``` @@ -92,6 +98,7 @@ Op specifier for `average` operation. Pass as `op` argument to `nisa.activation( `nki.language.bitwise_and()` **Signature:** + ```python language.bitwise_and ``` @@ -105,6 +112,7 @@ Op specifier for `bitwise_and` operation. Pass as `op` argument to `nisa.activat `nki.language.bitwise_or()` **Signature:** + ```python language.bitwise_or ``` @@ -118,6 +126,7 @@ Op specifier for `bitwise_or` operation. Pass as `op` argument to `nisa.activati `nki.language.bitwise_xor()` **Signature:** + ```python language.bitwise_xor ``` @@ -131,6 +140,7 @@ Op specifier for `bitwise_xor` operation. Pass as `op` argument to `nisa.activat `nki.language.bypass()` **Signature:** + ```python language.bypass ``` @@ -144,6 +154,7 @@ Op specifier for `bypass` operation. Pass as `op` argument to `nisa.activation() `nki.language.ceil()` **Signature:** + ```python language.ceil ``` @@ -157,6 +168,7 @@ Op specifier for `ceil` operation. Pass as `op` argument to `nisa.activation()`, `nki.language.copy()` **Signature:** + ```python language.copy ``` @@ -170,6 +182,7 @@ Op specifier for `copy` operation. Pass as `op` argument to `nisa.activation()`, `nki.language.cos()` **Signature:** + ```python language.cos ``` @@ -183,6 +196,7 @@ Op specifier for `cos` operation. Pass as `op` argument to `nisa.activation()`, `nki.language.divide()` **Signature:** + ```python language.divide ``` @@ -196,6 +210,7 @@ Op specifier for `divide` operation. Pass as `op` argument to `nisa.activation() `nki.language.equal()` **Signature:** + ```python language.equal ``` @@ -209,6 +224,7 @@ Op specifier for `equal` operation. Pass as `op` argument to `nisa.activation()` `nki.language.erf()` **Signature:** + ```python language.erf ``` @@ -222,6 +238,7 @@ Op specifier for `erf` operation. Pass as `op` argument to `nisa.activation()`, `nki.language.erf_dx()` **Signature:** + ```python language.erf_dx ``` @@ -235,6 +252,7 @@ Op specifier for `erf_dx` operation. Pass as `op` argument to `nisa.activation() `nki.language.exp()` **Signature:** + ```python language.exp ``` @@ -248,6 +266,7 @@ Op specifier for `exp` operation. Pass as `op` argument to `nisa.activation()`, `nki.language.floor()` **Signature:** + ```python language.floor ``` @@ -261,6 +280,7 @@ Op specifier for `floor` operation. Pass as `op` argument to `nisa.activation()` `nki.language.fmod()` **Signature:** + ```python language.fmod ``` @@ -274,6 +294,7 @@ Op specifier for `fmod` operation. Pass as `op` argument to `nisa.activation()`, `nki.language.gelu()` **Signature:** + ```python language.gelu ``` @@ -287,6 +308,7 @@ Op specifier for `gelu` operation. Pass as `op` argument to `nisa.activation()`, `nki.language.gelu_apprx_sigmoid()` **Signature:** + ```python language.gelu_apprx_sigmoid ``` @@ -300,6 +322,7 @@ Op specifier for `gelu_apprx_sigmoid` operation. Pass as `op` argument to `nisa. `nki.language.gelu_apprx_sigmoid_dx()` **Signature:** + ```python language.gelu_apprx_sigmoid_dx ``` @@ -313,6 +336,7 @@ Op specifier for `gelu_apprx_sigmoid_dx` operation. Pass as `op` argument to `ni `nki.language.gelu_apprx_tanh()` **Signature:** + ```python language.gelu_apprx_tanh ``` @@ -326,6 +350,7 @@ Op specifier for `gelu_apprx_tanh` operation. Pass as `op` argument to `nisa.act `nki.language.gelu_dx()` **Signature:** + ```python language.gelu_dx ``` @@ -339,6 +364,7 @@ Op specifier for `gelu_dx` operation. Pass as `op` argument to `nisa.activation( `nki.language.greater()` **Signature:** + ```python language.greater ``` @@ -352,6 +378,7 @@ Op specifier for `greater` operation. Pass as `op` argument to `nisa.activation( `nki.language.greater_equal()` **Signature:** + ```python language.greater_equal ``` @@ -365,6 +392,7 @@ Op specifier for `greater_equal` operation. Pass as `op` argument to `nisa.activ `nki.language.invert()` **Signature:** + ```python language.invert ``` @@ -378,6 +406,7 @@ Op specifier for `invert` operation. Pass as `op` argument to `nisa.activation() `nki.language.left_shift()` **Signature:** + ```python language.left_shift ``` @@ -391,6 +420,7 @@ Op specifier for `left_shift` operation. Pass as `op` argument to `nisa.activati `nki.language.less()` **Signature:** + ```python language.less ``` @@ -404,6 +434,7 @@ Op specifier for `less` operation. Pass as `op` argument to `nisa.activation()`, `nki.language.less_equal()` **Signature:** + ```python language.less_equal ``` @@ -417,6 +448,7 @@ Op specifier for `less_equal` operation. Pass as `op` argument to `nisa.activati `nki.language.log()` **Signature:** + ```python language.log ``` @@ -430,6 +462,7 @@ Op specifier for `log` operation. Pass as `op` argument to `nisa.activation()`, `nki.language.logical_and()` **Signature:** + ```python language.logical_and ``` @@ -443,6 +476,7 @@ Op specifier for `logical_and` operation. Pass as `op` argument to `nisa.activat `nki.language.logical_not()` **Signature:** + ```python language.logical_not ``` @@ -456,6 +490,7 @@ Op specifier for `logical_not` operation. Pass as `op` argument to `nisa.activat `nki.language.logical_or()` **Signature:** + ```python language.logical_or ``` @@ -469,6 +504,7 @@ Op specifier for `logical_or` operation. Pass as `op` argument to `nisa.activati `nki.language.logical_xor()` **Signature:** + ```python language.logical_xor ``` @@ -482,6 +518,7 @@ Op specifier for `logical_xor` operation. Pass as `op` argument to `nisa.activat `nki.language.maximum()` **Signature:** + ```python language.maximum ``` @@ -495,6 +532,7 @@ Op specifier for `maximum` operation. Pass as `op` argument to `nisa.activation( `nki.language.minimum()` **Signature:** + ```python language.minimum ``` @@ -508,6 +546,7 @@ Op specifier for `minimum` operation. Pass as `op` argument to `nisa.activation( `nki.language.mish()` **Signature:** + ```python language.mish ``` @@ -521,6 +560,7 @@ Op specifier for `mish` operation. Pass as `op` argument to `nisa.activation()`, `nki.language.mod()` **Signature:** + ```python language.mod ``` @@ -534,6 +574,7 @@ Op specifier for `mod` operation. Pass as `op` argument to `nisa.activation()`, `nki.language.multiply()` **Signature:** + ```python language.multiply ``` @@ -547,6 +588,7 @@ Op specifier for `multiply` operation. Pass as `op` argument to `nisa.activation `nki.language.negative()` **Signature:** + ```python language.negative ``` @@ -560,6 +602,7 @@ Op specifier for `negative` operation. Pass as `op` argument to `nisa.activation `nki.language.not_equal()` **Signature:** + ```python language.not_equal ``` @@ -573,6 +616,7 @@ Op specifier for `not_equal` operation. Pass as `op` argument to `nisa.activatio `nki.language.power()` **Signature:** + ```python language.power ``` @@ -586,6 +630,7 @@ Op specifier for `power` operation. Pass as `op` argument to `nisa.activation()` `nki.language.prelu()` **Signature:** + ```python language.prelu ``` @@ -599,6 +644,7 @@ Op specifier for `prelu` operation. Pass as `op` argument to `nisa.activation()` `nki.language.reciprocal()` **Signature:** + ```python language.reciprocal ``` @@ -612,6 +658,7 @@ Op specifier for `reciprocal` operation. Pass as `op` argument to `nisa.activati `nki.language.relu()` **Signature:** + ```python language.relu ``` @@ -625,6 +672,7 @@ Op specifier for `relu` operation. Pass as `op` argument to `nisa.activation()`, `nki.language.right_shift()` **Signature:** + ```python language.right_shift ``` @@ -638,6 +686,7 @@ Op specifier for `right_shift` operation. Pass as `op` argument to `nisa.activat `nki.language.rsqrt()` **Signature:** + ```python language.rsqrt ``` @@ -651,6 +700,7 @@ Op specifier for `rsqrt` operation. Pass as `op` argument to `nisa.activation()` `nki.language.sigmoid()` **Signature:** + ```python language.sigmoid ``` @@ -664,6 +714,7 @@ Op specifier for `sigmoid` operation. Pass as `op` argument to `nisa.activation( `nki.language.sign()` **Signature:** + ```python language.sign ``` @@ -677,6 +728,7 @@ Op specifier for `sign` operation. Pass as `op` argument to `nisa.activation()`, `nki.language.silu()` **Signature:** + ```python language.silu ``` @@ -690,6 +742,7 @@ Op specifier for `silu` operation. Pass as `op` argument to `nisa.activation()`, `nki.language.silu_dx()` **Signature:** + ```python language.silu_dx ``` @@ -703,6 +756,7 @@ Op specifier for `silu_dx` operation. Pass as `op` argument to `nisa.activation( `nki.language.sin()` **Signature:** + ```python language.sin ``` @@ -716,6 +770,7 @@ Op specifier for `sin` operation. Pass as `op` argument to `nisa.activation()`, `nki.language.softplus()` **Signature:** + ```python language.softplus ``` @@ -729,6 +784,7 @@ Op specifier for `softplus` operation. Pass as `op` argument to `nisa.activation `nki.language.sqrt()` **Signature:** + ```python language.sqrt ``` @@ -742,6 +798,7 @@ Op specifier for `sqrt` operation. Pass as `op` argument to `nisa.activation()`, `nki.language.square()` **Signature:** + ```python language.square ``` @@ -755,6 +812,7 @@ Op specifier for `square` operation. Pass as `op` argument to `nisa.activation() `nki.language.subtract()` **Signature:** + ```python language.subtract ``` @@ -768,6 +826,7 @@ Op specifier for `subtract` operation. Pass as `op` argument to `nisa.activation `nki.language.tan()` **Signature:** + ```python language.tan ``` @@ -781,6 +840,7 @@ Op specifier for `tan` operation. Pass as `op` argument to `nisa.activation()`, `nki.language.tanh()` **Signature:** + ```python language.tanh ``` @@ -794,6 +854,7 @@ Op specifier for `tanh` operation. Pass as `op` argument to `nisa.activation()`, `nki.language.trunc()` **Signature:** + ```python language.trunc ``` diff --git a/skills/neuron-nki-docs/references/programming/api/api-nki-language-types.md b/skills/neuron-nki-docs/references/programming/api/api-nki-language-types.md index 55ee552..4cb0acf 100644 --- a/skills/neuron-nki-docs/references/programming/api/api-nki-language-types.md +++ b/skills/neuron-nki-docs/references/programming/api/api-nki-language-types.md @@ -14,6 +14,7 @@ Data type conversion and management. `nki.language.bfloat16()` **Signature:** + ```python language.bfloat16 ``` @@ -22,11 +23,12 @@ language.bfloat16 --- -### nki.language.bool_ {#nki-language-bool_} +### nki.language.bool* {#nki-language-bool*} `nki.language.bool_()` **Signature:** + ```python language.bool_ ``` @@ -40,6 +42,7 @@ Boolean (True or False) stored as a byte `nki.language.float16()` **Signature:** + ```python language.float16 ``` @@ -53,6 +56,7 @@ language.float16 `nki.language.float32()` **Signature:** + ```python language.float32 ``` @@ -66,6 +70,7 @@ language.float32 `nki.language.float4_e2m1fn_x4()` **Signature:** + ```python language.float4_e2m1fn_x4 ``` @@ -79,6 +84,7 @@ language.float4_e2m1fn_x4 `nki.language.float8_e4m3()` **Signature:** + ```python language.float8_e4m3 ``` @@ -92,6 +98,7 @@ language.float8_e4m3 `nki.language.float8_e4m3fn()` **Signature:** + ```python language.float8_e4m3fn ``` @@ -105,6 +112,7 @@ language.float8_e4m3fn `nki.language.float8_e4m3fn_x4()` **Signature:** + ```python language.float8_e4m3fn_x4 ``` @@ -118,6 +126,7 @@ language.float8_e4m3fn_x4 `nki.language.float8_e5m2()` **Signature:** + ```python language.float8_e5m2 ``` @@ -131,6 +140,7 @@ language.float8_e5m2 `nki.language.float8_e5m2_x4()` **Signature:** + ```python language.float8_e5m2_x4 ``` @@ -144,6 +154,7 @@ language.float8_e5m2_x4 `nki.language.int16()` **Signature:** + ```python language.int16 ``` @@ -157,6 +168,7 @@ language.int16 `nki.language.int32()` **Signature:** + ```python language.int32 ``` @@ -170,6 +182,7 @@ language.int32 `nki.language.int8()` **Signature:** + ```python language.int8 ``` @@ -183,6 +196,7 @@ language.int8 `nki.language.tfloat32()` **Signature:** + ```python language.tfloat32 ``` @@ -196,6 +210,7 @@ language.tfloat32 `nki.language.uint16()` **Signature:** + ```python language.uint16 ``` @@ -209,6 +224,7 @@ language.uint16 `nki.language.uint32()` **Signature:** + ```python language.uint32 ``` @@ -222,6 +238,7 @@ language.uint32 `nki.language.uint8()` **Signature:** + ```python language.uint8 ``` diff --git a/skills/neuron-nki-docs/references/programming/api/api-nki-tensor.md b/skills/neuron-nki-docs/references/programming/api/api-nki-tensor.md index a12ac96..a19c987 100644 --- a/skills/neuron-nki-docs/references/programming/api/api-nki-tensor.md +++ b/skills/neuron-nki-docs/references/programming/api/api-nki-tensor.md @@ -14,6 +14,7 @@ View methods on `NkiTensor`. Each returns a new `NkiTensor` view that shares the `NkiTensor.ap(pattern, offset, scalar_offset, vector_offset, indirect_dim, dtype)` **Signature:** + ```python NkiTensor.ap(pattern, offset=None, scalar_offset=None, vector_offset=None, indirect_dim=0, dtype=None) ``` @@ -37,15 +38,15 @@ sb.ap(pattern=[[64, 128], [1, 16]], dtype=nl.bfloat16, - **pattern** — list of `[stride, count]` pairs defining the access pattern - **offset** — element offset added to the view's base storage offset. - When `None` (default), inherits the current view's storage offset - unchanged. Pass an explicit integer to compose with the base offset - (e.g. offset=0 keeps the base, offset=N shifts by N additional elements). + When `None` (default), inherits the current view's storage offset + unchanged. Pass an explicit integer to compose with the base offset + (e.g. offset=0 keeps the base, offset=N shifts by N additional elements). - **scalar_offset** — dynamic scalar index tensor for indirect access - **vector_offset** — per-partition index tensor for indirect access - **indirect_dim** — dimension in `self.shape` whose stride scales - the indirect scalar/vector offset (default 0) + the indirect scalar/vector offset (default 0) - **dtype** — reinterpret storage as this dtype (default: tensor's dtype) -**Returns:** new `NkiTensor` with the explicit access pattern + **Returns:** new `NkiTensor` with the explicit access pattern --- @@ -54,6 +55,7 @@ sb.ap(pattern=[[64, 128], [1, 16]], dtype=nl.bfloat16, `NkiTensor.broadcast(dim, size)` **Signature:** + ```python NkiTensor.broadcast(dim, size) ``` @@ -75,7 +77,7 @@ t.broadcast(1, 8) # shape becomes (128, 8, 64) - **dim** — dimension to broadcast - **size** — new size for the dimension -**Returns:** new `NkiTensor` view with the broadcasted dimension + **Returns:** new `NkiTensor` view with the broadcasted dimension --- @@ -84,6 +86,7 @@ t.broadcast(1, 8) # shape becomes (128, 8, 64) `NkiTensor.expand_dim(dim)` **Signature:** + ```python NkiTensor.expand_dim(dim) ``` @@ -103,7 +106,7 @@ t.expand_dim(1) # shape becomes (128, 1, 64) the indirect partition dim; see `is_indirect`) - **dim** — position at which to insert the new dimension -**Returns:** new `NkiTensor` view with an additional size-1 dimension + **Returns:** new `NkiTensor` view with an additional size-1 dimension --- @@ -112,6 +115,7 @@ t.expand_dim(1) # shape becomes (128, 1, 64) `NkiTensor.flatten_dims(start_dim, end_dim)` **Signature:** + ```python NkiTensor.flatten_dims(start_dim, end_dim) ``` @@ -137,7 +141,7 @@ t.flatten_dims(1, 2) # shape becomes (128, 6, 4) - **start_dim** — first dimension to merge (inclusive) - **end_dim** — last dimension to merge (inclusive) -**Returns:** new `NkiTensor` view with the merged dimension + **Returns:** new `NkiTensor` view with the merged dimension --- @@ -146,6 +150,7 @@ t.flatten_dims(1, 2) # shape becomes (128, 6, 4) `NkiTensor.get_pattern()` **Signature:** + ```python NkiTensor.get_pattern() ``` @@ -164,6 +169,7 @@ the same order as `shape` / `strides`. `NkiTensor.indirect(index, num_elem)` **Signature:** + ```python NkiTensor.indirect(index, num_elem=None) ``` @@ -181,10 +187,10 @@ comes from `index[i % G, i // G]` where G is the group size (16 for vector/scalar/gpsimd engines, 32 for tensor engine). - **index** — SBUF tensor containing free-dimension offsets, shape `(P, K)` - where `P == self.shape[0]`. + where `P == self.shape[0]`. - **num_elem** — number of offsets to use. Defaults to `index.size`. -**Returns:** new `NkiTensor` view with TI attached. Output shape is - `(P, num_elem)`. + **Returns:** new `NkiTensor` view with TI attached. Output shape is + `(P, num_elem)`. --- @@ -193,6 +199,7 @@ vector/scalar/gpsimd engines, 32 for tensor engine). `NkiTensor.is_contiguous()` **Signature:** + ```python NkiTensor.is_contiguous() ``` @@ -214,6 +221,7 @@ per-partition over the free dims only. `NkiTensor.is_indirect()` **Signature:** + ```python NkiTensor.is_indirect() ``` @@ -233,6 +241,7 @@ sliced or selected — use this query to guard against those chains. `NkiTensor.permute(dims)` **Signature:** + ```python NkiTensor.permute(dims) ``` @@ -256,7 +265,7 @@ t.permute((0, 2, 1)) # shape becomes (128, 8, 4) stays outermost; see `is_indirect`) - **dims** — tuple of dimension indices in the desired order -**Returns:** new `NkiTensor` view with reordered dimensions + **Returns:** new `NkiTensor` view with reordered dimensions --- @@ -265,6 +274,7 @@ t.permute((0, 2, 1)) # shape becomes (128, 8, 4) `NkiTensor.rearrange(src_pattern, dst_pattern, fixed_sizes)` **Signature:** + ```python NkiTensor.rearrange(src_pattern, dst_pattern, fixed_sizes=None) ``` @@ -286,7 +296,7 @@ t.rearrange(('b', ('h', 'w')), ('b', 'w', 'h'), {'h': 4}) - **src_pattern** — source dimension pattern (tuple of str or tuple-of-str) - **dst_pattern** — destination dimension pattern (same dimension names) - **fixed_sizes** — dict mapping dimension names to known sizes (for -1 inference) -**Returns:** new `NkiTensor` view with rearranged dimensions + **Returns:** new `NkiTensor` view with rearranged dimensions --- @@ -295,6 +305,7 @@ t.rearrange(('b', ('h', 'w')), ('b', 'w', 'h'), {'h': 4}) `NkiTensor.reshape(shape)` **Signature:** + ```python NkiTensor.reshape(shape) ``` @@ -321,7 +332,7 @@ t.reshape((128, 2, 12)) # split differently - Fails if the current layout is incompatible with the requested shape - **shape** — tuple of new dimension sizes -**Returns:** new `NkiTensor` view with the requested shape + **Returns:** new `NkiTensor` view with the requested shape --- @@ -330,6 +341,7 @@ t.reshape((128, 2, 12)) # split differently `NkiTensor.reshape_dim(dim, shape)` **Signature:** + ```python NkiTensor.reshape_dim(dim, shape) ``` @@ -354,7 +366,7 @@ t.reshape_dim(1, (4, -1)) # same result, 6 is inferred - **dim** — dimension to split - **shape** — tuple of sizes for the new dimensions (may contain one -1) -**Returns:** new `NkiTensor` view with the dimension split + **Returns:** new `NkiTensor` view with the dimension split --- @@ -363,6 +375,7 @@ t.reshape_dim(1, (4, -1)) # same result, 6 is inferred `NkiTensor.select(dim, index)` **Signature:** + ```python NkiTensor.select(dim, index) ``` @@ -396,7 +409,7 @@ hbm_t.select(0, idx) # shape becomes (128, 8) - **dim** — dimension to select from - **index** — integer index (static) or `NkiTensor` scalar (dynamic) -**Returns:** new `NkiTensor` view with the dimension removed + **Returns:** new `NkiTensor` view with the dimension removed --- @@ -405,6 +418,7 @@ hbm_t.select(0, idx) # shape becomes (128, 8) `NkiTensor.slice(dim, start, end, step)` **Signature:** + ```python NkiTensor.slice(dim, start, end, step=1) ``` @@ -433,7 +447,7 @@ t.slice(1, 0, 64, 2) # shape becomes (128, 32) - **start** — start index (inclusive) - **end** — end index (exclusive) - **step** — step size (default 1) -**Returns:** new `NkiTensor` view with the sliced dimension + **Returns:** new `NkiTensor` view with the sliced dimension --- @@ -442,6 +456,7 @@ t.slice(1, 0, 64, 2) # shape becomes (128, 32) `NkiTensor.squeeze_dim(dim)` **Signature:** + ```python NkiTensor.squeeze_dim(dim) ``` @@ -461,7 +476,7 @@ t.squeeze_dim(1) # shape becomes (128, 64) - After `vector_select`: `dim` must not be 0 - **dim** — dimension to remove (must have size 1) -**Returns:** new `NkiTensor` view with the dimension removed + **Returns:** new `NkiTensor` view with the dimension removed --- @@ -470,6 +485,7 @@ t.squeeze_dim(1) # shape becomes (128, 64) `NkiTensor.vector_select(dim, vector_offset)` **Signature:** + ```python NkiTensor.vector_select(dim, vector_offset) ``` @@ -503,7 +519,7 @@ hbm_t.vector_select(0, offsets) # shape becomes (128, 128, 8) - **dim** — dimension to apply indirect addressing (must be 0) - **vector_offset** — SBUF tensor with per-partition indices, shape `(num_partitions, 1)` -**Returns:** new `NkiTensor` view with dim 0 size set to `vector_offset.shape[0]` + **Returns:** new `NkiTensor` view with dim 0 size set to `vector_offset.shape[0]` --- @@ -512,6 +528,7 @@ hbm_t.vector_select(0, offsets) # shape becomes (128, 128, 8) `NkiTensor.view(dtype)` **Signature:** + ```python NkiTensor.view(dtype) ``` @@ -541,6 +558,6 @@ u.view(nl.float32) # shape becomes (128, 64), 4x contraction - Not supported after dynamic / vector select - **dtype** — target NKI dtype to reinterpret as -**Returns:** new `NkiTensor` view with the adjusted dtype and shape + **Returns:** new `NkiTensor` view with the adjusted dtype and shape --- diff --git a/skills/neuron-nki-docs/references/programming/api/api-nki-tools.md b/skills/neuron-nki-docs/references/programming/api/api-nki-tools.md index 4f7baa5..f5f16b6 100644 --- a/skills/neuron-nki-docs/references/programming/api/api-nki-tools.md +++ b/skills/neuron-nki-docs/references/programming/api/api-nki-tools.md @@ -15,10 +15,11 @@ Profiling, benchmarking, and simulation tools. nki.jit -nki.jit(*func=None*, ***kwargs*)[[source]](../../../_modules/nki.html#jit) +nki.jit(_func=None_, _\*\*kwargs_)[[source]](../../../\_modules/nki.html#jit) This decorator compiles a top-level NKI function to run on NeuronDevices. The NKI Compiler automatically detects the appropriate machine learning framework based on the kernel arguments: + - **Torch tensors**: uses TorchXLA integration. - **JAX arrays**: uses JAX integration. - **NumPy arrays**: runs the kernel in standalone mode without a machine learning framework. @@ -33,7 +34,7 @@ The `platform_target` and `mode` parameters are removed from `@nki.jit`. The com Parameters: -* **func** – Function that defines the custom operation. +- **func** – Function that defines the custom operation. Listing 11 Writing an addition kernel using `@nki.jit` diff --git a/skills/neuron-nki-docs/references/programming/api/api-overview.md b/skills/neuron-nki-docs/references/programming/api/api-overview.md index db97499..da0053e 100644 --- a/skills/neuron-nki-docs/references/programming/api/api-overview.md +++ b/skills/neuron-nki-docs/references/programming/api/api-overview.md @@ -7,36 +7,27 @@ The NKI Library provides pre-built reference kernels you can use directly in you ## Normalization and Quantization Kernels - | [ RMSNorm-Quant Kernel API Reference ](../../reference/library/rmsnorm-quant.md) | API reference for the RMSNorm-Quant kernel included in the NKI Library. The kernel performs optional RMS normalization followed by quantization to fp8 . | -| --- | --- | - +| -------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------- | ## QKV Projection Kernels - | [ QKV Kernel API Reference ](../../reference/library/qkv.md) | API reference for the QKV kernel included in the NKI Library. The kernel performs Query-Key-Value projection with optional normalization fusion. | -| --- | --- | - +| ------------------------------------------------------------ | ------------------------------------------------------------------------------------------------------------------------------------------------ | ## Attention Kernels - -| [ Attention CTE Kernel API Reference ](../../reference/library/attention-cte.md) | API reference for the Attention CTE kernel included in the NKI Library. The kernel implements attention specifically optimized for Context Encoding use cases. | -| --- | --- | +| [ Attention CTE Kernel API Reference ](../../reference/library/attention-cte.md) | API reference for the Attention CTE kernel included in the NKI Library. The kernel implements attention specifically optimized for Context Encoding use cases. | +| -------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | | [ Attention TKG Kernel API Reference ](../../reference/library/attention-tkg.md) | API reference for the Attention TKG kernel included in the NKI Library. The kernel implements attention specifically optimized for Token Generation (Decoding) use cases with small active sequence lengths. | - ## Multi-Layer Perceptron (MLP) Kernels - | [ MLP Kernel API Reference ](../../reference/library/mlp.md) | API reference for the MLP kernel included in the NKI Library. The kernel implements a Multi-Layer Perceptron with optional normalization fusion and various optimizations. | -| --- | --- | - +| ------------------------------------------------------------ | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ## Output Projection Kernels - | [ Output Projection CTE Kernel API Reference ](../../reference/library/output-projection-cte.md) | API reference for the Output Projection CTE kernel included in the NKI Library. The kernel computes the output projection operation optimized for Context Encoding use cases. | -| --- | --- | -| [ Output Projection TKG Kernel API Reference ](../../reference/library/output-projection-tkg.md) | API reference for the Output Projection TKG kernel included in the NKI Library. The kernel computes the output projection operation optimized for Token Generation use cases. | \ No newline at end of file +| ------------------------------------------------------------------------------------------------ | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| [ Output Projection TKG Kernel API Reference ](../../reference/library/output-projection-tkg.md) | API reference for the Output Projection TKG kernel included in the NKI Library. The kernel computes the output projection operation optimized for Token Generation use cases. | diff --git a/skills/neuron-nki-docs/references/programming/api/index.md b/skills/neuron-nki-docs/references/programming/api/index.md index 2e1dbd8..5e1b57b 100644 --- a/skills/neuron-nki-docs/references/programming/api/index.md +++ b/skills/neuron-nki-docs/references/programming/api/index.md @@ -19,7 +19,6 @@ Complete index of all NKI API functions organized by module. ## Function Index - ### N - `nki.isa.activate2` - [api-nki-isa-scalar.md](api-nki-isa-scalar.md#nki-isa-activate2) @@ -107,4 +106,4 @@ Complete index of all NKI API functions organized by module. - `nki.language.uint16` - [api-nki-language-misc.md](api-nki-language-misc.md#nki-language-uint16) - `nki.language.uint32` - [api-nki-language-misc.md](api-nki-language-misc.md#nki-language-uint32) - `nki.language.uint8` - [api-nki-language-misc.md](api-nki-language-misc.md#nki-language-uint8) -- `nki.language.zeros` - [api-nki-language-creation.md](api-nki-language-creation.md#nki-language-zeros) \ No newline at end of file +- `nki.language.zeros` - [api-nki-language-creation.md](api-nki-language-creation.md#nki-language-zeros) diff --git a/skills/neuron-nki-docs/references/programming/api/nki.api.shared.md b/skills/neuron-nki-docs/references/programming/api/nki.api.shared.md index 15027f1..f74bd15 100644 --- a/skills/neuron-nki-docs/references/programming/api/nki.api.shared.md +++ b/skills/neuron-nki-docs/references/programming/api/nki.api.shared.md @@ -3,120 +3,117 @@ NKI API Common Fields + ## Supported Data Types [Supported Data Types by NKI](#tbl-dtype) below lists all supported data types by NKI. Almost all of the NKI APIs accept a data type field, dtype, which must be a nki.language data type. - -| | Data Type | Accepted dtype Field by NKI APIs | -| --- | --- | --- | -| Integer | 8-bit unsigned integer | nki.language.uint8 | -| 8-bit signed integer | nki.language.int8 | -| 16-bit unsigned integer | nki.language.uint16 | -| 16-bit signed integer | nki.language.int16 | -| 32-bit unsigned integer | nki.language.uint32 | -| 32-bit signed integer | nki.language.int32 | -| Float | float8_e4m3 (1S,4E,3M) [ [ 2 ] ](#id2) | nki.language.float8_e4m3 | -| float8_e5m2 (1S,5E,2M) | nki.language.float8_e5m2 | -| float16 (1S,5E,10M) | nki.language.float16 | -| bfloat16 (1S,8E,7M) | nki.language.bfloat16 | -| tfloat32 (1S,8E,10M) | nki.language.tfloat32 | -| float32 (1S,8E,23M) | nki.language.float32 | -| Boolean | boolean stored as uint8 | nki.language.bool_ | - +| | Data Type | Accepted dtype Field by NKI APIs | +| ----------------------- | -------------------------------------- | -------------------------------- | +| Integer | 8-bit unsigned integer | nki.language.uint8 | +| 8-bit signed integer | nki.language.int8 | +| 16-bit unsigned integer | nki.language.uint16 | +| 16-bit signed integer | nki.language.int16 | +| 32-bit unsigned integer | nki.language.uint32 | +| 32-bit signed integer | nki.language.int32 | +| Float | float8_e4m3 (1S,4E,3M) [ [ 2 ] ](#id2) | nki.language.float8_e4m3 | +| float8_e5m2 (1S,5E,2M) | nki.language.float8_e5m2 | +| float16 (1S,5E,10M) | nki.language.float16 | +| bfloat16 (1S,8E,7M) | nki.language.bfloat16 | +| tfloat32 (1S,8E,10M) | nki.language.tfloat32 | +| float32 (1S,8E,23M) | nki.language.float32 | +| Boolean | boolean stored as uint8 | nki.language.bool\_ | + ## Supported Math Operators for NKI ISA [Supported Math Operators by NKI ISA](#tbl-aluop) below lists all the mathematical operator primitives supported by NKI. Many [nki.isa](nki.isa.md#id1) APIs (instructions) allow programmable operators through the `op` field. -The supported operators fall into two categories: *bitvec* and *arithmetic*. In general, instructions -using *bitvec* operators expect integer data types and treat input elements as bit patterns. On the other -hand, instructions using *arithmetic* operators accept any valid NKI data type and convert input elements +The supported operators fall into two categories: _bitvec_ and _arithmetic_. In general, instructions +using _bitvec_ operators expect integer data types and treat input elements as bit patterns. On the other +hand, instructions using _arithmetic_ operators accept any valid NKI data type and convert input elements into float32 before performing the operators. - -| | Operator | op | Legal Reduction op | -| --- | --- | --- | --- | -| Bitvec | Bitwise Not | nki.language.invert | N | -| Bitwise And | nki.language.bitwise_and | Y | -| Bitwise Or | nki.language.bitwise_or | Y | -| Bitwise Xor | nki.language.bitwise_xor | Y | -| Arithmetic Shift Left | nki.language.left_shift | N | -| Arithmetic Shift Right | Not supported | N | -| Logical Shift Left | nki.language.left_shift | N | -| Logical Shift Right | nki.language.right_shift | N | -| Arithmetic | Add | nki.language.add | Y | -| Subtract | nki.language.subtract | Y | -| Multiply | nki.language.multiply | Y | -| Max | nki.language.maximum | Y | -| Min | nki.language.minimum | Y | -| Is Equal to | nki.language.equal | N | -| Is Not Equal to | nki.language.not_equal | N | -| Is Greater than or Equal to | nki.language.greater_equal | N | -| Is Greater than to | nki.language.greater | N | -| Is Less than or Equal to | nki.language.less_equal | N | -| Is Less than | nki.language.less | N | -| Logical And | nki.language.logical_and | Y | -| Logical Or | nki.language.logical_or | Y | -| Logical Xor | nki.language.logical_xor | Y | -| Reverse Square Root | nki.language.rsqrt | N | -| Reciprocal | nki.language.reciprocal | N | -| Absolute | nki.language.abs | N | -| Absolute Maximum | nki.language.abs_max | N | -| Absolute Minimum | nki.language.abs_min | N | -| Square | nki.language.square | N | -| Relu | nki.language.relu | N | -| Power | nki.language.power | N | - +| | Operator | op | Legal Reduction op | +| --------------------------- | -------------------------- | ------------------- | ------------------ | +| Bitvec | Bitwise Not | nki.language.invert | N | +| Bitwise And | nki.language.bitwise_and | Y | +| Bitwise Or | nki.language.bitwise_or | Y | +| Bitwise Xor | nki.language.bitwise_xor | Y | +| Arithmetic Shift Left | nki.language.left_shift | N | +| Arithmetic Shift Right | Not supported | N | +| Logical Shift Left | nki.language.left_shift | N | +| Logical Shift Right | nki.language.right_shift | N | +| Arithmetic | Add | nki.language.add | Y | +| Subtract | nki.language.subtract | Y | +| Multiply | nki.language.multiply | Y | +| Max | nki.language.maximum | Y | +| Min | nki.language.minimum | Y | +| Is Equal to | nki.language.equal | N | +| Is Not Equal to | nki.language.not_equal | N | +| Is Greater than or Equal to | nki.language.greater_equal | N | +| Is Greater than to | nki.language.greater | N | +| Is Less than or Equal to | nki.language.less_equal | N | +| Is Less than | nki.language.less | N | +| Logical And | nki.language.logical_and | Y | +| Logical Or | nki.language.logical_or | Y | +| Logical Xor | nki.language.logical_xor | Y | +| Reverse Square Root | nki.language.rsqrt | N | +| Reciprocal | nki.language.reciprocal | N | +| Absolute | nki.language.abs | N | +| Absolute Maximum | nki.language.abs_max | N | +| Absolute Minimum | nki.language.abs_min | N | +| Square | nki.language.square | N | +| Relu | nki.language.relu | N | +| Power | nki.language.power | N | + ## Supported Activation Functions for NKI ISA [Supported Activation Functions by NKI ISA](#tbl-act-func) below lists all the activation function supported by the `nki.isa.activation` API. These activation functions are approximated with piece-wise polynomials on Scalar Engine. -*NOTE*: if input values fall outside the supported **Valid Input Range** listed below, +_NOTE_: if input values fall outside the supported **Valid Input Range** listed below, the Scalar Engine will generate invalid output results. - -| Function Name | Accepted op by Scalar Engine | Valid Input Range | -| --- | --- | --- | -| Identity | nki.language.copy | [-inf, inf] | -| Square | nki.language.square | [-inf, inf] | -| Sigmoid | nki.language.sigmoid | [-inf, inf] | -| Relu | nki.language.relu | [-inf, inf] | -| Gelu | nki.language.gelu | [-inf, inf] | -| Gelu Derivative | nki.language.gelu_dx | [-inf, inf] | -| Gelu with Tanh Approximation | nki.language.gelu_apprx_tanh | [-inf, inf] | -| Gelu with Sigmoid Approximation | nki.language.gelu_apprx_sigmoid | [-inf, inf] | -| Silu | nki.language.silu | [-inf, inf] | -| Silu Derivative | nki.language.silu_dx | [-inf, inf] | -| Tanh | nki.language.tanh | [-inf, inf] | -| Softplus | nki.language.softplus | [-inf, inf] | -| Mish | nki.language.mish | [-inf, inf] | -| Erf | nki.language.erf | [-inf, inf] | -| Erf Derivative | nki.language.erf_dx | [-inf, inf] | -| Exponential | nki.language.exp | [-inf, inf] | -| Natural Log | nki.language.log [2^-64, 2^64] | -| Sine | nki.language.sin | [-PI, PI] | -| Arctan | nki.language.arctan | [-PI/2, PI/2] | -| Square Root | nki.language.sqrt | [2^-116, 2^118] | -| Reverse Square Root | nki.language.rsqrt | [2^-87, 2^97] | -| Reciprocal | nki.language.reciprocal | ±[2^-42, 2^42] | -| Sign | nki.language.sign | [-inf, inf] | -| Absolute | nki.language.abs | [-inf, inf] | -| PReLU | nki.language.prelu | [-inf, inf] | -| Bypass (pass-through) | nki.language.bypass | [-inf, inf] | - +| Function Name | Accepted op by Scalar Engine | Valid Input Range | +| ------------------------------- | ------------------------------- | ----------------- | +| Identity | nki.language.copy | [-inf, inf] | +| Square | nki.language.square | [-inf, inf] | +| Sigmoid | nki.language.sigmoid | [-inf, inf] | +| Relu | nki.language.relu | [-inf, inf] | +| Gelu | nki.language.gelu | [-inf, inf] | +| Gelu Derivative | nki.language.gelu_dx | [-inf, inf] | +| Gelu with Tanh Approximation | nki.language.gelu_apprx_tanh | [-inf, inf] | +| Gelu with Sigmoid Approximation | nki.language.gelu_apprx_sigmoid | [-inf, inf] | +| Silu | nki.language.silu | [-inf, inf] | +| Silu Derivative | nki.language.silu_dx | [-inf, inf] | +| Tanh | nki.language.tanh | [-inf, inf] | +| Softplus | nki.language.softplus | [-inf, inf] | +| Mish | nki.language.mish | [-inf, inf] | +| Erf | nki.language.erf | [-inf, inf] | +| Erf Derivative | nki.language.erf_dx | [-inf, inf] | +| Exponential | nki.language.exp | [-inf, inf] | +| Natural Log | nki.language.log [2^-64, 2^64] | +| Sine | nki.language.sin | [-PI, PI] | +| Arctan | nki.language.arctan | [-PI/2, PI/2] | +| Square Root | nki.language.sqrt | [2^-116, 2^118] | +| Reverse Square Root | nki.language.rsqrt | [2^-87, 2^97] | +| Reciprocal | nki.language.reciprocal | ±[2^-42, 2^42] | +| Sign | nki.language.sign | [-inf, inf] | +| Absolute | nki.language.abs | [-inf, inf] | +| PReLU | nki.language.prelu | [-inf, inf] | +| Bypass (pass-through) | nki.language.bypass | [-inf, inf] | ## NKI Engine Selection for Operators Supported on Multiple Engines There is a tradeoff between precision and speed on different engines for operators with multiple engine options. Users can select which engine to map to based on their needs. We take reciprocal and reverse square root as two examples and explain the tradeoff below. -* Reciprocal can run on Scalar Engine or Vector Engine: +- Reciprocal can run on Scalar Engine or Vector Engine: Reciprocal can run on Vector Engine with `nki.isa.reciprocal` or on Scalar Engine with `nki.isa.activation(nl.reciprocal)`. Vector Engine performs reciprocal at a higher precision compared to Scalar Engine; however, the computation throughput of reciprocal on Vector Engine is about 8x lower than Scalar Engine for large @@ -125,22 +122,21 @@ cycles) dominates performance so Scalar Engine and Vector Engine have comparable **Estimated cycles on different engines:** - -| Cost (Engine Cycles) | Condition | -| --- | --- | -| max(MIN_II, N) | mapped to Scalar Engine nki.isa.scalar_engine | -| max(MIN_II, 8*N) | mapped to Vector Engine nki.isa.vector_engine | +| Cost (Engine Cycles) | Condition | +| -------------------- | --------------------------------------------- | +| max(MIN_II, N) | mapped to Scalar Engine nki.isa.scalar_engine | +| max(MIN_II, 8\*N) | mapped to Vector Engine nki.isa.vector_engine | where, -* `N` is the number of elements per partition in the input tile. +- `N` is the number of elements per partition in the input tile. -* `MIN_II` is the minimum instruction initiation interval for small input tiles. -`MIN_II` is roughly 64 engine cycles. +- `MIN_II` is the minimum instruction initiation interval for small input tiles. + `MIN_II` is roughly 64 engine cycles. **Note** `nki.isa.activation(op=nl.reciprocal)` doesn’t support setting bias on NeuronCore-v2. -* Reverse square root can run on GpSIMD Engine or Scalar Engine: +- Reverse square root can run on GpSIMD Engine or Scalar Engine: Reverse square root can run on GpSIMD Engine with `nki.isa.tensor_scalar(op0=nl.rsqrt, operand0=0.0)` or on Scalar Engine with `nki.isa.activation(nl.rsqrt)`. GpSIMD Engine performs reverse square root at a higher precision compared to Scalar Engine; however, the computation throughput of reverse square root on GpSIMD @@ -149,4 +145,4 @@ Engine is 4x lower than Scalar Engine. Footnotes [[2](#id1)] -S: sign bits, E: exponent bits, M: mantissa bits \ No newline at end of file +S: sign bits, E: exponent bits, M: mantissa bits diff --git a/skills/neuron-nki-docs/references/programming/api/nki.isa.md b/skills/neuron-nki-docs/references/programming/api/nki.isa.md index abaaeec..e41a69a 100644 --- a/skills/neuron-nki-docs/references/programming/api/nki.isa.md +++ b/skills/neuron-nki-docs/references/programming/api/nki.isa.md @@ -4,74 +4,69 @@ nki.isa ## NKI ISA - -| [ nc_matmul ](generated/nki.isa.nc_matmul.md#nki.isa.nc_matmul) | Compute dst = stationary.T @ moving matrix multiplication using Tensor Engine. | -| --- | --- | -| [ nc_matmul_mx ](generated/nki.isa.nc_matmul_mx.md#nki.isa.nc_matmul_mx) | Compute matrix multiplication of MXFP8/MXFP4 quantized matrices with integrated dequantization using Tensor Engine. | -| [ nc_transpose ](generated/nki.isa.nc_transpose.md#nki.isa.nc_transpose) | Perform a 2D transpose between the partition axis and the free axis of input data using Tensor or Vector Engine. | -| [ activation ](generated/nki.isa.activation.md#nki.isa.activation) | Apply an activation function on every element of the input tile using Scalar Engine, with an optional scale/bias operation before the activation and an optional reduction operation after the activation in the same instruction. | -| [ activation_reduce ](generated/nki.isa.activation_reduce.md#nki.isa.activation_reduce) | Perform the same computation as nisa.activation and also a reduction along the free dimension of the nisa.activation result using Scalar Engine. | -| [ activate2 ](generated/nki.isa.activate2.md#nki.isa.activate2) | Apply activation to result of two-stage tensor-scalar pipeline `(data op0 imm0) op1 imm1` with optional reduction, all in one Scalar Engine instruction. Trn3 only. | -| [ tensor_reduce ](generated/nki.isa.tensor_reduce.md#nki.isa.tensor_reduce) | Apply a reduction operation to the free axes of an input data tile using Vector Engine. | -| [ tensor_partition_reduce ](generated/nki.isa.tensor_partition_reduce.md#nki.isa.tensor_partition_reduce) | Apply a reduction operation across partitions of an input data tile using GpSimd Engine. | -| [ tensor_tensor ](generated/nki.isa.tensor_tensor.md#nki.isa.tensor_tensor) | Perform an element-wise operation of input two tiles using Vector Engine or GpSimd Engine. | -| [ tensor_tensor_scan ](generated/nki.isa.tensor_tensor_scan.md#nki.isa.tensor_tensor_scan) | Perform a scan operation of two input tiles using Vector Engine. | -| [ scalar_tensor_tensor ](generated/nki.isa.scalar_tensor_tensor.md#nki.isa.scalar_tensor_tensor) | Apply two math operators in sequence using Vector Engine: (data <op0> operand0) <op1> operand1 . | -| [ tensor_scalar ](generated/nki.isa.tensor_scalar.md#nki.isa.tensor_scalar) | Apply up to two math operators to the input data tile by broadcasting scalar/vector operands in the free dimension using Vector or Scalar or GpSimd Engine: (data <op0> operand0) <op1> operand1 . | -| [ tensor_scalar_reduce ](generated/nki.isa.tensor_scalar_reduce.md#nki.isa.tensor_scalar_reduce) | Perform the same computation as nisa.tensor_scalar with one math operator and also a reduction along the free dimension of the nisa.tensor_scalar result using Vector Engine. | -| [ tensor_scalar_cumulative ](generated/nki.isa.tensor_scalar_cumulative.md#nki.isa.tensor_scalar_cumulative) | Perform tensor-scalar arithmetic operation with cumulative reduction using Vector Engine. | -| [ tensor_copy ](generated/nki.isa.tensor_copy.md#nki.isa.tensor_copy) | Create a copy of src tile within NeuronCore on-chip SRAMs using Vector, Scalar or GpSimd Engine. | -| [ tensor_copy_predicated ](generated/nki.isa.tensor_copy_predicated.md#nki.isa.tensor_copy_predicated) | Conditionally copy elements from the src tile to the destination tile on SBUF / PSUM based on a predicate using Vector Engine. | -| [ reciprocal ](generated/nki.isa.reciprocal.md#nki.isa.reciprocal) | Compute element-wise reciprocal (1.0/x) of the input data tile using Vector Engine. | -| [ quantize_mx ](generated/nki.isa.quantize_mx.md#nki.isa.quantize_mx) | Quantize FP16/BF16 data to MXFP8 tensors (both data and scales) using Vector Engine. | -| [ iota ](generated/nki.isa.iota.md#nki.isa.iota) | Generate a constant literal pattern into SBUF using GpSimd Engine. | -| [ dropout ](generated/nki.isa.dropout.md#nki.isa.dropout) | Randomly replace some elements of the input tile data with zeros based on input probabilities using Vector Engine. | -| [ exponential ](generated/nki.isa.exponential.md#nki.isa.exponential) | Dedicated exponential instruction with max subtraction, faster than `nisa.activation(op=nl.exp)`. Trn3 (NeuronCore-v4) only. | -| [ affine_select ](generated/nki.isa.affine_select.md#nki.isa.affine_select) | Select elements between an input tile on_true_tile and a scalar value on_false_value according to a boolean predicate tile using GpSimd Engine. | -| [ range_select ](generated/nki.isa.range_select.md#nki.isa.range_select) | Select elements from on_true_tile based on comparison with bounds using Vector Engine. | -| [ select_reduce ](generated/nki.isa.select_reduce.md#nki.isa.select_reduce) | Selectively copy elements from either on_true or on_false to the destination tile based on a predicate using Vector Engine, with optional reduction (max). | -| [ sequence_bounds ](generated/nki.isa.sequence_bounds.md#nki.isa.sequence_bounds) | Compute the sequence bounds for a given set of segment IDs using GpSIMD Engine. | -| [ memset ](generated/nki.isa.memset.md#nki.isa.memset) | Initialize dst by filling it with a compile-time constant value , using Vector or GpSimd Engine. | -| [ bn_stats ](generated/nki.isa.bn_stats.md#nki.isa.bn_stats) | Compute mean- and variance-related statistics for each partition of an input tile data in parallel using Vector Engine. | -| [ bn_aggr ](generated/nki.isa.bn_aggr.md#nki.isa.bn_aggr) | Aggregate one or multiple bn_stats outputs to generate a mean and variance per partition using Vector Engine. | -| [ local_gather ](generated/nki.isa.local_gather.md#nki.isa.local_gather) | Gather SBUF data in src_buffer using index on GpSimd Engine. | -| [ dma_copy ](generated/nki.isa.dma_copy.md#nki.isa.dma_copy) | Copy data from src to dst using DMA engines with optional read-modify-write operations. | -| [ dma_transpose ](generated/nki.isa.dma_transpose.md#nki.isa.dma_transpose) | Perform a transpose on input src using DMA Engine. | -| [ dma_compute ](generated/nki.isa.dma_compute.md#nki.isa.dma_compute) | Perform math operations using compute logic inside DMA engines with element-wise scaling and reduction. | -| [ max8 ](generated/nki.isa.max8.md#nki.isa.max8) | Find the 8 largest values in each partition of the source tile. | -| [ nonzero_with_count ](generated/nki.isa.nonzero_with_count.md#nki.isa.nonzero_with_count) | Find indices of nonzero elements and their total count using GpSimd Engine. NeuronCore-v3+ only. | -| [ nc_n_gather ](generated/nki.isa.nc_n_gather.md#nki.isa.nc_n_gather) | Gather elements from data according to indices using GpSimd Engine. | -| [ nc_find_index8 ](generated/nki.isa.nc_find_index8.md#nki.isa.nc_find_index8) | Find indices of the 8 given vals in each partition of the data tensor. | -| [ nc_match_replace8 ](generated/nki.isa.nc_match_replace8.md#nki.isa.nc_match_replace8) | Replace first occurrence of each value in vals with imm in data using the Vector engine and return the replaced tensor. | -| [ nc_stream_shuffle ](generated/nki.isa.nc_stream_shuffle.md#nki.isa.nc_stream_shuffle) | Apply cross-partition data movement within a quadrant of 32 partitions from source tile src to destination tile dst using Vector Engine. | -| [ register_alloc ](generated/nki.isa.register_alloc.md#nki.isa.register_alloc) | Allocate a virtual register and optionally initialize it with an integer value x . | -| [ register_load ](generated/nki.isa.register_load.md#nki.isa.register_load) | Load a scalar value from memory (HBM or SBUF) into a virtual register. | -| [ register_move ](generated/nki.isa.register_move.md#nki.isa.register_move) | Move a value from a source VirtualRegister into a destination register. | -| [ register_store ](generated/nki.isa.register_store.md#nki.isa.register_store) | Store the value from a virtual register into memory (HBM/SBUF). | -| [ core_barrier ](generated/nki.isa.core_barrier.md#nki.isa.core_barrier) | Synchronize execution across multiple NeuronCores by implementing a barrier mechanism. | -| [ sendrecv ](generated/nki.isa.sendrecv.md#nki.isa.sendrecv) | Perform point-to-point communication between NeuronCores by sending and receiving data simultaneously using DMA engines. Uses `dma_engine` enum for engine selection. | -| [ rng ](generated/nki.isa.rng.md#nki.isa.rng) | Generate pseudo random numbers using the Vector or GpSimd Engine. | -| [ rand2 ](generated/nki.isa.rand2.md#nki.isa.rand2) | Generate pseudo random numbers with uniform distribution using Vector Engine. | -| [ rand_set_state ](generated/nki.isa.rand_set_state.md#nki.isa.rand_set_state) | Seed the pseudo random number generator (PRNG) inside the engine. | -| [ rand_get_state ](generated/nki.isa.rand_get_state.md#nki.isa.rand_get_state) | Store the current pseudo random number generator (PRNG) states from the engine to SBUF. | -| [ set_rng_seed ](generated/nki.isa.set_rng_seed.md#nki.isa.set_rng_seed) | Seed the pseudo random number generator (PRNG) inside the Vector Engine. | - +| [ nc_matmul ](generated/nki.isa.nc_matmul.md#nki.isa.nc_matmul) | Compute dst = stationary.T @ moving matrix multiplication using Tensor Engine. | +| ------------------------------------------------------------------------------------------------------------ | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| [ nc_matmul_mx ](generated/nki.isa.nc_matmul_mx.md#nki.isa.nc_matmul_mx) | Compute matrix multiplication of MXFP8/MXFP4 quantized matrices with integrated dequantization using Tensor Engine. | +| [ nc_transpose ](generated/nki.isa.nc_transpose.md#nki.isa.nc_transpose) | Perform a 2D transpose between the partition axis and the free axis of input data using Tensor or Vector Engine. | +| [ activation ](generated/nki.isa.activation.md#nki.isa.activation) | Apply an activation function on every element of the input tile using Scalar Engine, with an optional scale/bias operation before the activation and an optional reduction operation after the activation in the same instruction. | +| [ activation_reduce ](generated/nki.isa.activation_reduce.md#nki.isa.activation_reduce) | Perform the same computation as nisa.activation and also a reduction along the free dimension of the nisa.activation result using Scalar Engine. | +| [ activate2 ](generated/nki.isa.activate2.md#nki.isa.activate2) | Apply activation to result of two-stage tensor-scalar pipeline `(data op0 imm0) op1 imm1` with optional reduction, all in one Scalar Engine instruction. Trn3 only. | +| [ tensor_reduce ](generated/nki.isa.tensor_reduce.md#nki.isa.tensor_reduce) | Apply a reduction operation to the free axes of an input data tile using Vector Engine. | +| [ tensor_partition_reduce ](generated/nki.isa.tensor_partition_reduce.md#nki.isa.tensor_partition_reduce) | Apply a reduction operation across partitions of an input data tile using GpSimd Engine. | +| [ tensor_tensor ](generated/nki.isa.tensor_tensor.md#nki.isa.tensor_tensor) | Perform an element-wise operation of input two tiles using Vector Engine or GpSimd Engine. | +| [ tensor_tensor_scan ](generated/nki.isa.tensor_tensor_scan.md#nki.isa.tensor_tensor_scan) | Perform a scan operation of two input tiles using Vector Engine. | +| [ scalar_tensor_tensor ](generated/nki.isa.scalar_tensor_tensor.md#nki.isa.scalar_tensor_tensor) | Apply two math operators in sequence using Vector Engine: (data <op0> operand0) <op1> operand1 . | +| [ tensor_scalar ](generated/nki.isa.tensor_scalar.md#nki.isa.tensor_scalar) | Apply up to two math operators to the input data tile by broadcasting scalar/vector operands in the free dimension using Vector or Scalar or GpSimd Engine: (data <op0> operand0) <op1> operand1 . | +| [ tensor_scalar_reduce ](generated/nki.isa.tensor_scalar_reduce.md#nki.isa.tensor_scalar_reduce) | Perform the same computation as nisa.tensor_scalar with one math operator and also a reduction along the free dimension of the nisa.tensor_scalar result using Vector Engine. | +| [ tensor_scalar_cumulative ](generated/nki.isa.tensor_scalar_cumulative.md#nki.isa.tensor_scalar_cumulative) | Perform tensor-scalar arithmetic operation with cumulative reduction using Vector Engine. | +| [ tensor_copy ](generated/nki.isa.tensor_copy.md#nki.isa.tensor_copy) | Create a copy of src tile within NeuronCore on-chip SRAMs using Vector, Scalar or GpSimd Engine. | +| [ tensor_copy_predicated ](generated/nki.isa.tensor_copy_predicated.md#nki.isa.tensor_copy_predicated) | Conditionally copy elements from the src tile to the destination tile on SBUF / PSUM based on a predicate using Vector Engine. | +| [ reciprocal ](generated/nki.isa.reciprocal.md#nki.isa.reciprocal) | Compute element-wise reciprocal (1.0/x) of the input data tile using Vector Engine. | +| [ quantize_mx ](generated/nki.isa.quantize_mx.md#nki.isa.quantize_mx) | Quantize FP16/BF16 data to MXFP8 tensors (both data and scales) using Vector Engine. | +| [ iota ](generated/nki.isa.iota.md#nki.isa.iota) | Generate a constant literal pattern into SBUF using GpSimd Engine. | +| [ dropout ](generated/nki.isa.dropout.md#nki.isa.dropout) | Randomly replace some elements of the input tile data with zeros based on input probabilities using Vector Engine. | +| [ exponential ](generated/nki.isa.exponential.md#nki.isa.exponential) | Dedicated exponential instruction with max subtraction, faster than `nisa.activation(op=nl.exp)`. Trn3 (NeuronCore-v4) only. | +| [ affine_select ](generated/nki.isa.affine_select.md#nki.isa.affine_select) | Select elements between an input tile on_true_tile and a scalar value on_false_value according to a boolean predicate tile using GpSimd Engine. | +| [ range_select ](generated/nki.isa.range_select.md#nki.isa.range_select) | Select elements from on_true_tile based on comparison with bounds using Vector Engine. | +| [ select_reduce ](generated/nki.isa.select_reduce.md#nki.isa.select_reduce) | Selectively copy elements from either on_true or on_false to the destination tile based on a predicate using Vector Engine, with optional reduction (max). | +| [ sequence_bounds ](generated/nki.isa.sequence_bounds.md#nki.isa.sequence_bounds) | Compute the sequence bounds for a given set of segment IDs using GpSIMD Engine. | +| [ memset ](generated/nki.isa.memset.md#nki.isa.memset) | Initialize dst by filling it with a compile-time constant value , using Vector or GpSimd Engine. | +| [ bn_stats ](generated/nki.isa.bn_stats.md#nki.isa.bn_stats) | Compute mean- and variance-related statistics for each partition of an input tile data in parallel using Vector Engine. | +| [ bn_aggr ](generated/nki.isa.bn_aggr.md#nki.isa.bn_aggr) | Aggregate one or multiple bn_stats outputs to generate a mean and variance per partition using Vector Engine. | +| [ local_gather ](generated/nki.isa.local_gather.md#nki.isa.local_gather) | Gather SBUF data in src_buffer using index on GpSimd Engine. | +| [ dma_copy ](generated/nki.isa.dma_copy.md#nki.isa.dma_copy) | Copy data from src to dst using DMA engines with optional read-modify-write operations. | +| [ dma_transpose ](generated/nki.isa.dma_transpose.md#nki.isa.dma_transpose) | Perform a transpose on input src using DMA Engine. | +| [ dma_compute ](generated/nki.isa.dma_compute.md#nki.isa.dma_compute) | Perform math operations using compute logic inside DMA engines with element-wise scaling and reduction. | +| [ max8 ](generated/nki.isa.max8.md#nki.isa.max8) | Find the 8 largest values in each partition of the source tile. | +| [ nonzero_with_count ](generated/nki.isa.nonzero_with_count.md#nki.isa.nonzero_with_count) | Find indices of nonzero elements and their total count using GpSimd Engine. NeuronCore-v3+ only. | +| [ nc_n_gather ](generated/nki.isa.nc_n_gather.md#nki.isa.nc_n_gather) | Gather elements from data according to indices using GpSimd Engine. | +| [ nc_find_index8 ](generated/nki.isa.nc_find_index8.md#nki.isa.nc_find_index8) | Find indices of the 8 given vals in each partition of the data tensor. | +| [ nc_match_replace8 ](generated/nki.isa.nc_match_replace8.md#nki.isa.nc_match_replace8) | Replace first occurrence of each value in vals with imm in data using the Vector engine and return the replaced tensor. | +| [ nc_stream_shuffle ](generated/nki.isa.nc_stream_shuffle.md#nki.isa.nc_stream_shuffle) | Apply cross-partition data movement within a quadrant of 32 partitions from source tile src to destination tile dst using Vector Engine. | +| [ register_alloc ](generated/nki.isa.register_alloc.md#nki.isa.register_alloc) | Allocate a virtual register and optionally initialize it with an integer value x . | +| [ register_load ](generated/nki.isa.register_load.md#nki.isa.register_load) | Load a scalar value from memory (HBM or SBUF) into a virtual register. | +| [ register_move ](generated/nki.isa.register_move.md#nki.isa.register_move) | Move a value from a source VirtualRegister into a destination register. | +| [ register_store ](generated/nki.isa.register_store.md#nki.isa.register_store) | Store the value from a virtual register into memory (HBM/SBUF). | +| [ core_barrier ](generated/nki.isa.core_barrier.md#nki.isa.core_barrier) | Synchronize execution across multiple NeuronCores by implementing a barrier mechanism. | +| [ sendrecv ](generated/nki.isa.sendrecv.md#nki.isa.sendrecv) | Perform point-to-point communication between NeuronCores by sending and receiving data simultaneously using DMA engines. Uses `dma_engine` enum for engine selection. | +| [ rng ](generated/nki.isa.rng.md#nki.isa.rng) | Generate pseudo random numbers using the Vector or GpSimd Engine. | +| [ rand2 ](generated/nki.isa.rand2.md#nki.isa.rand2) | Generate pseudo random numbers with uniform distribution using Vector Engine. | +| [ rand_set_state ](generated/nki.isa.rand_set_state.md#nki.isa.rand_set_state) | Seed the pseudo random number generator (PRNG) inside the engine. | +| [ rand_get_state ](generated/nki.isa.rand_get_state.md#nki.isa.rand_get_state) | Store the current pseudo random number generator (PRNG) states from the engine to SBUF. | +| [ set_rng_seed ](generated/nki.isa.set_rng_seed.md#nki.isa.set_rng_seed) | Seed the pseudo random number generator (PRNG) inside the Vector Engine. | ## NKI ISA Config Enums - -| [ engine ](generated/nki.isa.engine.md#nki.isa.engine) | Neuron Device engines | -| --- | --- | -| [ reduce_cmd ](generated/nki.isa.reduce_cmd.md#nki.isa.reduce_cmd) | Engine Register Reduce commands | -| [ dge_mode ](generated/nki.isa.dge_mode.md#nki.isa.dge_mode) | Neuron Descriptor Generation Engine Mode | +| [ engine ](generated/nki.isa.engine.md#nki.isa.engine) | Neuron Device engines | +| ------------------------------------------------------------------ | ------------------------------------------------------------------------------------------------------------------------- | +| [ reduce_cmd ](generated/nki.isa.reduce_cmd.md#nki.isa.reduce_cmd) | Engine Register Reduce commands | +| [ dge_mode ](generated/nki.isa.dge_mode.md#nki.isa.dge_mode) | Neuron Descriptor Generation Engine Mode | | [ dma_engine ](generated/nki.isa.dma_engine.md#nki.isa.dma_engine) | DMA transfer engine selection (`dma_engine.dma` for shared DMA, `dma_engine.gpsimd_dma` for GPSIMD's internal DMA engine) | -| [ oob_mode ](generated/nki.isa.oob_mode.md#nki.isa.oob_mode) | Out-of-bounds handling mode (`oob_mode.error`, `oob_mode.skip`) | -| [ nc_version ](generated/nki.isa.nc_version.md#nki.isa.nc_version) | NeuronCore version enum | - +| [ oob_mode ](generated/nki.isa.oob_mode.md#nki.isa.oob_mode) | Out-of-bounds handling mode (`oob_mode.error`, `oob_mode.skip`) | +| [ nc_version ](generated/nki.isa.nc_version.md#nki.isa.nc_version) | NeuronCore version enum | ## Target - -| [ nc_version ](generated/nki.isa.nc_version.md#nki.isa.nc_version) | NeuronCore version | -| --- | --- | -| [ get_nc_version ](generated/nki.isa.get_nc_version.md#nki.isa.get_nc_version) | Returns the nc_version of the current target context. | \ No newline at end of file +| [ nc_version ](generated/nki.isa.nc_version.md#nki.isa.nc_version) | NeuronCore version | +| ------------------------------------------------------------------------------ | ----------------------------------------------------- | +| [ get_nc_version ](generated/nki.isa.get_nc_version.md#nki.isa.get_nc_version) | Returns the nc_version of the current target context. | diff --git a/skills/neuron-nki-docs/references/programming/api/nki.language.md b/skills/neuron-nki-docs/references/programming/api/nki.language.md index 3bba949..59d5086 100644 --- a/skills/neuron-nki-docs/references/programming/api/nki.language.md +++ b/skills/neuron-nki-docs/references/programming/api/nki.language.md @@ -4,74 +4,62 @@ nki.language ## Creation operations - -| [ ndarray ](generated/nki.language.ndarray.md#nki.language.ndarray) | Create a new tensor of given shape and dtype on the specified buffer. | -| --- | --- | -| [ zeros ](generated/nki.language.zeros.md#nki.language.zeros) | Create a new tensor of given shape and dtype on the specified buffer, filled with zeros. | - +| [ ndarray ](generated/nki.language.ndarray.md#nki.language.ndarray) | Create a new tensor of given shape and dtype on the specified buffer. | +| ------------------------------------------------------------------- | ---------------------------------------------------------------------------------------- | +| [ zeros ](generated/nki.language.zeros.md#nki.language.zeros) | Create a new tensor of given shape and dtype on the specified buffer, filled with zeros. | ## Tensor manipulation operations - | [ ds ](generated/nki.language.ds.md#nki.language.ds) | Construct a dynamic slice for simple tensor indexing. | -| --- | --- | - +| ---------------------------------------------------- | ----------------------------------------------------- | ## Iterators In NKI 0.3.0, all range iterators are unified and have identical effect. Use standard Python `range` for all loops. The legacy `nl.affine_range`, `nl.sequential_range`, and `nl.static_range` are retained as aliases but have no distinct behavior. -| [ static_range ](generated/nki.language.static_range.md#nki.language.static_range) | Legacy alias for `range`. | -| --- | --- | -| [ affine_range ](generated/nki.language.affine_range.md#nki.language.affine_range) | Legacy alias for `range`. | +| [ static_range ](generated/nki.language.static_range.md#nki.language.static_range) | Legacy alias for `range`. | +| ---------------------------------------------------------------------------------------------- | ------------------------- | +| [ affine_range ](generated/nki.language.affine_range.md#nki.language.affine_range) | Legacy alias for `range`. | | [ sequential_range ](generated/nki.language.sequential_range.md#nki.language.sequential_range) | Legacy alias for `range`. | - ## Memory Hierarchy - -| [ psum ](generated/nki.language.psum.md#nki.language.psum) | PSUM - Only visible to each individual kernel instance in the SPMD grid | -| --- | --- | -| [ sbuf ](generated/nki.language.sbuf.md#nki.language.sbuf) | State Buffer - Only visible to each individual kernel instance in the SPMD grid | -| [ hbm ](generated/nki.language.hbm.md#nki.language.hbm) | HBM - Alias of private_hbm | -| [ private_hbm ](generated/nki.language.private_hbm.md#nki.language.private_hbm) | HBM - Only visible to each individual kernel instance in the SPMD grid | -| [ shared_hbm ](generated/nki.language.shared_hbm.md#nki.language.shared_hbm) | Shared HBM - Visible to all kernel instances in the SPMD grid | - +| [ psum ](generated/nki.language.psum.md#nki.language.psum) | PSUM - Only visible to each individual kernel instance in the SPMD grid | +| ------------------------------------------------------------------------------- | ------------------------------------------------------------------------------- | +| [ sbuf ](generated/nki.language.sbuf.md#nki.language.sbuf) | State Buffer - Only visible to each individual kernel instance in the SPMD grid | +| [ hbm ](generated/nki.language.hbm.md#nki.language.hbm) | HBM - Alias of private_hbm | +| [ private_hbm ](generated/nki.language.private_hbm.md#nki.language.private_hbm) | HBM - Only visible to each individual kernel instance in the SPMD grid | +| [ shared_hbm ](generated/nki.language.shared_hbm.md#nki.language.shared_hbm) | Shared HBM - Visible to all kernel instances in the SPMD grid | ## Others - -| [ program_id ](generated/nki.language.program_id.md#nki.language.program_id) | Index of the current SPMD program along the given axis in the launch grid. | -| --- | --- | -| [ num_programs ](generated/nki.language.num_programs.md#nki.language.num_programs) | Number of SPMD programs along the given axes in the launch grid. | -| [ program_ndim ](generated/nki.language.program_ndim.md#nki.language.program_ndim) | Number of dimensions in the SPMD launch grid. | +| [ program_id ](generated/nki.language.program_id.md#nki.language.program_id) | Index of the current SPMD program along the given axis in the launch grid. | +| ---------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------- | +| [ num_programs ](generated/nki.language.num_programs.md#nki.language.num_programs) | Number of SPMD programs along the given axes in the launch grid. | +| [ program_ndim ](generated/nki.language.program_ndim.md#nki.language.program_ndim) | Number of dimensions in the SPMD launch grid. | | [ device_print ](generated/nki.language.device_print.md#nki.language.device_print) | Print a message with a string print_prefix followed by the value of a tile tensor . | - ## Data Types - -| [ bool_ ](generated/nki.language.bool_.md#nki.language.bool_) | Boolean (True or False) stored as a byte | -| --- | --- | -| [ int8 ](generated/nki.language.int8.md#nki.language.int8) | 8-bit signed integer number | -| [ int16 ](generated/nki.language.int16.md#nki.language.int16) | 16-bit signed integer number | -| [ int32 ](generated/nki.language.int32.md#nki.language.int32) | 32-bit signed integer number | -| [ uint8 ](generated/nki.language.uint8.md#nki.language.uint8) | 8-bit unsigned integer number | -| [ uint16 ](generated/nki.language.uint16.md#nki.language.uint16) | 16-bit unsigned integer number | -| [ uint32 ](generated/nki.language.uint32.md#nki.language.uint32) | 32-bit unsigned integer number | -| [ float16 ](generated/nki.language.float16.md#nki.language.float16) | 16-bit floating-point number | -| [ float32 ](generated/nki.language.float32.md#nki.language.float32) | 32-bit floating-point number | -| [ bfloat16 ](generated/nki.language.bfloat16.md#nki.language.bfloat16) | 16-bit floating-point number (1S,8E,7M) | -| [ tfloat32 ](generated/nki.language.tfloat32.md#nki.language.tfloat32) | 32-bit floating-point number (1S,8E,10M) | -| [ float8_e4m3 ](generated/nki.language.float8_e4m3.md#nki.language.float8_e4m3) | 8-bit floating-point number (1S,4E,3M) | -| [ float8_e5m2 ](generated/nki.language.float8_e5m2.md#nki.language.float8_e5m2) | 8-bit floating-point number (1S,5E,2M) | -| [ float8_e5m2_x4 ](generated/nki.language.float8_e5m2_x4.md#nki.language.float8_e5m2_x4) | 4x packed float8_e5m2 elements, custom data type for nki.isa.nc_matmul_mx on NeuronCore-v4 | +| [ bool\_ ](generated/nki.language.bool_.md#nki.language.bool_) | Boolean (True or False) stored as a byte | +| ---------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------- | +| [ int8 ](generated/nki.language.int8.md#nki.language.int8) | 8-bit signed integer number | +| [ int16 ](generated/nki.language.int16.md#nki.language.int16) | 16-bit signed integer number | +| [ int32 ](generated/nki.language.int32.md#nki.language.int32) | 32-bit signed integer number | +| [ uint8 ](generated/nki.language.uint8.md#nki.language.uint8) | 8-bit unsigned integer number | +| [ uint16 ](generated/nki.language.uint16.md#nki.language.uint16) | 16-bit unsigned integer number | +| [ uint32 ](generated/nki.language.uint32.md#nki.language.uint32) | 32-bit unsigned integer number | +| [ float16 ](generated/nki.language.float16.md#nki.language.float16) | 16-bit floating-point number | +| [ float32 ](generated/nki.language.float32.md#nki.language.float32) | 32-bit floating-point number | +| [ bfloat16 ](generated/nki.language.bfloat16.md#nki.language.bfloat16) | 16-bit floating-point number (1S,8E,7M) | +| [ tfloat32 ](generated/nki.language.tfloat32.md#nki.language.tfloat32) | 32-bit floating-point number (1S,8E,10M) | +| [ float8_e4m3 ](generated/nki.language.float8_e4m3.md#nki.language.float8_e4m3) | 8-bit floating-point number (1S,4E,3M) | +| [ float8_e5m2 ](generated/nki.language.float8_e5m2.md#nki.language.float8_e5m2) | 8-bit floating-point number (1S,5E,2M) | +| [ float8_e5m2_x4 ](generated/nki.language.float8_e5m2_x4.md#nki.language.float8_e5m2_x4) | 4x packed float8_e5m2 elements, custom data type for nki.isa.nc_matmul_mx on NeuronCore-v4 | | [ float8_e4m3fn_x4 ](generated/nki.language.float8_e4m3fn_x4.md#nki.language.float8_e4m3fn_x4) | 4x packed float8_e4m3fn elements, custom data type for nki.isa.nc_matmul_mx on NeuronCore-v4 | | [ float4_e2m1fn_x4 ](generated/nki.language.float4_e2m1fn_x4.md#nki.language.float4_e2m1fn_x4) | 4x packed float4_e2m1fn elements, custom data type for nki.isa.nc_matmul_mx on NeuronCore-v4 | - ## Constants - | [ tile_size ](generated/nki.language.tile_size.md#nki.language.tile_size) | Tile size constants. | -| --- | --- | \ No newline at end of file +| ------------------------------------------------------------------------- | -------------------- | diff --git a/skills/neuron-nki-docs/references/programming/api/nki.md b/skills/neuron-nki-docs/references/programming/api/nki.md index 812731f..869f243 100644 --- a/skills/neuron-nki-docs/references/programming/api/nki.md +++ b/skills/neuron-nki-docs/references/programming/api/nki.md @@ -4,6 +4,5 @@ nki ## Decorators - | [ jit ](generated/nki.jit.md#nki.jit) | This decorator compiles a top-level NKI function to run on NeuronDevices. | -| --- | --- | \ No newline at end of file +| ------------------------------------- | ------------------------------------------------------------------------- | diff --git a/skills/neuron-nki-docs/references/programming/api/view-old-api-ref-pages.md b/skills/neuron-nki-docs/references/programming/api/view-old-api-ref-pages.md index 6bc3c8c..a25aba2 100644 --- a/skills/neuron-nki-docs/references/programming/api/view-old-api-ref-pages.md +++ b/skills/neuron-nki-docs/references/programming/api/view-old-api-ref-pages.md @@ -5,24 +5,23 @@ If you landed on this page, you probably followed a link to an older NKI API ref To view older versions of the NKI API docs, follow these steps: -* Go to the AWS Neuron NKI API docs landing page: [https://awsdocs-neuron.readthedocs-hosted.com/en/latest/nki/api/index.html](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/nki/api/index.html) +- Go to the AWS Neuron NKI API docs landing page: [https://awsdocs-neuron.readthedocs-hosted.com/en/latest/nki/api/index.html](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/nki/api/index.html) -* Click on the version selector dropdown at the bottom-right of the page (it will say latest by default). +- Click on the version selector dropdown at the bottom-right of the page (it will say latest by default). ![../../_images/api-doc-version-selector.jpg](../../_images/api-doc-version-selector.jpg) -* Select the version that corresponds to the older NKI API docs you want to view. The last published version of the NKI API docs before the NKI 0.2.0 (Beta 2) release is version 2.26.1. +- Select the version that corresponds to the older NKI API docs you want to view. The last published version of the NKI API docs before the NKI 0.2.0 (Beta 2) release is version 2.26.1. ## Legacy NKI API reference pages AWS Neuron SDK version: 2.26.1 (last version before NKI 0.2.0 / Beta 2) - -| NKI API namespace | Link | -| --- | --- | -| nki | [nki API reference for version 2.26.1](https://awsdocs-neuron.readthedocs-hosted.com/en/v2.26.1/nki/api/nki.html) | -| nki.language | [nki.language API reference for version 2.26.1](https://awsdocs-neuron.readthedocs-hosted.com/en/v2.26.1/nki/api/nki.language.html) | -| nki.isa | [nki.utils API reference for version 2.26.1](https://awsdocs-neuron.readthedocs-hosted.com/en/v2.26.1/nki/api/nki.isa.html) | -| nki.compiler | [nki.compiler API reference for version 2.26.1](https://awsdocs-neuron.readthedocs-hosted.com/en/v2.26.1/nki/api/nki.compiler.html) | -| NKI API common fields | [NKI API common fields reference for version 2.26.1](https://awsdocs-neuron.readthedocs-hosted.com/en/v2.26.1/nki/api/nki.api.shared.html) | -| NKI API error documentation | [NKI API error documentation for version 2.26.1](https://awsdocs-neuron.readthedocs-hosted.com/en/v2.26.1/nki/api/nki.errors.html) | \ No newline at end of file +| NKI API namespace | Link | +| --------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------ | +| nki | [nki API reference for version 2.26.1](https://awsdocs-neuron.readthedocs-hosted.com/en/v2.26.1/nki/api/nki.html) | +| nki.language | [nki.language API reference for version 2.26.1](https://awsdocs-neuron.readthedocs-hosted.com/en/v2.26.1/nki/api/nki.language.html) | +| nki.isa | [nki.utils API reference for version 2.26.1](https://awsdocs-neuron.readthedocs-hosted.com/en/v2.26.1/nki/api/nki.isa.html) | +| nki.compiler | [nki.compiler API reference for version 2.26.1](https://awsdocs-neuron.readthedocs-hosted.com/en/v2.26.1/nki/api/nki.compiler.html) | +| NKI API common fields | [NKI API common fields reference for version 2.26.1](https://awsdocs-neuron.readthedocs-hosted.com/en/v2.26.1/nki/api/nki.api.shared.html) | +| NKI API error documentation | [NKI API error documentation for version 2.26.1](https://awsdocs-neuron.readthedocs-hosted.com/en/v2.26.1/nki/api/nki.errors.html) | diff --git a/skills/neuron-nki-docs/references/programming/data-representation-overview.md b/skills/neuron-nki-docs/references/programming/data-representation-overview.md index 20dd759..e921d88 100644 --- a/skills/neuron-nki-docs/references/programming/data-representation-overview.md +++ b/skills/neuron-nki-docs/references/programming/data-representation-overview.md @@ -47,14 +47,12 @@ underlying memory type of `hbm`. These tensors are placed in the HBM memory prior to calling the NKI kernel. Intermediate tensors can be allocated using the tensor creation APIs, for instance: - ```python # Allocate 3D tensor on the SBUF x = nl.ndarray((128, 32, 512), dtype=nl.float32, buffer=nl.sbuf) ``` - The above code creates a new 3D tensor on the SBUF memory with shape 128x32x512, and with an element type of 32-bit floats. The physical location of this tensor will be assigned by the NKI compiler, and the total amount of memory used will be: -`8,388,608 = 128 * 32 * 512 * 4` \ No newline at end of file +`8,388,608 = 128 * 32 * 512 * 4` diff --git a/skills/neuron-nki-docs/references/programming/framework_custom_op.md b/skills/neuron-nki-docs/references/programming/framework_custom_op.md index aa799c3..d9da1e9 100644 --- a/skills/neuron-nki-docs/references/programming/framework_custom_op.md +++ b/skills/neuron-nki-docs/references/programming/framework_custom_op.md @@ -17,7 +17,6 @@ This effectively calculates: `a * b * (a + b)`. We define a common NKI kernel for addition. This is a tiled variation of the addition kernel from Quickstart: Build and Run a Kernel. - ```python import os import nki as nki @@ -39,14 +38,14 @@ def nki_tensor_add(a_input, b_input): Returns: c_output: an output tensor """ - # Create output tensor shared between all SPMD instances as + # Create output tensor shared between all SPMD instances as # result tensor (uninitialized) c_output = nl.ndarray(a_input.shape, dtype=a_input.dtype, buffer=nl.shared_hbm) # Extract the dimensions for the a_input shape. M, N = a_input.shape - # Set the tile dimensions, while the TILE_N is not, strictly speaking, limited to + # Set the tile dimensions, while the TILE_N is not, strictly speaking, limited to # 512 for the additiona operation, we stick with this size for simplicity. TILE_M = 128 TILE_N = 512 @@ -91,12 +90,10 @@ def nki_tensor_add(a_input, b_input): return c_output ``` - ### PyTorch We can perform `(a + b) * a * b` using native PyTorch code. - ```python import torch from torch_xla.core import xla_model as xm @@ -111,13 +108,11 @@ out = a * b * c print(out) ``` - Now let’s replace the tensor addition (`c = a + b`) with a NKI kernel. To do this we replace the `+` operator with a call to the NKI kernel caller (`nki_tensor_add`), and everything else works as before. - ```python device = xm.xla_device() a = torch.randn(256, 1024, dtype=torch.float32).to(device) @@ -127,7 +122,6 @@ out = a * b * c print(out) ``` - To understand what happens under the hood when we compile the above code, we can print HLO IR graph generated by XLA by setting the `NEURON_FRAMEWORK_DEBUG` environment variable, which preserves the HLO in @@ -135,14 +129,12 @@ binary form, and the `XLA_SAVE_TENSORS_FILE`, which presents a textual representation of the HLO. For example, you may add the following lines to your code: - ```python import os os.environ['NEURON_FRAMEWORK_DEBUG'] = "1" os.environ["XLA_SAVE_TENSORS_FILE"] = "example1.pbtxt" ``` - A `example1.pbtxt.0` file is then written in your run directory that has the corresponding human-readable HLO IR. @@ -152,7 +144,6 @@ mapped to an HLO `xla::_op_CallImpl` instruction, representing the custo that `xla::_op_CallImpl` is then consumed by the next instruction in line #15 as usual. - ```python [ScheduleSyncTensorsGraph] TensorsGraphInfo: @@ -177,7 +168,6 @@ that `xla::_op_CallImpl` is then consumed by the next instruction in lin ## END_GRAPH ``` - The Neuron compiler replaces the above custom call with the corresponding NKI kernel implementation while optimizing the rest of the compute graph as usual. At the end of the compilation process, a single @@ -189,7 +179,6 @@ including the NKI kernel. For more information about NEFF files, see [Neuron Com We can perform `(a + b) * a * b` using native JAX code. - ```python import jax import jax.numpy as jnp @@ -208,12 +197,10 @@ b = jax.random.normal(seed_b, (256, 1024), dtype=jnp.float32) print(jax_customop_tutorial(a, b)) ``` - Similar to the PyTorch example above, let’s replace the tensor addition `(c = a + b)` with the addition NKI kernel. To do this we replace the `+` operator with a call to the NKI kernel caller (`nki_tensor_add`), and everything else works as before. - ```python import jax import jax.numpy as jnp @@ -231,11 +218,9 @@ b = jax.random.normal(seed_b, (256, 1024), dtype=jnp.float32) print(jax_customop_tutorial(a, b)) ``` - To understand what happens under the hood when we compile the above code, we can print the HLO IR graph by adding the following snippet to your code: - ```python print(jax.jit(jax_customop_tutorial) .lower(a, b) @@ -245,14 +230,12 @@ print(jax.jit(jax_customop_tutorial) ) ``` - Let’s examine the XLA output of this example. In line #8 we can identify that the tensor addition is now mapped to an HLO `custom-call` instruction, similar to PyTorch. The output of that `custom-call` is then consumed by the next instruction in line #9 as usual. - ```python HloModule jit_jax_customop_tutorial, entry_computation_layout={(f32[256,1024]{1,0}, f32[256,1024]{1,0})->(f32[256,1024]{1,0})}, allow_spmd_sharding_propagation_to_parameters={}, allow_spmd_sharding_propagation_to_output={true} @@ -267,7 +250,6 @@ ENTRY %main.12 (Arg_0.1: f32[256,1024], Arg_1.2: f32[256,1024]) -> (f32[256,1024 } ``` - The Neuron compiler replaces the above custom-call with the corresponding NKI kernel implementation while optimizing the rest of the compute graph as usual. At the end of the compilation process, a single @@ -301,7 +283,6 @@ function. The gradients of both input tensors in `y = a + b` are ones, so the `backward()` function propagates the `dy` gradients from the previous backward function. - ```python import torch import torch_xla.core.xla_model as xm @@ -332,7 +313,6 @@ loss.backward() xm.mark_step() ``` - ### JAX We define a `custom_vjp` function `nki_add_func` by using @@ -345,7 +325,6 @@ the gradients through. Finally, to start training, we execute the forward pass by calling `nki_add_func(a, b) * x * y`. To get the gradients, we call `jax.grad` directly with a loss function. - ```python @jax.custom_vjp def nki_add_func(a, b): @@ -370,4 +349,4 @@ def jax_customop_tutorial_and_grad(a, b): return out, *grad c, grad_a, grad_b = jax_customop_tutorial_and_grad(a, b) -``` \ No newline at end of file +``` diff --git a/skills/neuron-nki-docs/references/programming/indexing-overview.md b/skills/neuron-nki-docs/references/programming/indexing-overview.md index 4c31a95..686ca42 100644 --- a/skills/neuron-nki-docs/references/programming/indexing-overview.md +++ b/skills/neuron-nki-docs/references/programming/indexing-overview.md @@ -6,10 +6,9 @@ This topic covers basic tensor indexing and how it applies to developing with th ## Basic Tensor Indexing NKI supports basic indexing of tensors using integers as indexes. For example, -we can index a 3-dimensional tensor with a single integer to get get a *view* +we can index a 3-dimensional tensor with a single integer to get get a _view_ of a portion of the original tensor. - ```python x = nl.ndarray((2, 2, 2), dtype=nl.float32, buffer=nl.shared_hbm) @@ -18,12 +17,10 @@ x = nl.ndarray((2, 2, 2), dtype=nl.float32, buffer=nl.shared_hbm) assert x[1].shape == [2, 2] ``` - NKI also supports creating views from sub-ranges of the original tensor dimension. This is done with the standard Python **slicing** syntax. For example: - ```python x = nl.ndarray((2, 128, 1024), dtype=nl.float32, buffer=nl.shared_hbm) @@ -38,7 +35,6 @@ assert x[1, :, 0:512].shape == [128, 512] assert x[:, 1, 0:2].shape == [2, 2] ``` - When indexing into tensors, NeuronCore offers much more flexible memory access in its on-chip SRAMs along the free dimension. You can use this to efficiently stride the SBUF/PSUM memories at high performance for all NKI APIs that access @@ -65,7 +61,6 @@ below visualizes the input and output tensors. Fig. 16 Tensor split to even and odd columns - ```python import nki import nki.language as nl @@ -122,7 +117,6 @@ if __name__ == "__main__": print(in_tensor, out1_tensor, out2_tensor) ``` - The main concept in this example is that we are using slices to access the even and odd columns of the input tensor. For the partition dimension, we use the slice expression :, which selects all of the rows of the input tensor. For @@ -135,11 +129,11 @@ each step. The odd columns are similar, except we start at index 1. In this example we transpose a tensor along two of its axes. Note, there are two main types of transposition in NKI: -* Transpose between the partition-dimension axis and one of the free-dimension axes, which is achieved via the -[nki.isa.nc_transpose](api/api-nki-isa-tensor.md#nki-isa-nc_transpose) API. +- Transpose between the partition-dimension axis and one of the free-dimension axes, which is achieved via the + [nki.isa.nc_transpose](api/api-nki-isa-tensor.md#nki-isa-nc_transpose) API. -* Transpose between two free-dimension axes, which is achieved via a [nki.isa.dma_copy](api/api-nki-isa-memory.md#nki-isa-dma_copy) API, -with indexing manipulation in the transposed axes to re-arrange the data. +- Transpose between two free-dimension axes, which is achieved via a [nki.isa.dma_copy](api/api-nki-isa-memory.md#nki-isa-dma_copy) API, + with indexing manipulation in the transposed axes to re-arrange the data. In this example, we’ll focus on the second case: consider a three-dimensional input tensor `[P, F1, F2]`, where the `P` axis is mapped @@ -154,7 +148,6 @@ below illustrates the input and output tensor layouts. Fig. 17 Tensor F1:F2 Transpose - ```python import nki import nki.language as nl @@ -221,7 +214,6 @@ def tensor_transpose2D_kernel_(in_tensor, shape2D): return out_tensor ``` - The main concept introduced in this example is a 2D memory access pattern per partition, via additional indices. We copy `in_tile` into `out_tile`, while traversing the memory in different access patterns @@ -248,7 +240,6 @@ below illustrates the input and output tensor layouts. Fig. 18 2D-Pooling Operation (reducing on axes F2 and F4) - ```python import nki import nki.language as nl @@ -307,4 +298,4 @@ if __name__ == "__main__": out_tensor = tensor_maxpool_kernel_(in_tensor, POOL_SIZE) print(in_tensor, out_tensor) # an implicit XLA barrier/mark-step -``` \ No newline at end of file +``` diff --git a/skills/neuron-nki-docs/references/programming/lnc.md b/skills/neuron-nki-docs/references/programming/lnc.md index 2d12a65..9321be6 100644 --- a/skills/neuron-nki-docs/references/programming/lnc.md +++ b/skills/neuron-nki-docs/references/programming/lnc.md @@ -18,11 +18,11 @@ performance of your kernel. NKI gives you a few mechanisms to for using Logical Neuron Cores (LNC). We will look briefly at each of these, specifically we will describe: -* How to launch a kernel on multiple cores +- How to launch a kernel on multiple cores -* How to tell if a kernel is running on multiple cores +- How to tell if a kernel is running on multiple cores -* How to tell which core a kernel is running on +- How to tell which core a kernel is running on ## Launching a kernel on multiple cores @@ -30,18 +30,15 @@ To launch a NKI kernel on multiple cores, you specify the number of cores to use, in square brackets, when calling the kernel. For example, suppose we have a kernel called lnc_test, and we want to launch this kernel on two cores. - ```python # Launch lnc_test on 2 cores lnc_test[2](input) ``` - The bracket syntax must contain only one number, the number of cores to use. If no brackets are given the number of cores defaults to 1. If the number is too large for the current architecture, then you will receive an error. - ```python # Launch lnc_test on 1 core lnc_test(input) @@ -56,7 +53,6 @@ lnc_test[2](input) lnc_test[8](input) ``` - ## Programming for multiple cores When writing a NKI kernel for multiple cores, there are two important APIs that @@ -68,7 +64,6 @@ is running on. If LNC is not being used, this API will return 1. So, we can tell if we are running on multiple cores by inspecting the result of this variable: - ```python @nki.jit def lnc_test(input): @@ -86,7 +81,6 @@ lnc_test(input) lnc_test[2](input) ``` - The program_id API will return the logical core id that the current instance is running on. In the case of LNC=2, this API will return either 0 or 1. When not using LNC, this API will return 0. This API can be used to @@ -97,7 +91,6 @@ compute the reciprocal of all of the elements of this tensor. We can write a kernel function that is LNC-aware and can make use of extra cores when available. - ```python def lnc_test(input): # Check the first dimension is 2 for this example @@ -125,7 +118,6 @@ def lnc_test(input): return output ``` - The code above has two cases, one for when we are not using LNC (num_programs returns 1), and one for when we are using LNC=2 (num_programs returns 2). In the non-LNC case, there is a for loop that @@ -146,4 +138,4 @@ This means that each core is executing the same basic control flow as the other cores. Most of the time, this requirement will be automatically satisfied by the NKI compiler. However, if you use dynamic control flow, and this control-flow is different on the different cores, then the behavior is -undefined, and you will likely receive an error at runtime. \ No newline at end of file +undefined, and you will likely receive an error at runtime. diff --git a/skills/neuron-nki-docs/references/programming/memory-hierarchy-overview.md b/skills/neuron-nki-docs/references/programming/memory-hierarchy-overview.md index e5107e4..d0808ae 100644 --- a/skills/neuron-nki-docs/references/programming/memory-hierarchy-overview.md +++ b/skills/neuron-nki-docs/references/programming/memory-hierarchy-overview.md @@ -8,13 +8,12 @@ for use in your Machine Leaning models. ## Memory hierarchy The diagram in [Fig. 19](#nki-fig-pm-memory), below, shows the four-level memory hierarchy available to a single NeuronCore. The latency -ranges provided in the figure are approximate and are intended to calibrate the programmer’s mental model (see [NeuronDevice Architecture Guide](../architecture/trainium_inferentia2_arch.md) for the exact values). Memories closer to the top of the figure are the closer to the compute engines; therefore, they are designed to provide the highest bandwidth and lowest latency. However, the faster memories also have smaller capacities compared to memories near the bottom. This set of memories is the *Memory Hierarchy* for the Trainium devices. +ranges provided in the figure are approximate and are intended to calibrate the programmer’s mental model (see [NeuronDevice Architecture Guide](../architecture/trainium_inferentia2_arch.md) for the exact values). Memories closer to the top of the figure are the closer to the compute engines; therefore, they are designed to provide the highest bandwidth and lowest latency. However, the faster memories also have smaller capacities compared to memories near the bottom. This set of memories is the _Memory Hierarchy_ for the Trainium devices. Unlike memory hierarchies for traditional processors (such as CPUs and GPUs), all of the memories available to a NeuronCore are software-managed. This means the contents of the memories are managed either directly by the programmer, or by the Neuron SDK tool chain, rather than being managed by the hardware. In other words, NeuronCore does not have a hardware cache system that performs data movement across memories in a way that is opaque to the program. All memory movement is explicit in the program itself. These explicit memory movements may be specified by writing a NKI kernel, or they may be computed by the Neuron Graph Compiler as part of the optimization process. In the following section we will discuss each memory in turn. - > **Figure: pm memory** > > A pyramid-shaped memory hierarchy diagram showing the four levels of memory in the Neuron system, from fastest/smallest (PSUM) at top to slowest/largest (Host CPU DRAM) at bottom, with capacity and bandwidth specifications for each level. @@ -22,6 +21,7 @@ In the following section we will discuss each memory in turn. > This diagram illustrates the complete memory hierarchy for Neuron-based systems, organized as a pyramid with the fastest, smallest memory at the top and progressively larger, slower memory toward the bottom. Color coding and arrows indicate data flow patterns. > > **Level 1 - PSUM (Top, Orange/Yellow):** +> > - Capacity: ~2 MB > - Bandwidth: ~10 TB/sec > - Purpose: Partial sum accumulator for matrix multiplication results @@ -31,13 +31,15 @@ In the following section we will discuss each memory in turn. > - Classification: Memory within NeuronCore (on-chip) > > **Level 2 - SBUF (Yellow):** -> - Capacity: ~25 MB +> +> - Capacity: ~25 MB > - Bandwidth: ~10 TB/sec > - Purpose: State Buffer for operand staging > - Classification: Memory within NeuronCore (on-chip) > - Both PSUM and SBUF are bracketed as "Memory within NeuronCore (on-chip)" > > **Level 3 - Device Memory HBM (Green):** +> > - Capacity: ~50 GB > - Bandwidth: ~0.5 TB/sec per NC (NeuronCore) > - Purpose: High Bandwidth Memory for device-level storage @@ -47,6 +49,7 @@ In the following section we will discuss each memory in turn. > - Classification: Memory within Neuron Device > > **Level 4 - Host CPU Memory DRAM (Bottom, Blue):** +> > - Capacity: ~1 TB > - Bandwidth: ~16 GB/sec > - Purpose: System memory for host CPU @@ -55,16 +58,19 @@ In the following section we will discuss each memory in turn. > - Red "End compute graph" arrow pointing down > > **Right Side Annotations:** +> > - "Memory within NeuronCore (on-chip)" - brackets PSUM and SBUF > - "Memory within Neuron Device" - brackets HBM > - Implicit: Host memory is outside Neuron Device > > **Key Bandwidth Insights:** +> > - 625x bandwidth difference between on-chip SBUF (~10 TB/s) and HBM (~16 GB/s effective) > - On-chip memory is precious but extremely fast > - HBM provides large capacity but requires careful data staging > > **Key Elements:** +> > - **PSUM (~2 MB, ~10 TB/s)**: Fastest, for matmul accumulation > - **SBUF (~25 MB, ~10 TB/s)**: On-chip operand storage > - **HBM (~50 GB, ~0.5 TB/s)**: Device memory, requires DMA @@ -73,13 +79,12 @@ In the following section we will discuss each memory in turn. > - **Red arrows**: Data storing/output flow > - **Pyramid shape**: Visualizes capacity/speed tradeoff - Fig. 19 NeuronCore Memory Hierarchy with Capacity and Bandwidth Ranges ### NeuronCore external memory The two memories at the bottom of the hierarchy, host memory and device memory, -are both considered *external* memory for a NeuronCore. These memories are +are both considered _external_ memory for a NeuronCore. These memories are **linear memory**, where multi-dimensional tensors must be stored in a flattened manner. @@ -101,7 +106,7 @@ the internal memory back to the HBM. ### NeuronCore internal memory The two memories at the top of the hierarchy, SBUF and PSUM, are both -considered *internal* (or *on-chip*) memory for a NeuronCore. Both memories are +considered _internal_ (or _on-chip_) memory for a NeuronCore. Both memories are **two-dimensional** memory, organized in **128 partitions**. The partitions size of PSUM is typically much smaller than SBUF, and PSUM/SBUF partition sizes vary with NeuronCore generations. @@ -128,6 +133,5 @@ to evict MatMult results back to SBUF as soon as possible. > **Note** > > Note -> -> -> To optimize kernel performance, it is good practice for NKI programmers to be mindful of SBUF and PSUM usage through careful [tiling](tiling-overview.md#nki-about-tiling) and loop fusion. If the total size of the live data being used by a NKI kernel overflows the capacity of any on-chip memory, the Neuron compiler will insert the necessary spills or refills between that memory and the next-tier memory in the hierarchy. \ No newline at end of file +> +> To optimize kernel performance, it is good practice for NKI programmers to be mindful of SBUF and PSUM usage through careful [tiling](tiling-overview.md#nki-about-tiling) and loop fusion. If the total size of the live data being used by a NKI kernel overflows the capacity of any on-chip memory, the Neuron compiler will insert the necessary spills or refills between that memory and the next-tier memory in the hierarchy. diff --git a/skills/neuron-nki-docs/references/programming/nki-aps.md b/skills/neuron-nki-docs/references/programming/nki-aps.md index 5ebc10d..170e051 100644 --- a/skills/neuron-nki-docs/references/programming/nki-aps.md +++ b/skills/neuron-nki-docs/references/programming/nki-aps.md @@ -33,7 +33,6 @@ how the NKI API abstracts this information. The NKI API for access pattern is a direct reflection of the hardware capability. The `nl.ndarray` has an `ap` method. - ```python def ap(self, pattern: List[Tuple[int, int]], offset: Optional[int] = 0, @@ -44,24 +43,23 @@ def ap(self, pattern: List[Tuple[int, int]], pass ``` - The parameters have the following definitions: -* `pattern`: A list of two-element tuples, each tuple describes the access on one dimension. The first element represents the element stepping and the second element represents the number of elements in each dimension. This tuple is referred to as `[step, num]` going forward. +- `pattern`: A list of two-element tuples, each tuple describes the access on one dimension. The first element represents the element stepping and the second element represents the number of elements in each dimension. This tuple is referred to as `[step, num]` going forward. The shape of a pattern is the collection of num. For example, given pattern `[[w_step, w_num], [z_step, z_num], [y_step, y_num], [x_step, x_num]]`, the shape is `[w_num, z_num, y_num, x_num]`. -* It is worth mentioning that the order of the pattern specified here is in the opposite order to what is actually accepted by the hardware. Therefore, the order of the tuples shown on the profiler will be to the opposite order of what is specified here. +- It is worth mentioning that the order of the pattern specified here is in the opposite order to what is actually accepted by the hardware. Therefore, the order of the tuples shown on the profiler will be to the opposite order of what is specified here. -* `offset`: The offset to start the access in terms of number of elements from the beginning of the tensor. The default value is 0. +- `offset`: The offset to start the access in terms of number of elements from the beginning of the tensor. The default value is 0. -* `scalar_offset`: An SBUF memory location that specifies the location to start the access in terms of number of elements on the `indirect_dim` of the access pattern. At most one of the `scalar_offset` and `vector_offset` can be specified. +- `scalar_offset`: An SBUF memory location that specifies the location to start the access in terms of number of elements on the `indirect_dim` of the access pattern. At most one of the `scalar_offset` and `vector_offset` can be specified. -* `vector_offset`: An SBUF memory location that specifies the location to start the access in terms of number of elements from the beginning of the indirect dimension specified by `indirect_dim`. At most one of the `scalar_offset` and `vector_offset` can be specified. +- `vector_offset`: An SBUF memory location that specifies the location to start the access in terms of number of elements from the beginning of the indirect dimension specified by `indirect_dim`. At most one of the `scalar_offset` and `vector_offset` can be specified. -* `indirect_dim`: The indirect dimension on which to apply `scalar_offset` and `vector_offset`. +- `indirect_dim`: The indirect dimension on which to apply `scalar_offset` and `vector_offset`. -* `dtype`: The data type of the access pattern. The default value is the `dtype` of the tensor being accessed. +- `dtype`: The data type of the access pattern. The default value is the `dtype` of the tensor being accessed. ## Semantics of the Access Pattern @@ -76,7 +74,6 @@ Given a tensor, the Access Pattern conceptually flattens the tensor to 1d, and then uses a loop to fetch elements from the tensor to construct a view. Consider the following NKI code: - ```python t = nl.ndarray((p_count, N), dtype=nl.float32, buffer=nl.sbuf) access = t.ap( @@ -85,10 +82,8 @@ access = t.ap( offset) ``` - The above represents the following access on the tensor `t`, written below in pseudo-code. - ```python access = nl.ndarray((p_size, z_num, y_num, x_num), dtype=nl.float32, buffer=nl.sbuf) for w in range(p_size): @@ -100,19 +95,17 @@ for w in range(p_size): + (y * y_step) + (x * x_step)] ``` - The access pattern has the following properties: 1. Recall from the hardware capability, the access pattern in each partition -must be identical. Therefore, the step of the first tuple in the AP must be -equal to the number of elements in the free dimension of the tensor. + must be identical. Therefore, the step of the first tuple in the AP must be + equal to the number of elements in the free dimension of the tensor. 2. The shape of the result view is always the same as the shape of the pattern. Note that calling `.ap` on a tensor does not do any computation directly. It describes how to get data. The engines will consume data when the AP is passed into a `nki.isa` instruction. - ```python src = nl.ndarray((16, 32), dtype=nl.float32, buffer=nl.sbuf) dst = nl.ndarray((16, 32), dtype=nl.float32, buffer=nl.sbuf) @@ -123,13 +116,11 @@ dst_access = dst.ap([32, 16], [1, 32]) # no computation happens nisa.dma_copy(dst_access, src_access) ``` - ## A Concrete Example Given a tensor `t` of size (16P, 16F), to iterate all the elements in `t[0:16, 8:16]` the access pattern can be written as: - ```python t = nl.ndarray((16, 16), dtype=nl.float32, buffer=nl.sbuf) access = t.ap(pattern=[[16, 16], [1, 8]], offset=8) @@ -145,8 +136,6 @@ for w in range(16): access[w, z] = t_flatten[idx] ``` - - > **Figure: memory access visualization 1** > > A memory access pattern visualization showing a 16x16 element grid where half the columns (8-15) are accessed, demonstrating a strided memory access pattern with offset and step parameters used in NKI DMA operations. @@ -154,21 +143,25 @@ for w in range(16): > This diagram visualizes how NKI accesses memory elements in a structured pattern, useful for understanding DMA transfer patterns and memory layout optimization. > > **Title and Layout:** +> > - Title: "Memory Access Visualization" > - Grid: 16 rows (labeled 0-15) x 16 columns (labeled 0-15) > - Total grid size: 256 elements > > **Memory Access Pattern:** > The grid shows two distinct regions: +> > 1. **Unaccessed region (columns 0-7)**: White/empty squares indicating elements not accessed > 2. **Accessed region (columns 8-15)**: Green squares with white dots indicating accessed elements > > **Pattern Annotations:** +> > - **"Offset=8"** (top, orange text with arrow): Points to column 8, indicating the starting column for access > - **"In each row, read 8 elements, step=1 at a time"** (right side, orange text): Describes the access pattern within each row > - **"Jump 16 elements at a time, repeat 16 times"** (left side, orange text with curved arrow): Describes how the pattern moves between rows > > **Access Pattern Details:** +> > - Starts at column 8 (offset = 8) > - Reads 8 consecutive elements per row (columns 8-15) > - Steps by 1 within each row @@ -176,16 +169,19 @@ for w in range(16): > - Repeats for all 16 rows > > **Legend (Bottom):** +> > - Green circle with dot: "Accessed" - elements that are read/written > - "N Overlapping" (orange): Indicates when elements are accessed multiple times > - White square: "Unaccessed" - elements not accessed > > **Statistics (Bottom):** +> > - Total elements: 128 (the accessed portion) > - Index range: 8 - 255 > - Memory region: 16 x 16 = 256 elements > > **Key Elements:** +> > - **Offset=8**: Starting position for memory access > - **8 elements per row**: Width of access pattern > - **16 row stride**: Distance between consecutive row accesses @@ -193,7 +189,6 @@ for w in range(16): > - **Contiguous column access**: Elements 8-15 in each row accessed sequentially > - **Strided row access**: Rows accessed with 16-element stride - ## Restriction on SBUF/PSUM Tensors For SBUF/PSUM tensors, the first tuple must always be the access for the @@ -205,7 +200,6 @@ the leading dimension must be `f_dim0 * f_dim1`. The following example is not allowed because it reads every other partition. - ```python t = nl.ndarray((16, 32), dtype=nl.float32, buffer=nl.sbuf) @@ -213,8 +207,6 @@ t = nl.ndarray((16, 32), dtype=nl.float32, buffer=nl.sbuf) t.ap(pattern=[[64, 8], [1, 32]], offset=0) ``` - - > **Figure: memory access visualization 2** > > A memory access pattern visualization showing a 32x16 element grid with alternating accessed and unaccessed rows, demonstrating a strided row access pattern where every other row is accessed completely. @@ -222,12 +214,14 @@ t.ap(pattern=[[64, 8], [1, 32]], offset=0) > This diagram visualizes a memory access pattern that accesses alternating rows in a larger memory region, common in certain tensor operations where data is interleaved or when accessing every other row. > > **Title and Layout:** +> > - Title: "Memory Access Visualization" > - Grid: 16 rows (labeled 0-15) x 32 columns (labeled 0-31) > - Total grid size: 512 elements > > **Memory Access Pattern:** > The grid displays an alternating row access pattern: +> > - **Row 0**: Fully accessed (all 32 elements - green with white dots) > - **Row 1**: Unaccessed (all 32 elements - white/empty) > - **Row 2**: Fully accessed @@ -246,23 +240,27 @@ t.ap(pattern=[[64, 8], [1, 32]], offset=0) > - **Row 15**: Unaccessed > > **Pattern Characteristics:** +> > - Every even-numbered row (0, 2, 4, 6, 8, 10, 12, 14) is fully accessed > - Every odd-numbered row (1, 3, 5, 7, 9, 11, 13, 15) is completely skipped > - All 32 columns are accessed within each accessed row > - Creates a "striped" pattern in the visualization > > **Legend (Bottom):** +> > - Green circle with dot: "Accessed" - elements that are read/written > - "N Overlapping" (orange): Indicates multiple access (not present in this pattern) > - White square: "Unaccessed" - elements not accessed > > **Statistics (Bottom):** +> > - Total elements: 256 (half of the 512-element region) > - Index range: 0 - 479 > - Memory region: 32 x 16 = 512 elements > - Out of bounds: 0 indices > > **Key Elements:** +> > - **Alternating row pattern**: Every other row accessed > - **32 columns per row**: Full row width accessed when row is active > - **8 accessed rows**: Rows 0, 2, 4, 6, 8, 10, 12, 14 @@ -270,20 +268,17 @@ t.ap(pattern=[[64, 8], [1, 32]], offset=0) > - **256 total accessed**: Half of total elements > - **Row stride of 2**: Access jumps every 2 rows - ## Restriction on Nested Indexing The `.ap` method is only allowed on `nl.ndarray` and cannot be called on a tile produced by it. For example, the following would result in an error. - ```python t = nl.ndarray((128, 256), dtype=nl.float32, buffer=nl.sbuf) t.ap(pattern=[[256, 128],[1, 256]], offset=0).ap(pattern=[[256, 64], [1, 64]], offset=0) ^-- cannot specify an access pattern on an already indexed tensor ``` - ## Reinterpret Cast with `ap` The `dtype` parameter can be used for reinterpret casting the tensor. @@ -291,11 +286,10 @@ Since both the pattern and the offset are in terms of number of elements, not bytes, the count must be computed accordingly. See the following example of reinterpret cast from `INT32` to `BF16`. - ```python t = nl.ndarray((128, 256), dtype=nl.int32, buffer=nl.sbuf) cast_to_bf16 = t.ap(pattern=[ [512, 128], [1, 512] ], # notice the number of elements is doubled due to dtype size change offset = 0) # cast_to_bf16 has shape (128, 512) -``` \ No newline at end of file +``` diff --git a/skills/neuron-nki-docs/references/programming/nki-compiler.md b/skills/neuron-nki-docs/references/programming/nki-compiler.md index 3ec0482..88794bf 100644 --- a/skills/neuron-nki-docs/references/programming/nki-compiler.md +++ b/skills/neuron-nki-docs/references/programming/nki-compiler.md @@ -10,6 +10,7 @@ The NKI language allows kernel writers to have direct, fine grained control over The diagram below shows the detailed compilation flow inside the Neuron compilers and how they work together to build the overall binary that is executable on Neuron hardware. The NKI Compiler first parses the kernel code into an AST representation for semantic analysis. It then performs a small number of middle end and back end transformations on the AST, optimizing resource allocations and instruction scheduling, producing optimized NKI IR that gets integrated back into the overall model. ! + > **Figure: nki compiler 1** > > A detailed flowchart diagram showing the Neuron Compiler architecture with numbered steps illustrating two parallel compilation paths: the NKI PyTorch/JAX flow and the NKI bare-metal flow. @@ -31,6 +32,7 @@ The diagram below shows the detailed compilation flow inside the Neuron compiler > A legend on the right indicates green arrows represent "NKI PyTorch/JAX flow" and blue arrows represent "NKI Bare-metal flow (No Framework)". > > **Key Elements:** +> > - **Step 1**: NKI Kernels conversion to Neuron IR (both paths) > - **Step 2**: Entry into NKI Compiler / Optimized NKI IR output > - **Step 3**: PyTorch/JAX Model + NKI Kernels combination @@ -45,15 +47,12 @@ The diagram below shows the detailed compilation flow inside the Neuron compiler > - **Green flow**: PyTorch/JAX integration path > - **Blue flow**: Bare-metal direct compilation path - > **Note** > > Important -> -> +> > While the NKI language looks and feels like Python, it is not actually Python code. When the Python interpreter encounters a top level function decorated with `@nki.jit`, it invokes the NKI Compiler to handle compilation of that function. - ```python # this is a Python function that calls 'kernel', which is a NKI kernel def a_function(x,y,z): @@ -66,7 +65,6 @@ def kernel(x,y,z): # this is kernel code ``` - Using Python features within NKI kernels that are not supported will result in useful errors from the NKI Compiler indicating that the feature is not a valid NKI feature. Neuron has intentionally constrained the NKI language to be as minimal as possible while serving the needs of building high performance kernels for today’s popular models and will continue to grow and evolve the language over time. ## NKI Compiler Open Source @@ -114,6 +112,6 @@ users a more predictable and performant result. ## Further reading -* [Neuron Graph Compiler](api/index.md) +- [Neuron Graph Compiler](api/index.md) -* [About Neuron Kernel Interface (NKI)](api/index.md) \ No newline at end of file +- [About Neuron Kernel Interface (NKI)](api/index.md) diff --git a/skills/neuron-nki-docs/references/programming/nki-dma-overview.md b/skills/neuron-nki-docs/references/programming/nki-dma-overview.md index 04fc309..36681da 100644 --- a/skills/neuron-nki-docs/references/programming/nki-dma-overview.md +++ b/skills/neuron-nki-docs/references/programming/nki-dma-overview.md @@ -27,6 +27,7 @@ DMA transfers are submitted to DMA queues for the DMA Engines to consume. There When moving data in or out of SBUF, optimal performance is achieved with transfers maximizing the number of partitions with 4KiB or larger per partition. Given 16x DMA engines and 128 SBUF partitions, each DMA engine is typically responsible for moving data for eight SBUF partitions (128 partitions / 16 DMA engines). The figure below visualizes the DMA throughput across different number of bytes per partition (“Free Bytes”), for a fixed partition dimension size of 128: ! + > **Figure: nki dma intro 1** > > A line graph showing DMA throughput in GB/s as a function of bytes per partition, demonstrating performance scaling characteristics when the partition dimension (p_dim) is fixed at 128. @@ -36,6 +37,7 @@ When moving data in or out of SBUF, optimal performance is achieved with transfe > The X-axis represents "Bytes per partition" with data points at powers of 2: 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384, and 32768 bytes. The Y-axis shows "DMA Throughput (GB/s)" with a grid background for reference. > > The curve shows three distinct regions: +> > 1. **Low throughput region** (32-128 bytes): Throughput remains relatively flat and low, indicating overhead-dominated transfers > 2. **Rapid scaling region** (256-2048 bytes): Throughput increases steeply, showing efficient bandwidth utilization as transfer sizes grow > 3. **Saturation region** (4096-32768 bytes): Throughput plateaus near the maximum achievable bandwidth, with diminishing returns for larger transfers @@ -43,6 +45,7 @@ When moving data in or out of SBUF, optimal performance is achieved with transfe > Each data point is marked with a green circle and labeled with the corresponding bytes-per-partition value. The title indicates this test was conducted with p_dim (partition dimension) fixed at 128. > > **Key Elements:** +> > - **Title**: "DMA Throughput varying Bytes per Partition for p_dim = 128" > - **X-axis**: Bytes per partition (32 to 32768, powers of 2) > - **Y-axis**: DMA Throughput (GB/s) @@ -50,7 +53,6 @@ When moving data in or out of SBUF, optimal performance is achieved with transfe > - **Data points**: 11 measurements from 32 to 32768 bytes > - **Key insight**: Throughput saturates around 4096+ bytes per partition - The points on the graph refer to various Free (Dimension) Byte values (that is, bytes per partition). We see that at 4096 free bytes, we are able to nearly saturate DMA bandwidth. Another key consideration for performance is overhead to initiate a DMA transfer. Small, frequent transfers incur significant overhead causing us to be latency bound, while larger transfers help amortize these costs, moving to a more bandwidth bound regime. For optimal performance, it’s important to batch data movements into larger transfers whenever possible. @@ -77,7 +79,6 @@ Here is a diagram with the expected behavior: Here is the kernel to perform the DMA transfer. - ```python import nki.language as nl import nki.isa as nisa @@ -120,12 +121,12 @@ if __name__ == "__main__": print(out_tensor) # an implicit XLA barrier/mark-step ``` - #### Profile The above code runs on a single NeuronCore-v3, in a Trn2 instance. Here we can look at the profile, to validate the expected behavior. Refer to the [Neuron Explorer user guide](api/index.md) for guidance on how to generate a profile. ! + > **Figure: nki dma intro 3** > > A Neuron profiler trace screenshot showing DMA load and store operations with annotated details about operation duration, semaphore IDs, and expected transfer sizes. @@ -135,6 +136,7 @@ The above code runs on a single NeuronCore-v3, in a Trn2 instance. Here we can l > At the top of the trace, two highlighted regions are visible: "Load Operation" on the left side (earlier in time) and "Store Operation" on the right side (later in time), both outlined with red/orange borders for emphasis. > > A detailed popup annotation box appears near the Load Operation, containing key profiling information including: +> > - DMA Operation Duration > - Semaphore ID for the DMA Transfer > - Expected 32KB write in a single transfer @@ -146,6 +148,7 @@ The above code runs on a single NeuronCore-v3, in a Trn2 instance. Here we can l > The dark background with contrasting colored elements (red/orange highlights, blue annotation boxes) makes it easy to identify the key DMA events and their relationships in the execution timeline. > > **Key Elements:** +> > - **Load Operation**: First DMA operation (highlighted on left) > - **Store Operation**: Second DMA operation (highlighted on right) > - **DMA Operation Duration**: Time taken for the transfer @@ -155,7 +158,6 @@ The above code runs on a single NeuronCore-v3, in a Trn2 instance. Here we can l > - **Timeline tracks**: Multiple horizontal tracks showing operation timing > - **Time scale**: Bottom axis showing execution time progression - This is exactly what we expected based on our analysis. From the profile, we can see that the first DMA engine takes 1416 ns to load 32 KiB from HBM to SBUF and also a small 4B semaphore update. Even though the remaining 15 DMA engines do not perform useful data movement, they also perform a small 4B semaphore update writes. This allows the NeuronCore to always monitor a semaphore increment of 16 to signal DMA transfer completion, regardless of the tensor shapes in the transfer. This is good, but this example only uses a single DMA engine. In the next example, we increase partition dimension to increase the number of DMA Engines in use. @@ -174,7 +176,6 @@ Here is a diagram of the expected transfer: #### Example - ```python import nki.language as nl import nki.isa as nisa @@ -220,10 +221,10 @@ if __name__ == "__main__": print(out_tensor) # an implicit XLA barrier/mark-step ``` - #### Profile ! + > **Figure: nki dma intro 5** > > A Neuron profiler trace screenshot showing parallel DMA operations across multiple DMA engines, with annotated details highlighting operation duration, semaphore IDs, and transfer sizes. @@ -231,6 +232,7 @@ if __name__ == "__main__": > This dark-themed profiler interface displays a comprehensive timeline view of DMA engine activity across multiple parallel channels. The interface shows individual track rows for each DMA engine along with other NeuronCore components. > > On the left side, track labels identify each component: +> > - DMA-E64(nc0) through DMA-E79(nc0): 16 DMA engine tracks for NeuronCore 0 > - Scalar(nc0): Scalar engine track > - GpSimd(nc0): GPSIMD engine track @@ -239,11 +241,12 @@ if __name__ == "__main__": > - Pending_DMA_Count(nc0): DMA queue depth track > > Two regions are highlighted with red/orange borders: -> 1. "Load Operation and Semaphore Update" (left side): Shows a staggered pattern of DMA load operations across engines E64-E79, with each engine's operation appearing as a small horizontal bar at slightly different times, creating a diagonal pattern. > +> 1. "Load Operation and Semaphore Update" (left side): Shows a staggered pattern of DMA load operations across engines E64-E79, with each engine's operation appearing as a small horizontal bar at slightly different times, creating a diagonal pattern. > 2. "Store Operation and Semaphore Update" (right side): Shows the corresponding store operations, also with a staggered pattern across the DMA engines. > > A detailed popup annotation near the top provides specific operation details: +> > - Time: 34,051ns - 34,148ns > - Duration: 97ns > - Semaphore ID: 518 (cpSimd[Dynamic]) @@ -260,6 +263,7 @@ if __name__ == "__main__": > The timeline scale at the bottom shows time from approximately 32,667ns to 38,685ns. > > **Key Elements:** +> > - **DMA-E64 to DMA-E79**: 16 parallel DMA engine tracks > - **Load Operations**: Left region showing parallel loads with staggered timing > - **Store Operations**: Right region showing parallel stores @@ -269,12 +273,12 @@ if __name__ == "__main__": > - **Staggered pattern**: Visual representation of parallel DMA scheduling > - **Scalar, GpSimd, State Buffer tracks**: Additional NeuronCore component timing - In the above profile, we can see that all 16 DMA engines are active, as each DMA engine is reading 8 rows from HBM and writing to 8 corresponding partition lanes in SBUF. Similarly, we see the reverse also applies from SBUF, back to HBM. By mousing over an individual DMA operation, we see each DMA engine corresponds to a single 2KiB read (8 rows x 128 elements x 2B), as we expect! Using the same profile from the 128x128 DMA example, lets look at the DMA Trigger and the associated Transfer. You can trace the DMA trigger instruction and the associated DMA transfer via the profiler. This would be useful if you wanted to understand the why a DMA was triggered when, and any preceding dependencies. ! + > **Figure: nki dma intro 6** > > A Neuron profiler trace screenshot showing detailed DMA instruction information with a popup displaying semaphore settings, memory patterns, timing data, and source code location. @@ -282,6 +286,7 @@ Using the same profile from the 128x128 DMA example, lets look at the DMA Trigge > This dark-themed profiler interface displays a timeline trace with multiple component tracks and a detailed information popup for a DMA operation. The view shows system-level profiling data including cumulative HBM throughput. > > The track labels on the left show: +> > - qpSimdDynamic (nc0): GPSIMD dynamic operations > - Scalar(nc0): Scalar engine activity > - GpSimd(nc0): GPSIMD engine activity @@ -296,6 +301,7 @@ Using the same profile from the 128x128 DMA example, lets look at the DMA Trigge > A red annotation arrow points to "DMA Trigger" in the trace, indicating the start of a DMA operation. > > The detailed popup (purple/lavender background) displays comprehensive instruction information: +> > - Name: semaphore=8 sem_increment=16 src_elem_size=256 > - dst_elem_size=256 src_pattern=[256,1][128,1] dst_pattern=[262144,1][128,1] > - src_table_offset_imm=0x8 src_table_index=0 @@ -313,6 +319,7 @@ Using the same profile from the 128x128 DMA example, lets look at the DMA Trigge > - Bir ID: sy00d0:28 > > **Key Elements:** +> > - **DMA Trigger**: Annotation showing DMA operation start point > - **DMA_DIRECT2D opcode**: Direct 2D DMA transfer instruction > - **Source/Destination patterns**: Memory access patterns for the transfer @@ -322,8 +329,8 @@ Using the same profile from the 128x128 DMA example, lets look at the DMA Trigge > - **Duration: 0 ns**: Trigger event (not the full transfer time) > - **HBM Throughput track**: Memory bandwidth visualization - ! + > **Figure: nki dma intro 7** > > A Neuron profiler trace screenshot showing a detailed DMA operation popup with timing, semaphore ID, DMA queue assignment, and transfer size information for a 32 KiB data transfer. @@ -331,6 +338,7 @@ Using the same profile from the 128x128 DMA example, lets look at the DMA Trigge > This dark-themed profiler interface displays a timeline view focused on a specific DMA operation with a detailed information popup. The trace shows various NeuronCore component tracks alongside throughput metrics. > > The track labels on the left include: +> > - qpSimdDynamic (nc0): GPSIMD dynamic operations track > - Scalar(nc0): Scalar engine track > - GpSimd(nc0): GPSIMD engine track (shows a highlighted purple bar indicating the selected operation) @@ -343,6 +351,7 @@ Using the same profile from the 128x128 DMA example, lets look at the DMA Trigge > - DMA Throughput (nc0): Per-core DMA throughput > > The detailed popup (purple/lavender background) shows: +> > - Time: 34,039 ns - 34,507 ns > - Duration: 468 ns > - Semaphore ID: 518 (qpSimdDynamic) - with annotation arrow labeled "Semaphore ID" @@ -358,6 +367,7 @@ Using the same profile from the 128x128 DMA example, lets look at the DMA Trigge > The DMA Throughput track at the bottom shows activity spikes corresponding to the data transfer periods. > > **Key Elements:** +> > - **Duration: 468 ns**: Time taken for the DMA operation > - **Semaphore ID: 518**: Synchronization identifier (qpSimdDynamic) > - **DMA Queue: qpSimdDynamic**: Queue assignment for the transfer @@ -367,5 +377,4 @@ Using the same profile from the 128x128 DMA example, lets look at the DMA Trigge > - **DMA Throughput**: Bottom track showing bandwidth utilization > - **Time range**: 32,836 ns to 38,685 ns visible on timeline - -We can see the first DMA is triggered from qGpSimdDynamic (First screenshot). We can look at GPSimd to see the corresponding trigger (second screenshot). \ No newline at end of file +We can see the first DMA is triggered from qGpSimdDynamic (First screenshot). We can look at GPSimd to see the corresponding trigger (second screenshot). diff --git a/skills/neuron-nki-docs/references/programming/nki-introduction.md b/skills/neuron-nki-docs/references/programming/nki-introduction.md index d4e948a..83c7886 100644 --- a/skills/neuron-nki-docs/references/programming/nki-introduction.md +++ b/skills/neuron-nki-docs/references/programming/nki-introduction.md @@ -5,8 +5,7 @@ Neuron Kernel Interface (NKI) Documentation > **Note** > > NKI Versions -> -> +> > NKI is now Generally Available (GA) with NKI 0.3.0 as the current version. Read more about [NKI versions](../optimization/nki-beta-versions.md). The Neuron Kernel Interface (NKI) is a Python-embedded Domain Specific Language (DSL) that gives developers direct access to Neuron’s Instruction Set Architecture (NISA). NKI provides the ease-of-programming offered by tile-level operations and full access to the Neuron Instruct Set Architecture within a familiar pythonic programming environment. It provides the flexibility to implement architecture-specific optimizations rapidly, at a speed difficult to achieve in higher-level DSLs and frameworks. This has enabled developers to achieve optimal performance across a wide spectrum of machine learning models on Trainium, including Transformers, Mixture-of-Experts, State Space Models, and more. @@ -15,11 +14,11 @@ In addition to directly exposing NISA, NKI provides easy-to-use APIs for control NKI currently supports multiple NeuronDevice generations: -* Trainium/Inferentia2, available on AWS `trn1`, `trn1n` and `inf2` instances +- Trainium/Inferentia2, available on AWS `trn1`, `trn1n` and `inf2` instances -* Trainium2, available on AWS `trn2` instances and UltraServers +- Trainium2, available on AWS `trn2` instances and UltraServers -* Trainium3, available on AWS `trn3` instances and UltraServers +- Trainium3, available on AWS `trn3` instances and UltraServers Explore the comprehensive guides below to learn how to implement and optimize your kernels for AWS Neuron accelerators: @@ -45,4 +44,4 @@ API documentation for the set of pre-built kernels in the NKI Library. [Profiling a NKI Kernel with Neuron Explorer](../optimization/use-neuron-profile.md) -[NKI Performance Optimizations](../optimization/nki_perf_guide.md#nki-perf-guide) \ No newline at end of file +[NKI Performance Optimizations](../optimization/nki_perf_guide.md#nki-perf-guide) diff --git a/skills/neuron-nki-docs/references/programming/nki-language-guide.md b/skills/neuron-nki-docs/references/programming/nki-language-guide.md index 4398648..2ca499e 100644 --- a/skills/neuron-nki-docs/references/programming/nki-language-guide.md +++ b/skills/neuron-nki-docs/references/programming/nki-language-guide.md @@ -5,7 +5,6 @@ The Neuron Kernel Interface (NKI) language is designed for writing kernel functi Let us start by looking at a simple NKI function. - ```python @nki.jit def my_function(x : tensor, y : tensor) -> tensor: @@ -20,24 +19,21 @@ def my_function(x : tensor, y : tensor) -> tensor: return output ``` - The first thing you may notice about this NKI function is that it looks very much like a Python function. In fact, all NKI functions are syntactically valid Python functions. However, it is important to understand that NKI functions are not Python functions: they will be compiled by the NKI compiler and run on the Trainium accelerator. Because of this, not all Python constructs and libraries are supported within a NKI function. The second thing to notice is that NKI has a sequential programming model. This means that the logical order of operations follows the syntactic order of the statements in the function. As you learn more about the Trainium hardware, you will see that the hardware can often do many things at the same time across the different compute engines on the Trainium devices. When we compile NKI functions, we will respect the sequential order of operations written by the programmer. The compiler may reorder operations that have no data dependencies, but this is functionally transparent to NKI programmers. Later you will see how to control which engines operations run on and even how to influence the ordering of operations with no data dependencies for better performance, but all of this is done in the context of the sequential ordering of the code. The third thing to notice about this simple function is that is has a print statement. You may be wondering: When does this print happen? Does the Trainium hardware output a string, where does it go? What about all those different engines we just talked about and the sequential ordering? The answer to these questions reveal a very important aspect of NKI programming. The answer is that the print is evaluated by the compiler at compile time, not at runtime. So, when you compile this NKI function, the NKI compiler will output a string like: - ```text adding tensors of type float16 and shape (128,512) ``` - However, when we run this compiled function on Trainium devices they will not output anything. This is usually what you want. The compiler gives important debugging information during compilation, but when you deploy your function across 1000 Trainium devices, they will not waste any time generating debug output. Note, there is a special print function that does run on the Trainium devices, called device_print, that can be used if this is really what you need, see the API references for more information. -We have just seen that the print statement is evaluated at compile-time, and not at runtime. In fact, most things in NKI programs are evaluated at compile time. In general, calls to nki.isa.* functions will result in on-device operations, and (almost) all other things will be evaluated by the compiler at compile time. We will discuss some exceptions to this rule below, but for now it is generally the case that only the nki.isa.* calls result in run-time operations, and everything else is evaluated by the compiler at compile-time. +We have just seen that the print statement is evaluated at compile-time, and not at runtime. In fact, most things in NKI programs are evaluated at compile time. In general, calls to nki.isa._ functions will result in on-device operations, and (almost) all other things will be evaluated by the compiler at compile time. We will discuss some exceptions to this rule below, but for now it is generally the case that only the nki.isa._ calls result in run-time operations, and everything else is evaluated by the compiler at compile-time. -This leads us to our the last observation about NKI functions. The nki.isa.* APIs are the heart of the matter. These APIs are designed to expose the underlying hardware capabilities in as direct a way as possible. If you write a nki.isa function, then the hardware will execute that operation at that point in the program. The NKI language simply provides a convenient way to specify which ISA operations you want to run on your data. +This leads us to our the last observation about NKI functions. The nki.isa.\* APIs are the heart of the matter. These APIs are designed to expose the underlying hardware capabilities in as direct a way as possible. If you write a nki.isa function, then the hardware will execute that operation at that point in the program. The NKI language simply provides a convenient way to specify which ISA operations you want to run on your data. In the rest of this guide we will focus on the NKI language, starting with the values you can manipulate in a NKI function. We will then cover tensor indexing, control flow, and end with a discussion of class support and interoperation with Python. @@ -45,31 +41,30 @@ In the rest of this guide we will focus on the NKI language, starting with the v The NKI language supports six types of values: -* The special None value +- The special None value -* Boolean values (True and False) +- Boolean values (True and False) -* 32-bit integer values +- 32-bit integer values -* 32-bit IEEE floating-point values +- 32-bit IEEE floating-point values -* String literals +- String literals -* Tensors (on-device tensor memory) +- Tensors (on-device tensor memory) In addition, NKI supports the following container types: -* Tuples of any fixed length +- Tuples of any fixed length -* Lists of arbitrary length +- Lists of arbitrary length -* Dictionaries with string-value keys +- Dictionaries with string-value keys -* Simple user-defined classes +- Simple user-defined classes NKI values and containers are very similar to their Python equivalents. For instance, you can use most of the Python standard list functions, and they work in the same way as in Python. - ```python l = [1,2,3] # create a list with 3 elements l.append(4.1) # append a value to the list @@ -89,10 +84,8 @@ for x in l.reverse(): print(x) ``` - The NKI dictionary type is also similar to the Python version, but with the restriction that the keys must be string values. - ```python d = dict() # create an empty dictionary d['a'] = 1 # set a value in the dictionary @@ -112,28 +105,26 @@ if d.pop('a'): a = d.setdefault('a', default=1) ``` - We will discuss user-defined classes later in the guide. For now, lets take a close look at the most important value in NKI, the tensor. ## Tensor Values The NKI tensor type is a representation of an on-chip tensor. That is, a value of type tensor is really a reference to some region of memory on the Trainium device at runtime. At compile-time, we do not yet know the precise location nor the precise contents of this tensor, and therefore, code evaluated at compile-time will not be able to query the precise location nor the contents. At compile-time we can only query meta-data about the tensor, such as its shape and element type. Each tensor value supports the following meta-data as (read-only) fields: -* t.dtype - The element type of the tensor, e.g. “float16” +- t.dtype - The element type of the tensor, e.g. “float16” -* t.shape - The shape of the tensor, e.g. (128,64,64) +- t.shape - The shape of the tensor, e.g. (128,64,64) -* t.address - The (virtual) address of the tensor (discussed below) +- t.address - The (virtual) address of the tensor (discussed below) -* t.offset - The access pattern offset (discussed below) +- t.offset - The access pattern offset (discussed below) -* t.pattern - The access pattern (discussed below) +- t.pattern - The access pattern (discussed below) -* t.buffer - The memory buffer this tensor lives in (discussed below) +- t.buffer - The memory buffer this tensor lives in (discussed below) The most commonly used fields are dtype and shape. We have already seen an example of using these fields to check that argument tensors are compatible in our simple example. Another common case is using a dimension of a shape to iterate over a tensor: - ```python # assume t is a 3-dimensional tensor, we can iterate over the # 2-D subtensors @@ -141,14 +132,12 @@ for i in range(t.shape[0]): my_function(t[i]) ``` - Note, because the shape is part of the meta-data of the tensor, the expression t.shape[0] is a compile-time constant. Therefore, the bounds of the for-loop are known at compile time. The compiler will unroll this loop into a sequence of calls to my_function, one for each subtensor of t. ## Creating Tensors The easiest way to create tensors is using the nki.language.ndarray API. This function takes a shape, a dtype, and a memory type, and creates a reference to a memory region in the given memory type large enough to hold the tensor. - ```python # A matrix of 128x128 16-bit float values in the SBUF memory t = nl.ndarray((128,128), nl.float16, nl.sbuf) @@ -157,10 +146,8 @@ assert t.dtype == nl.float16 assert t.buffer == nl.sbuf ``` - You can also create a tensor from an existing tensor using the reshape method. The reshape method will create a new reference to the same memory with a different shape. - ```python # create an alternate view of t with shape 128x2x64 u = t.reshape((128,2,64)) @@ -169,14 +156,12 @@ u = t.reshape((128,2,64)) v = t.reshape((128,32)) ``` - When using reshape the new tensor must use the same or less memory than the original tensor. So, the tensor v, defined above, corresponds to one quarter of the original tensor t. ## Creating Tensors (the hard way) The function nl.ndarray is an easy way to create tensors that covers the most common cases. For more precise control, you can also create tensors by first defining a memory region, and then creating a view of the memory region. There are several memory regions you can choose, but we will focus on the SBUF region, the most common case. To create a memory region, we start with an existing memory region and define which part of the existing region we want to use. The special region sbuf refers to the entire device SBUF memory, so we can start with that. Once we have a memory region, we can create a tensor by calling the view method. - ```python # create a memory region in the SBUF of size 128x64 bytes region = sbuf.ptr(size=(128, 64)) @@ -185,29 +170,23 @@ region = sbuf.ptr(size=(128, 64)) t = region.view(nl.float16, (128, 32)) ``` - Note, that the combination of ptr and view is similar to ndarray. In fact, this is what ndarray is, a view of a region that is just large enough to fit the desired tensor. In fact, you can pass a region directly to ndarray if you like, as long as it is big enough to hold the resulting tensor. - ```python # equivalent to region.view above t = nl.ndarray((128,32), nl.float16, buffer=region) ``` - So far, we haven’t done anything that we couldn’t do with ndarray. However, the ptr method has another argument, offset which lets us specify the (relative) offset of the region. Tensors built this way are known as “allocated tensors,” because we have given the compiler some direction about how to allocate the tensors. - ```python # create a tensor at offset 128 bytes from the beginning of the SBUF memory. region = sbuf.ptr(size=(128,64), offset=(0,128)) t = region.view(nl.float16, (128,32)) ``` - Note, the offset is a virtual offset. This will be the location of the tensor relative to the overall memory assigned to your kernel function by the compiler. This is useful if you want to control the relative location of two tensors. For example, to create two tensors that are right next to each other in the SBUF, you could use: - ```python region1 = sbuf.ptr(size=(128,64), offset=(0,0)) region2 = sbuf.ptr(size=(128,64), offset=(0,64)) @@ -216,10 +195,8 @@ t1 = region1.view(nl.float16, (128,32)) t2 = region2.view(nl.float16, (128,32)) ``` - There is actually another way to achieve the same result, but using multiple views of a single region: - ```python region = sbuf.ptr(size=(128,128)) @@ -230,12 +207,10 @@ t1 = region1.view(nl.float16, (128,32)) t2 = region2.view(nl.float16, (128,32)) ``` - In the above, we first create a region large enough to hold both tensors. Then we create two regions inside of the first region which each take up half of the space. Then, we create our two tensors in these regions. The main difference between the first and the second approach is that in the first approach, the two tensors have a fixed address relative to the rest of the memory of the kernel. In the second approach, the two tensors have a fixed address relative to each other, but not to the rest of the memory of the kernel. The region offset may be changed by the compiler, because it is not specified, but the offsets of region1 and region2, within region are fixed. As a final note on creating tensors, you may have noticed that the lower-level creation routines allow you to create two tensors in the same memory region with different shapes, as long as they both fit in the memory. For example: - ```python region2 = region.ptr(size=(128,64)) @@ -244,10 +219,8 @@ t1 = region.view(nl.float16, (128,32)) t2 = region.view(nl.float16, (128,2,16)) ``` - In fact, the tensor reshape method is just a short-hand notation for view: - ```python # this is just a short-hand u = t.reshape(shape) @@ -256,18 +229,16 @@ u = t.reshape(shape) u = t.address.reshape(t.dtype, shape) ``` - This is a common theme with the NKI tensor creation APIs: there are several nice convenience functions available, but everything can be achieved with the more primitive ptr and view methods. ## Tensor Indexing In the previous section we noted that there are six read-only fields you can query on a tensor value. We discussed four of them, but not offset or pattern. These last two fields are related to tensor indexing. Before we talk about these fields, first lets look at the most common way of indexing tensors using integers and slices. -Suppose you have a tensor t with shape 64x64x64 that is in the SBUF memory. The SBUF memory is a two dimensional memory, so the underlying storage for this 3-D tensor is a 2-D region of the SBUF. Recall, in the SBUF, the first dimension is called the partition dimension and the second dimension if called the free dimension. By convention, the first dimension of a tensor always corresponds to the partition dimension, and the remaining dimension are layed out in the free dimension. Note, this is a change from NKI Beta 1 where the partition dimension could be mapped to any dimension of the tensor. The first dimension is always the partition dimension. Therefore, in our example, we have 64 partitions, each with 64*64=4096 elements. +Suppose you have a tensor t with shape 64x64x64 that is in the SBUF memory. The SBUF memory is a two dimensional memory, so the underlying storage for this 3-D tensor is a 2-D region of the SBUF. Recall, in the SBUF, the first dimension is called the partition dimension and the second dimension if called the free dimension. By convention, the first dimension of a tensor always corresponds to the partition dimension, and the remaining dimension are layed out in the free dimension. Note, this is a change from NKI Beta 1 where the partition dimension could be mapped to any dimension of the tensor. The first dimension is always the partition dimension. Therefore, in our example, we have 64 partitions, each with 64\*64=4096 elements. We can refer to specific elements of the tensor using an index expression. - ```python # 10th element in partition 0 u = t[0,0,10] @@ -279,10 +250,8 @@ u = t[0,1,0] u = t[63,63,63] ``` - It is more common to refer to whole sub-tensors rather then single elements, and for this we can use slices. A slice is an expression of the form start:stop:step, which describes a range of elements starting with index start, up to (but not including) index stop, and incrementing by step. If any of start, stop, or step are not specified, defaults will be used. - ```python # All first 64 elements of every partition u = t[0:64, 0, 0:64] @@ -294,10 +263,8 @@ u = t[:, 0, :] u = t[:, :, ::2] ``` - Finally, you can also use the ellipsis (…) to indicate defaults for a range of dimensions. - ```python # the whole tensor t u = t[...] @@ -310,10 +277,8 @@ u = t[:,...] u = t[0,...,:] ``` - Note, when you index into a tensor, the result is another tensor. So, in the examples above, the tensor u also has the normal tensor fields and capabilities. This means you can query the shape of the result, or further index the tensor u. - ```python u = t[0,...] assert u.shape = (64,64) @@ -322,10 +287,8 @@ v = u[0:32, :] assert v.shape = (32, 64) ``` - In addition to querying the shape, you can also query the hardware access pattern that corresponds to the tensor value. For example, the code below will display the access pattern that would be used to query u, which is a sub-tensor of t. - ```python u = t[0,...] @@ -334,23 +297,19 @@ print(u.offset) print(u.pattern) ``` - For advanced use cases, the hardware access pattern can be specified directly. - ```python # Specify HW access pattern directly u = t.ap(offset = 0, pattern = [...]) ``` - For more details on hardware access patterns, see the architecture guide. ## Control Flow NKI supports basic control flow constructs, including if-statements, for-loops over ranges, lists or tuples, and while loops. All of these constructs work similarly their equivalents in Python. For example, the code below uses a simple loop with an nested if statement to process the even and odd elements of a list differently. - ```python def kernel(outputs, inputs): for i in range(len(inputs)): @@ -360,38 +319,30 @@ def kernel(outputs, inputs): nki.isa.reciprocal(dst=outputs[i], data=inputs[i]) ``` - The loop and if-statement above will ultimately be evaluated by the NKI compiler. This means the the ISA instructions will be output to the final executable as a linear sequence. For example, suppose we call kernel with these arguments. - ```python kernel([a,b,c], [x,y,z]) ``` - where a,b,c and x,y,z are tensors. Then, this call is equivalent to the code: - ```python nki.isa.nc_transpose(dst=a, data=x) nki.isa.reciprocal(dst=b, data=y) nki.isa.nc_transpose(dst=c, data=z) ``` - We will see in the next section how to write loops that run on the Trainium hardware. First, let’s look at some more common uses of control flow in NKI kernels. The recommended way to write loops in NKI is using the standard Python `range`: - ```python for i in range(...): ... ``` - NKI also provides `nl.affine_range`, `nl.sequential_range`, and `nl.static_range` as legacy aliases. In NKI 0.3.0, all of these range functions have identical effect — they are all equivalent to `range`. The recommended approach is to simply use `range`. A for-loop can also iterate over a list or tuple, similar to Python. The two loops below both print the numbers 1-3 in sequence. - ```python l = [1,2,3] for x in l: @@ -402,10 +353,8 @@ for x in t: print(x) ``` - Finally, NKI also supports while loops. Again these loops are similar to Python, and will be unrolled by the compiler, just like the for-loops. - ```python # print the numbers 0-9 x = 0 @@ -414,43 +363,37 @@ while x < 10: x += 1 ``` - ## Dynamic Control Flow > **NKI 0.6.0+ note:** The NKI frontend is moving from Parsing to Tracing (Tracing available in > 0.6.0, default in 0.7.0, parser removed in 0.8.0). The `for i in dynamic_range(...)` and bare > `while reg:` forms shown below are **parser-only and removed under tracing** — a runtime register > has no value at trace time. Replace them with the structured constructs `nl.fori_loop(lower, upper, -> body_fun, step=1)` (counted loop with a runtime bound) and `nl.while_loop(init, body_fun)` +body_fun, step=1)` (counted loop with a runtime bound) and `nl.while_loop(init, body_fun)` > (data-dependent loop); both compile on the parser and the tracer. In the previous section we looked at control-flow constructs that are ultimately expanded at compile-time. NKI also supports dynamic control-flow, or control-flow that runs on the device. Dynamic control-flow is not expanded by the compiler, but lowered to equivalent Trainium control-flow instructions. The most basic dynamic loop is a for-loop with static bounds. A dynamic loop with static bounds can be written using the standard for-loop with a dynamic_range hint. - ```python # create a dynamic loop that runs "on chip" for i in dynamic_range(10): process_tensor(t[i]) ``` - The for loop above will lower to a loop on the Trainium device. The loop will execute its body (process_tensor), 10 times and then continue. Because this is a dynamic loop, the loop index, i, will be stored in a hardware register during evaluation. Therefore, the type of i is register in NKI. Register values can be used to index tensors, and passed to nki.isa APIs. We can also use registers to create dynamic loops with dynamic bounds. - ```python count = nki.isa.register_alloc(count_tensor) for i in dynamic_range(count): process_tensor(t[i]) ``` - The loop above uses a register value as the upper bound. This register is initialized with the register_alloc function, which can take a SBUF tensor as an argument. In this case register_alloc will load a value from the SBUf tensor count_tensor and store it in the register count. The for loop will then execute count times. There are four register APIs that can be used to create, and load and store values to and from registers. - ```python # allocate a new register with initial value # either from constant integer, or a SBUF tensor @@ -469,10 +412,8 @@ def register_load(dst: register, src: tensor): ... def register_store(dst: tensor, src: register): ... ``` - Using the APIs above, we can also create dynamic while loops. A dynamic while loop is specified using the standard while-loop with a condition that is a single register value. The NKI compiler will preserve while loops with register conditions, and not unroll them. - ```python # suppose cond is an SBUF tensor, perhaps declared as cond = nl.ndarray((1, 1), buffer=nl.shared_hbm, dtype=np.int32) @@ -489,14 +430,12 @@ while reg: register_load(reg, cond) ``` - The code above uses a 1x1 SBUF tensor called cond to store the condition. We update this tensor in the body of the loop and then use register_load to update the register. When the register reg holds the value 0 the loop will terminate. ## Class Support NKI has basic support for user-defined classes. In NKI all classes are similar to Python data classes. When you declare a class for use in a NKI kernel, the class must inherit from NKIObject and no other classes. This restriction is to ensure the NKI compiler only brings in class definitions that are intended for NKI. A simple NKI class can be declared similar to a Python data class: - ```python @dataclass class C(NKIObject): @@ -513,9 +452,7 @@ c.toggle() print(c.x, c.y) ``` - -The @dataclass decorator is optional; classes with and without the @dataclass decorator will be compiled in the same way by the NKI compiler. The compiler will create the initializer functions __init__ and __post_init__, if they are not provided by the user. For the class above, the default initializers are: - +The @dataclass decorator is optional; classes with and without the @dataclass decorator will be compiled in the same way by the NKI compiler. The compiler will create the initializer functions **init** and **post_init**, if they are not provided by the user. For the class above, the default initializers are: ```python # default if not provided by the user @@ -529,10 +466,8 @@ def __post_init__(self): pass ``` - Classes can be declared in Python and passed as arguments to NKI functions. When a class is used as an argument to a NKI kernel, the NKI kernel will import the definition of the Python class, and convert the Python class instance to a NKI instance using the objects dictionary. Currently, NKI does not look at slots or other object features, only the object dictionary. For example, consider the code shown below. - ```python class A(NKIObject): x : int = 1 @@ -545,10 +480,8 @@ def kernel(a : A): ... kernel(A(1)) ``` - The class A is instantiated in Python as an argument to the kernel function. The NKI compiler will take this object and translate it to an instance of A on the NKI side. Roughly this translation is done by translating the object dictionary, in pseudo-code: - ```python # pseudo-code "copy constuct" A on NKI side def kernel(python_a : A): @@ -558,12 +491,10 @@ def kernel(python_a : A): nki_a.__dict__ = python_a.__dict__ ``` - ## Enumerations In addition to the basic data classes described, NKI also supports basic enumerations. For example, the following can be used in NK kernel functions. - ```python class E(Enum): x = 1 @@ -578,10 +509,8 @@ def f(e : E): f(E.x) ``` - Similar to Python, the NKI compiler will translate the enumration class E to the following: - ```python class E(NKIObject): x = E("x", 1) @@ -593,5 +522,4 @@ class E(NKIObject): self.value = value ``` - -Equality in NKI is structural, so no additional code is needed to replicate the behavior of == and != for objects of type E. No other binary operators on enum values are supported. \ No newline at end of file +Equality in NKI is structural, so no additional code is needed to replicate the behavior of == and != for objects of type E. No other binary operators on enum values are supported. diff --git a/skills/neuron-nki-docs/references/programming/quickstart-implement-run-kernel.md b/skills/neuron-nki-docs/references/programming/quickstart-implement-run-kernel.md index bb5218f..9e7fce8 100644 --- a/skills/neuron-nki-docs/references/programming/quickstart-implement-run-kernel.md +++ b/skills/neuron-nki-docs/references/programming/quickstart-implement-run-kernel.md @@ -5,19 +5,19 @@ The Neuron Kernel Interface (NKI) lets you write low-level kernels that use the When you have completed it, you will have a simple kernel that adds two input tensors and returns the result and a test program in PyTorch or JAX. -* This quickstart is for: Customers new to NKI +- This quickstart is for: Customers new to NKI -* Time to complete: ~10 minutes +- Time to complete: ~10 minutes ## Prerequisites Before you begin, you will need an Inf2, Trn1, Trn2, or Trn3 EC2 instance. -* Your EC2 instance should have the Neuron SDK and NKI library installed on them. If you used the Deep Learning AMI (DLAMI), these will be available by activating a PyTorch or JAX environment with Python’s venv. +- Your EC2 instance should have the Neuron SDK and NKI library installed on them. If you used the Deep Learning AMI (DLAMI), these will be available by activating a PyTorch or JAX environment with Python’s venv. -* You will need a text editor or IDE for editing code. +- You will need a text editor or IDE for editing code. -* A basic familiarity with Python and either PyTorch or JAX will be helpful, though not strictly required. +- A basic familiarity with Python and either PyTorch or JAX will be helpful, though not strictly required. ## Before you start @@ -27,21 +27,18 @@ Make sure you are logged in to your EC2 instance and have activated either a PyT In this step you create the `add_kernel.py` file and add imports for the `nki`, `nki.language`, and `nki.isa` libraries. - ```python import nki import nki.language as nl import nki.isa as nisa ``` - Open your favorite editor or IDE and create the `add_kernel.py` code file, and then add the imports for the NKI libraries. ## Step 2: Create the nki_tensor_add_kernel In this step, you define the `nki_tensor_add_kernel` function. - ```python import os os.environ["NEURON_PLATFORM_TARGET_OVERRIDE"] = "trn1" @@ -53,7 +50,6 @@ def nki_tensor_add_kernel(a_input, b_input): """ ``` - Add the `nki_tensor_add_kernel` function definition above. Make sure you annotate it with the `@nki.jit` decorator as in the example above. ## Step 3: Check input size and shapes @@ -62,7 +58,6 @@ In this step, you add a couple of assertions to check that `a_input` and `b_inpu Add the following assertions to your `nki_tensor_add_kernel` function in `add_kernel.py`. - ```python # Check both input tensor shapes are the same for element-wise operation. assert a_input.shape == b_input.shape @@ -72,7 +67,6 @@ assert a_input.shape == b_input.shape assert a_input.shape[0] <= nl.tile_size.pmax ``` - The first assertion checks that `a_input` and `b_input` have the same shape. The second assertion checks that the inputs will fit in within the tile size of the on-chip memory. If an input is larger than the on-chip tile size, you must tile the input. To keep this example simple we will avoid discussing tiling further in this quick start. ## Step 4: Read input into the on-chip memory @@ -81,7 +75,6 @@ In this step, you will add code to read the inputs from HBM into on-chip memory. The `nki_tensor_add_kernel` function will receive inputs from the HBM memory and must move them into on-chip memory to operate over their values. You first create space in the on-chip memory and then copy the value into on-chip memory for each input. See [Memory Hierarchy](memory-hierarchy-overview.md) for more details on the memory hierarchy. - ```python # Allocate space for the input tensors in SBUF and copy the inputs from HBM # to SBUF with DMA copy. Note: 'sbuf' is a keyword in NKI. @@ -92,14 +85,12 @@ b_tile = sbuf.view(dtype=b_input.dtype, shape=b_input.shape) nisa.dma_copy(dst=b_tile, src=b_input) ``` - The `sbuf.view` function allows you to allocate tensors in SBUF. The `sbuf` keyword is available in any NKI kernel. Here you allocate `a_tile` and `b_tile` and use the `nisa.dma_copy` [instruction](api/api-nki-isa-memory.md#nki-isa-dma_copy) to copy tensors between HBM and SBUF memories. You first supply the destination for the copy, `a_tile` and `b_tile`. Then you provide the source for the copy, `a_input` and `b_input`, as seen in this example. ## Step 5: Add the two tensors In this step, you add code to allocate a destination tensor in SBUF and put the results of adding these two tensor in the new tensor. - ```python # Allocate space for the result and use tensor_tensor to perform # element-wise addition. Note: the first argument of 'tensor_tensor' @@ -108,14 +99,12 @@ c_tile = sbuf.view(dtype=a_input.dtype, shape=a_input.shape) nisa.tensor_tensor(dst=c_tile, data1=a_tile, data2=b_tile, op=nl.add) ``` - As in step 4, you allocate a space for the `c_tile` in SBUF, using `sbuf.view`. Since the shape of the output will be the same shape as the inputs, you can use the `a_input` data type and shape for the allocation. You use the `nisa.tensor_tensor` [instruction](api/api-nki-isa-tensor.md#nki-isa-tensor_tensor) to perform element-wise calculation on two tensors. The first argument of `tensor_tensor` is the destination tensor, `c_tile`, and the sources, `a_tile` and `b_tile`, follow it. You must also provide an op which tells `tensor_tensor` which operation to perform on the inputs. In this case, you use `op=nl.add` to specify addition. ## Step 6: Copy the result to HBM In this step, you will allocate space for the output tensor in HBM and copy the result from SBUF to the new tensor. This is the inverse of what you did with the input, where you copied the inputs from HBM into SBUF. - ```python # Create a tensor in HBM and copy the result into HBM. Note: Simlar to # 'sbuf', 'hbm' is a keyword in NKI. @@ -123,23 +112,19 @@ c_output = hbm.view(dtype=a_input.dtype, shape=a_input.shape) nisa.dma_copy(dst=c_output, src=c_tile) ``` - You use the hbm keyword to create tensors in HBM, similar to how you allocated space in SBUF with the sbuf keyword. You then copy the result in `c_tile` into `c_output`. Remember that `c_output` is the destination and `c_tile` is the source for the `dma_copy` instruction. The copy is needed because outputs, like inputs, need to be in HBM. ## Step 7: Return the output In this step, you will return the result. - ```python # Return kernel output as function output. return c_output ``` - You should now have an `add_kernel.py` file that looks as follows. - ```python import os import nki @@ -184,14 +169,12 @@ def nki_tensor_add_kernel(a_input, b_input): return c_output ``` - ## Step 8: Create a PyTorch or JAX test program In this step, you create a test program as a Python script using either PyTorch or JAX. PyTorchJAXYou can create a file called `test_program.py` with the following content. - ```python import torch from torch_xla.core import xla_model as xm @@ -210,12 +193,10 @@ c = nki_tensor_add_kernel(a, b) print(c) ``` - You use the `xla_device` function to look up device information. You use the device to move tensors created in PyTorch onto the Neuron device. You call the `nki_tensor_add_kernel(a, b)` function to invoke the kernel. The `print` function tells PyTorch to trace the model, causing the kernel to be compiled and run on the Neuron device. You can create a file called `test_program.py` with the following content. - ```python import jax.numpy as jnp from add_kernel import nki_tensor_add_kernel @@ -231,7 +212,6 @@ c = nki_tensor_add_kernel(a, b) print(c) ``` - You create input tensors using the `jax.numpy` library. You call the `nki_tensor_add_kernel function` to invoke the kernel. The `print` function prints the result to the console. All complete! Now, let’s confirm everything works. @@ -240,17 +220,14 @@ All complete! Now, let’s confirm everything works. You can confirm the success of the kernel by running the driver you created in step 8. - ```bash NEURON_PLATFORM_TARGET_OVERRIDE=trn2 python test_program.py ``` - Note that the `NEURON_PLATFORM_TARGET_OVERRIDE` environment variable sets the target architecture. In this example it is set to `trn2` which creates a binary suitable for running on Trn2 machines. For Trn1 / Inf2, specify `trn1`; and for Trn3 specify `trn3`. Note: the `platform_target` parameter on `@nki.jit` was removed in NKI 0.3.0; use this environment variable instead. Whether you used PyTorch or JAX for the driver, you should see the following result. - ```text [[2. 2. 2.] [2. 2. 2.] @@ -258,10 +235,10 @@ Whether you used PyTorch or JAX for the driver, you should see the following res [2. 2. 2.]] ``` - You will also see some additional output depending on whether you used PyTorch or JAX. PyTorchJAX + ```text driver.py:6: DeprecationWarning: Use torch_xla.device instead device = xm.xla_device() @@ -289,8 +266,6 @@ tensor([[2., 2., 2.], nrtucode: internal error: 54 object(s) leaked, improper teardown ``` - - ```text Compiler status PASS The KLR format is located at: final_klir_filepath='/tmp/nki_tensor_add_kernelq3uk7mz0.klir' @@ -310,16 +285,15 @@ Compiler status PASS [2. 2. 2.]] ``` - Congratulations! You have now your first NKI kernel written and running. If you encountered any issues, see the Common issues section below. ## Common issues Uh oh! Did you encounter an error or other issue while working through this quickstart? Here are some commonly encountered issues and how to address them. -* `nki`, `jax`, `torch`, etc. library not found: You may need to activate the PyTorch or JAX environment. +- `nki`, `jax`, `torch`, etc. library not found: You may need to activate the PyTorch or JAX environment. -* No neuron device available: You may not have the `neuron` kernel module loaded. Make sure the `neuron` module is loaded with `sudo modprobe neuron`. +- No neuron device available: You may not have the `neuron` kernel module loaded. Make sure the `neuron` module is loaded with `sudo modprobe neuron`. ## Clean up @@ -329,12 +303,12 @@ When you are finished with this example, you can deactivate your `venv` with `de Now that you’ve completed this quickstart, take your work and dive into other topics that build off of it. -* [NKI Language Guide](nki-language-guide.md) +- [NKI Language Guide](nki-language-guide.md) -* [NKI Tutorials](api/index.md) +- [NKI Tutorials](api/index.md) ## Further reading -* [NKI API Reference Manual](api/index.md) +- [NKI API Reference Manual](api/index.md) -* NKI Developer Guides \ No newline at end of file +- NKI Developer Guides diff --git a/skills/neuron-nki-docs/references/programming/setup-env.md b/skills/neuron-nki-docs/references/programming/setup-env.md index 06b9f27..96a19e1 100644 --- a/skills/neuron-nki-docs/references/programming/setup-env.md +++ b/skills/neuron-nki-docs/references/programming/setup-env.md @@ -11,7 +11,7 @@ Next, you’ll install the Neuron SDK (if not included in the AMI), and you will ## Prerequisites -* You need an AWS login to launch an Inf2 / Trn1 / Trn2 / Trn3 EC2 instance. +- You need an AWS login to launch an Inf2 / Trn1 / Trn2 / Trn3 EC2 instance. ## Instructions @@ -19,9 +19,10 @@ Amazon Linux 2023Ubuntu 22You can set up an environment to use NKI in several wa DLAMIStandard AMIUpgrade -* Launch the instance using the Neuron Deep Learning AMI. +- Launch the instance using the Neuron Deep Learning AMI. ! + > **Figure: nki setup 1** > > An AWS EC2 console screenshot showing the Application and OS Images (AMI) selection page with the Deep Learning AMI Neuron (Amazon Linux 2023) selected for launching a Trainium/Inferentia instance. @@ -29,17 +30,21 @@ DLAMIStandard AMIUpgrade > This screenshot displays the AWS EC2 Launch Instance wizard at the AMI selection step, showing how to choose the Deep Learning AMI for NKI development on Neuron hardware. > > **Page Header:** +> > - Title: "Application and OS Images (Amazon Machine Image)" with "Info" link > - Description text explaining that an AMI contains the operating system, application server, and applications for the instance > > **Search Bar:** +> > - Search field with placeholder "Search our full catalog including 1000s of application and OS images" > > **Tab Navigation:** +> > - "Recents" and "Quick Start" tabs (Quick Start selected) > > **Quick Start OS Options (Icon Grid):** > Seven operating system options displayed as clickable tiles: +> > - **Amazon Linux** (aws logo) - selected/highlighted > - **macOS** (Apple logo) > - **Ubuntu** (Ubuntu logo) @@ -47,11 +52,11 @@ DLAMIStandard AMIUpgrade > - **Red Hat** (Red Hat logo) > - **SUSE Linux** (SUSE logo) > - **Debian** (Debian logo) -> > - "Browse more AMIs" link with description "Including AMIs from AWS, Marketplace and the Community" > > **Amazon Machine Image (AMI) Selection:** > Selected AMI shown in a box: +> > - **Name**: Deep Learning AMI Neuron (Amazon Linux 2023) > - **AMI ID**: ami-00534fb2eb3269cfb (64-bit (x86)) > - **Virtualization**: hvm @@ -59,11 +64,13 @@ DLAMIStandard AMIUpgrade > - **Root device type**: ebs > > **Description Section:** +> > - Release notes link: https://docs.aws.amazon.com/dlami/latest/devguide/appendix-ami-release-notes.html > - Supported EC2 instances: Trn1, Trn1n, Inf2, Trn2 > - User Guide link: https://awsdocs-neuron.readthedocs.com/en/latest/dlami/index.html > > **AMI Details Row:** +> > - Architecture: 64-bit (x86) > - AMI ID: ami-00534fb2eb3269cfb > - Publish Date: 2025-10-30 @@ -71,6 +78,7 @@ DLAMIStandard AMIUpgrade > - Verified provider badge (checkmark) > > **Key Elements:** +> > - **Deep Learning AMI Neuron**: Pre-configured AMI for Neuron development > - **Amazon Linux 2023**: Base operating system > - **Supported instances**: Trn1, Trn1n, Inf2, Trn2 (Trainium and Inferentia) @@ -84,12 +92,11 @@ Once the instance is launched, an environment can be activated with the NKI libr Note: If you are looking to use the Neuron DLAMI in your cloud automation flows, Neuron also supports SSM parameters to easily retrieve the latest DLAMI id. -* Launch the instance using the Amazon Linux 2023 +- Launch the instance using the Amazon Linux 2023 Select the desired region from the EC2 Console and choose “Launch Instance”. In the “Quick Start” tab, select “Amazon Linux”, then in the AL2023 AMI. Select an Inf2 / Trn1 / Trn1n / Trn2 instance type. For more details see the Inf2, Trn1, or Trn2 EC2 pages. Note: You will need to allocate at least 85 GB of storage. -* Install Drivers and Tools - +- Install Drivers and Tools ```bash # Configure Linux for Neuron repository updates @@ -125,10 +132,10 @@ sudo dnf install aws-neuronx-tools-2.* -y export PATH=/opt/aws/neuron/bin:$PATH ``` - -* Set up either a PyTorch or JAX environment to use with NKI +- Set up either a PyTorch or JAX environment to use with NKI PyTorchJAX + ```bash # Install External Dependency sudo dnf install -y libxcrypt-compat @@ -163,8 +170,6 @@ pip install awscli pip install neuronx-cc==2.* torch-neuronx==2.8.* torchvision nki ``` - - ```bash # Install External Dependency sudo dnf install -y libxcrypt-compat @@ -183,27 +188,23 @@ source aws_neuron_venv_jax/bin/activate pip install -U pip ``` - Neuron provides two different ways to install the JAX package. The first is a common package with jax-neuronx packaged together and tested with all the necessary dependencies including jax, jaxlib, libneuronxla, neuronx-cc, and nki. This package can be installed as follows. - ```bash pip install jax-neuronx[stable] --extra-index-url=https://pip.repos.neuron.amazonaws.com ``` - Alternatively, jax, jaxlib, libneuronxla, neuronx-cc, and nki can be installed separately, with jax-neuronx being an optional addition. This version can be installed as follows. - ```bash pip install jax==0.4.38 jaxlib==0.4.38 pip install jax-neuronx libneuronxla neuronx-cc==2.* nki --extra-index-url=https://pip.repos.neuron.amazonaws.com ``` - Upgrading an existing AL2023 install of of the Neuron SDK with NKI can be done with for PyTorch or JAX. PyTorchJAX + ```bash # Install External Dependency sudo dnf install -y libxcrypt-compat @@ -228,8 +229,6 @@ pip install awscli pip install --upgrade neuronx-cc==2.* torch-neuronx==2.8.* torchvision nki ``` - - ```bash # Install External Dependency sudo dnf install -y libxcrypt-compat @@ -242,31 +241,27 @@ pip install wget pip install awscli ``` - JAX upgrade can be done with either the combined jax-neuronx package which is tested to work together as follows. - ```bash pip install --upgrade jax-neuronx[stable] --extra-index-url=https://pip.repos.neuron.amazonaws.com ``` - Alternatively, jax, jaxlib, libneuronxla, neuronx-cc, and nki can be upgraded separately, with jax-neuronx being an optional addition. This version can be installed as follows. - ```bash pip install jax==0.4.38 jaxlib==0.4.38 pip install --upgrade jax-neuronx libneuronxla neuronx-cc==2.* nki --extra-index-url=https://pip.repos.neuron.amazonaws.com ``` - The easiest way to set up an environment to use NKI is by using the Neuron Multi-framework Deep Learning AMI (DLAMI). The DLAMI provides Python virtual environments (using venv) for a variety of frameworks including PyTorch and JAX and is updated with each new release of the Neuron SDK. For customers that prefer to manage the environment directly, it is also possible to start with an standard Ubuntu 22 AMI and install the Neuron SDK and NKI library directly. Customers who already have an environment configured can follow the instructions in the upgrade tab to upgrade to the latest SDK. DLAMIStandard AMIUpgrade -* Launch the instance using the Neuron Deep Learning AMI +- Launch the instance using the Neuron Deep Learning AMI ! + > **Figure: nki setup 2** > > An AWS EC2 console screenshot showing the Application and OS Images (AMI) selection page with the Deep Learning AMI Neuron (Ubuntu 22.04) selected as an alternative option for launching a Trainium/Inferentia instance. @@ -274,17 +269,21 @@ DLAMIStandard AMIUpgrade > This screenshot displays the AWS EC2 Launch Instance wizard at the AMI selection step, showing the Ubuntu-based Deep Learning AMI option for NKI development. > > **Page Header:** +> > - Title: "Application and OS Images (Amazon Machine Image)" with "Info" link > - Description text explaining AMI contents and selection options > > **Search Bar:** +> > - Search field with placeholder "Search our full catalog including 1000s of application and OS images" > > **Tab Navigation:** +> > - "Recents" and "Quick Start" tabs (Quick Start selected) > > **Quick Start OS Options (Icon Grid):** > Seven operating system options displayed as clickable tiles: +> > - **Amazon Linux** (aws logo) > - **macOS** (Apple logo) > - **Ubuntu** (Ubuntu logo) - selected/highlighted with blue border @@ -292,11 +291,11 @@ DLAMIStandard AMIUpgrade > - **Red Hat** (Red Hat logo) > - **SUSE Linux** (SUSE logo) > - **Debian** (Debian logo) -> > - "Browse more AMIs" link with description text > > **Amazon Machine Image (AMI) Selection:** > Selected AMI shown in a box: +> > - **Name**: Deep Learning AMI Neuron (Ubuntu 22.04) > - **AMI ID**: ami-00652e4ca97ea8199 (64-bit (x86)) > - **Virtualization**: hvm @@ -304,11 +303,13 @@ DLAMIStandard AMIUpgrade > - **Root device type**: ebs > > **Description Section:** +> > - Release notes link: https://docs.aws.amazon.com/dlami/latest/devguide/appendix-ami-release-notes.html > - Supported EC2 instances: Trn1, Trn1n, Inf2, Trn2 > - User Guide link: https://awsdocs-neuron.readthedocs.com/en/latest/dlami/index.html > > **AMI Details Row:** +> > - Architecture: 64-bit (x86) > - AMI ID: ami-00652e4ca97ea8199 > - Publish Date: 2025-10-30 @@ -316,6 +317,7 @@ DLAMIStandard AMIUpgrade > - Verified provider badge (checkmark) > > **Key Elements:** +> > - **Deep Learning AMI Neuron (Ubuntu)**: Ubuntu-based alternative to Amazon Linux AMI > - **Ubuntu 22.04**: Base operating system (LTS release) > - **Supported instances**: Trn1, Trn1n, Inf2, Trn2 (same as Amazon Linux version) @@ -330,12 +332,11 @@ Once the instance is launched, an environment can be activated with the NKI libr Note: If you are looking to use the Neuron DLAMI in your cloud automation flows, Neuron also supports SSM parameters to easily retrieve the latest DLAMI id. -* Launch the instance using the Ubuntu 22 +- Launch the instance using the Ubuntu 22 Select the desired region from the EC2 Console and choose “Launch Instance”. In the “Quick Start” tab, select “Ubuntu”, then in the Ubuntu Server 22 AMI. Select an Inf2 / Trn1 / Trn1n / Trn2 instance type. For more details see the Inf2, Trn1, or Trn2 EC2 pages. Note: You will need to allocate at least 50 GB of storage. -* Install Drivers and Tools - +- Install Drivers and Tools ```bash # Configure Linux for Neuron repository updates @@ -368,10 +369,10 @@ sudo apt-get install aws-neuronx-tools=2.* -y export PATH=/opt/aws/neuron/bin:$PATH ``` - -* Set up either a PyTorch or JAX environment to use with NKI +- Set up either a PyTorch or JAX environment to use with NKI PyTorchJAX + ```bash # Install Python venv sudo apt-get install -y python3.10-venv g++ @@ -400,8 +401,6 @@ python -m pip install awscli python -m pip install neuronx-cc==2.* torch-neuronx==2.8.* torchvision nki ``` - - ```bash # Install Python venv sudo apt-get install -y python3.10-venv g++ @@ -414,27 +413,23 @@ source aws_neuron_venv_jax/bin/activate python -m pip install -U pip ``` - Neuron provides two different ways to install the JAX package. The first is a common package with jax-neuronx packaged together and tested with all the necessary dependencies including jax, jaxlib, libneuronxla, neuronx-cc, and nki. This package can be installed as follows. - ```bash pip install jax-neuronx[stable] --extra-index-url=https://pip.repos.neuron.amazonaws.com ``` - Alternatively, jax, jaxlib, libneuronxla, neuronx-cc, and nki can be installed separately, with jax-neuronx being an optional addition. This version can be installed as follows. - ```bash pip install jax==0.4.38 jaxlib==0.4.38 pip install jax-neuronx libneuronxla neuronx-cc==2.* nki --extra-index-url=https://pip.repos.neuron.amazonaws.com ``` - Upgrading an existing Ubuntu 22 install of of the Neuron SDK with NKI can be done with for PyTorch or JAX. PyTorchJAX + ```bash # Install Python venv sudo apt-get install -y python3.10-venv g++ @@ -463,8 +458,6 @@ pip install awscli pip install neuronx-cc==2.* torch-neuronx==2.8.* torchvision nki ``` - - ```bash # Update Python venv sudo apt-get install -y python3.10-venv g++ @@ -474,24 +467,19 @@ source aws_neuron_venv_jax/bin/activate pip install -U pip ``` - Neuron provides two different ways to install the JAX package. The first is a common package with jax-neuronx packaged together and tested with all the necessary dependencies including jax, jaxlib, libneuronxla, neuronx-cc, and nki. This package can be installed as follows. - ```bash pip install --upgrade jax-neuronx[stable] --extra-index-url=https://pip.repos.neuron.amazonaws.com ``` - Alternatively, jax, jaxlib, libneuronxla, neuronx-cc, and nki can be installed separately, with jax-neuronx being an optional addition. This version can be installed as follows. - ```bash pip install jax==0.4.38 jaxlib==0.4.38 pip install --upgrade jax-neuronx libneuronxla neuronx-cc==2.* nki --extra-index-url=https://pip.repos.neuron.amazonaws.com ``` - ## Confirm your work To test the NKI environment is set up and ready to use, a `venv` that contains the `nki` library must be activated. Select the tab below that corresponds to how you installed the Neuron SDK above. @@ -499,55 +487,49 @@ To test the NKI environment is set up and ready to use, a `venv` that contains t Deep Learning AMIStandard AMIThe Deep Learning AMI provides a number of environments for PyTorch, JAX, and other supported ML frameworks. Any of the PyTorch or JAX venvs supplied as a part of the Deep Learning AMI will include the `nki` library. See the Neuron DLAMI overview for the full list of environments. For simplicity, the JAX and PyTorch tabs below each choose the plain JAX and PyTorch venv respectively. PyTorchJAX + ```bash source /opt/aws_neuronx_venv_pytorch_2_9/bin/activate ``` - - ```bash source /opt/aws_neuronx_venv_jax_0_6/bin/activate ``` - The venv created in the setup step above can be activate as follows. PyTorchJAX + ```bash source aws_neuronx_venv_pytorch/bin/activate ``` - - ```bash source aws_neuronx_venv_jax/bin/activate ``` - Once the `venv` is activated, Python can be used to test that the library is available. - ```bash python -c 'import nki' ``` - If the environment is setup correctly, Python should return without reporting any errors. ## Common issues Uh oh! Did you encounter an error or other issue while working through this task? Here are some commonly encountered issues and how to address them. -* Python reports an error trying to import NKI when using a Deep Learning AMI: Make sure a PyTorch or JAX `venv` (provided as part of the Deep Learning AMI) is activated. Your shell prompt should reflect this by starting with `(aws_neuronx_venv_) ...` +- Python reports an error trying to import NKI when using a Deep Learning AMI: Make sure a PyTorch or JAX `venv` (provided as part of the Deep Learning AMI) is activated. Your shell prompt should reflect this by starting with `(aws_neuronx_venv_) ...` -* Python reports an error trying to import NKI in the `venv` created as part of the Standard AMI install: +- Python reports an error trying to import NKI in the `venv` created as part of the Standard AMI install: Make sure the `venv` you created is activated. Your shell prompt should reflect this by starting with `() ...` -* Make sure that the NKI library installation (with `pip`) from the previous instructions succeeded. +- Make sure that the NKI library installation (with `pip`) from the previous instructions succeeded. ## Related information -* [Neuron DLAMI User Guide](api/index.md) +- [Neuron DLAMI User Guide](api/index.md) -* [Neuron Setup Guide](api/index.md) \ No newline at end of file +- [Neuron Setup Guide](api/index.md) diff --git a/skills/neuron-nki-docs/references/programming/tiling-overview.md b/skills/neuron-nki-docs/references/programming/tiling-overview.md index 0a329e4..bf6075a 100644 --- a/skills/neuron-nki-docs/references/programming/tiling-overview.md +++ b/skills/neuron-nki-docs/references/programming/tiling-overview.md @@ -7,7 +7,6 @@ This topic covers tiling and how it applies to developing NKI kernels with the A All NKI APIs operate on tiles. A tile is just a tensor that resides in either the SBUF or PSUM memory with a size and layout that satisfies the constraints of the Neuron instruction set architecture (NeuronCore ISA). Since the SBUF and PSUM memories have 128 partitions, most APIs are limited to tiles with a first dimension (also called the “Partition Dimension”) no larger than 128 elements. So, for example, to compute the reciprocal of a matrix of size 256x256, you will need to split the computation up into (at least) two parts: - ```python # Example how to split 256x256 into tiles with 128 partition dimensions # Assume input and output are tensors of size 256 x 256 @@ -31,10 +30,8 @@ nki.isa.reciprocal(dst=out_tile, data=in_tile) nki.isa.dma_copy(dst=output[P_DIM:256, 0:256], src=out_tile) ``` - In the code above, we allocate two SBUF tensors to store our tiles: one for the input and one for the result. These two tiles are available within the kernel that they are declared in, and will be automatically recycled by the compiler when no longer needed. Then we copy the first 128 rows of our matrix from the input in HBM to the input tile in SBUF, and compute the reciprocal placing the result into the output tile in SBUF. Finally, we copy the result back to the output tensor, in HBM. Of course, this could also be done with a loop, as shown below. - ```python # allocate memory for input and output tiles in_tile = nl.ndarray((P_DIM, 256), dtype=nl.float32, buffer=nl.sbuf) @@ -47,7 +44,6 @@ for i in range(input.shape[0] // P_DIM): nki.isa.dma_copy(dst=output[s, 0:256], src=out_tile) ``` - We will provide more discussion of the indexing in Tensor Indexing. Next, let’s discuss two important considerations when working with tile-based operations in NKI: [data layout](#nki-tile-layout) and [tile size](#nki-tile-size) constraints. ## Layout considerations @@ -66,9 +62,9 @@ To summarize, the partition and free dimensions of a NKI tensor dictate how the The NeuronCore compute engines impose two layout constraints (LC): -* **[Layout Constraint #1]** For matrix multiplication operations, the contraction axis of both input tiles must be mapped to the Partition (P or P_DIM) dimension which is typically 128 for current hardware. +- **[Layout Constraint #1]** For matrix multiplication operations, the contraction axis of both input tiles must be mapped to the Partition (P or P_DIM) dimension which is typically 128 for current hardware. -* **[Layout Constraint #2]** For operations that are not matrix multiplication operations, such as scalar or vector operations, the parallel axis should be mapped to the Partition (`P` or `P_DIM`) dimension. +- **[Layout Constraint #2]** For operations that are not matrix multiplication operations, such as scalar or vector operations, the parallel axis should be mapped to the Partition (`P` or `P_DIM`) dimension. Layout Constraint #1 means that to perform a matrix multiplication of shapes `[M, K]` and `[K, N]` that contracts on K to generate `[M, N]`, Tensor Engine (the engine performing this matmul operation) requires the K dimension to be mapped to the partition dimension in SBUF for both input matrices. Therefore, you need to pass shapes `[K, M]` and `[K, N]` into the [nki.isa.nc_matmul](api/api-nki-isa-tensor.md#nki-isa-nc_matmul) API, as the partition dimension is always the left-most dimension for an input tile to any NKI compute API. @@ -80,15 +76,14 @@ LC#2, on the other hand, is applicable to many instructions supported on Vector, Besides layout constraints, NeuronCore hardware further imposes three tile-size constraints (TC) in NKI: -* **[Tile-Size Constraint#1]** The P dimension size of a tile in both SBUF and PSUM must never exceed `nki.tile_size.pmax == 128`. +- **[Tile-Size Constraint#1]** The P dimension size of a tile in both SBUF and PSUM must never exceed `nki.tile_size.pmax == 128`. -* **[Tile-Size Constraint#2]** For tiles in PSUM, the F dimension size must not exceed `nki.tile_size.psum_fmax == 512`. +- **[Tile-Size Constraint#2]** For tiles in PSUM, the F dimension size must not exceed `nki.tile_size.psum_fmax == 512`. -* **[TileSize Constraint#3]** Matrix multiplication input tiles F dimension size must not exceed `nki.tile_size.gemm_stationary_fmax == 128` on the left-hand side (LHS), or `nki.tile_size.gemm_moving_fmax == 512` on the right-hand side (RHS). +- **[TileSize Constraint#3]** Matrix multiplication input tiles F dimension size must not exceed `nki.tile_size.gemm_stationary_fmax == 128` on the left-hand side (LHS), or `nki.tile_size.gemm_moving_fmax == 512` on the right-hand side (RHS). Programmers are responsible for breaking up your tensors according to these tile-size constraints. For example, below is a simple kernel that applies the exponential function to every element of an input tensor. The kernel expects a shape of `(128, 512)` for both input and output tensors: - ```python import nki.isa as nisa import nki.language as nl @@ -122,10 +117,8 @@ def tensor_kernel(in_tensor): return out_tensor ``` - As expected, the output tensor is an element-wise exponentiation of the input-tensor (a tensor of ones): - ```python tensor([[2.7188, 2.7188, 2.7188, ..., 2.7188, 2.7188, 2.7188], [2.7188, 2.7188, 2.7188, ..., 2.7188, 2.7188, 2.7188], @@ -137,12 +130,10 @@ tensor([[2.7188, 2.7188, 2.7188, ..., 2.7188, 2.7188, 2.7188], device='xla:1', dtype=torch.bfloat16) ``` - Now let’s examine what happens if the input/output tensor shapes do not match the shape of the compute kernel. As an example, we can change the input and output tensor shape from `[128,512]` to `[256,512]`: Since the compute kernel is expecting `(128, 512)` input/output tensors, but we used a `(256, 512)` input/output tensor instead, the bottom half of the output tensor becomes garbage data: - ```python tensor([[2.7188, 2.7188, 2.7188, ..., 2.7188, 2.7188, 2.7188], [2.7188, 2.7188, 2.7188, ..., 2.7188, 2.7188, 2.7188], @@ -154,20 +145,16 @@ tensor([[2.7188, 2.7188, 2.7188, ..., 2.7188, 2.7188, 2.7188], device='xla:1', dtype=torch.bfloat16) ``` - We could try to fix this by changing the tile size inside the compute kernel to `(256, 512)` as well, and see what happens: (**Note**: This violates tile-size constraint #1!) Here, the Neuron Graph Compiler identifies the tile-size constraint violation and fails compilation with the following exception: - ```python Size of partition dimension 256 exceeds architecture limitation of 128. ``` - Now, let’s see how to build a kernel that properly handles `(256, 512)` input/output tensors with a simple loop. We can use the `nki.language.tile_size.pmax` constant defined in NKI as the maximum partition dimension size in a tile. - ```python import nki.isa as nisa import nki.language as nl @@ -206,12 +193,10 @@ def tensor_exp_kernel_(in_tensor): return out_tensor ``` - The `range(2)` call returns `[0, 1]`, just as in standard Python. While the code above does handle `(256, 512)` tensors correctly, it is rather inflexible since it only supports an input shape of `(256, 512)`. Therefore, as a last step, we extend this kernel to handle varying input/output sizes: - ```python import nki.isa as nisa import nki.language as nl @@ -253,9 +238,8 @@ def tensor_exp_kernel_(in_tensor): return out_tensor ``` - The above example handles cases where `in_tensor.shape[0]` is not a multiple of 128 by using the standard Python `min` function to make sure the tensor access is in bounds. ## Further reading -* [Logical Neuron Cores (LNC)](lnc.md#nki-about-lnc) \ No newline at end of file +- [Logical Neuron Cores (LNC)](lnc.md#nki-about-lnc) diff --git a/skills/neuron-nki-docs/references/programming/tutorial-use-a-prebuilt-kernel.md b/skills/neuron-nki-docs/references/programming/tutorial-use-a-prebuilt-kernel.md index f6881e3..be0dbc2 100644 --- a/skills/neuron-nki-docs/references/programming/tutorial-use-a-prebuilt-kernel.md +++ b/skills/neuron-nki-docs/references/programming/tutorial-use-a-prebuilt-kernel.md @@ -5,15 +5,14 @@ This tutorial demonstrates how to leverage pre-built kernels from the NKI Librar To accelerate a compute workload on Trainium with a kernel from NKI Library, you will need the following: -* A reference implementation in PyTorch +- A reference implementation in PyTorch -* A matching kernel in NKI Library with input parameters in the supported range +- A matching kernel in NKI Library with input parameters in the supported range ## Creating a Reference Implementation Here is an example of a reference implementation of the MLP layer in a typical transformer: - ```python class MLPReference(nn.Module): def __init__(self, hidden_size: int, intermediate_size: int, dtype=torch.bfloat16): @@ -28,10 +27,8 @@ class MLPReference(nn.Module): return self.down_proj(gate_output * up_output) ``` - This will serve as a baseline for numerical correctness and optionally for performance. To get the CPU torch reference output, execute the forward pass. The reference output will be used later to confirm that the kernel has been integrated properly. - ```python model = MLPReference(hidden_size, intermediate_size, dtype=torch.bfloat16) model.eval() @@ -40,20 +37,16 @@ with torch.no_grad(): reference_output = model(input_tensor) ``` - ## Using the NKI Library MLP Kernel After it has been confirmed that the reference implementation is working and reasonable (non-zero numbers, non-NaN, etc.), we can try using the MLP kernel from NKI Library: - ```python from nkilib.core.mlp.mlp import fused_mlp_isa_kernel ``` - Use the API documentation to fill out the arguments and to ensure the input parameters are within the supported space for the kernel. Keep in mind that in the newest release of NKI, the SPMD launch grid (for LNC sharding) can be passed simply as an integer, and the output directly stored as an assignment like the following example. Move the output to CPU so that this can be compared. - ```python with torch.no_grad(): nki_output = fused_mlp_isa_kernel[LNC_DEGREE]( @@ -69,12 +62,10 @@ with torch.no_grad(): nki_output_cpu = nki_output.cpu() ``` - ## Comparing Outputs Finally, confirm that the kernel output matches the CPU output. - ```python print(f"\nReference output:\n{reference_output}") print(f"\nNKI output:\n{nki_output_cpu}") @@ -85,10 +76,8 @@ assert nki_output_cpu.shape == reference_output.shape, f"Shape mismatch: {nki_ou torch.testing.assert_close(nki_output_cpu, reference_output, rtol=1e-2, atol=1e-2) ``` - You should see something like this: - ```text Compiler status PASS @@ -124,14 +113,12 @@ tensor([[[0.0000, 0.0039, 0.0000, ..., 0.0029, 0.0020, 0.0020], PASSED ``` - Now this can be used in place of the MLP layer in your torch model definition. ## Complete Example The full script is available below. Make sure to set the environment variable `NEURON_PLATFORM_TARGET_OVERRIDE` before running the script (for example: `NEURON_PLATFORM_TARGET_OVERRIDE=trn2 python mlp_comparison.py`). - ```python # MLP CTE new frontend kernel torch demo script. # Set `NEURON_PLATFORM_TARGET_OVERRIDE` before executing this script (e.g., trn2) @@ -221,4 +208,4 @@ def test_mlp_kernel_accuracy(batch_size, seq_len, hidden_size, intermediate_size if __name__ == "__main__": pytest.main([__file__, "-v", "-s", "-x"]) -``` \ No newline at end of file +``` diff --git a/skills/neuron-nki-docs/references/programming/tutorials/average_pool2d.md b/skills/neuron-nki-docs/references/programming/tutorials/average_pool2d.md index 820628f..4b80445 100644 --- a/skills/neuron-nki-docs/references/programming/tutorials/average_pool2d.md +++ b/skills/neuron-nki-docs/references/programming/tutorials/average_pool2d.md @@ -11,9 +11,9 @@ dimensionality reduction. We implement a 2D AveragePool operation, which is used in many vision neural networks. In doing so, we learn about: -* NKI syntax and programming model. +- NKI syntax and programming model. -* multi-dimensional memory access patterns in NKI. +- multi-dimensional memory access patterns in NKI. The 2D AveragePool operation takes `C x [H,W]` matrices and reduces each matrix along the `H` and `W` @@ -33,7 +33,6 @@ Fig. 26 2D-Pooling Operation (reducing on axes F2 and F4) ### Compute kernel - ```python import nki import nki.language as nl @@ -86,12 +85,10 @@ def tensor_avgpool_kernel(in_tensor, pool_size): return out_tensor ``` - ### Launching kernel and testing correctness To execute the kernel, we prepare tensors `in_tensor` and call `tensor_avgpool_kernel`: - ```python import torch from torch_xla.core import xla_model as xm @@ -119,14 +116,12 @@ if __name__ == "__main__": print("NKI and Torch differ") ``` - ## JAX ### Compute kernel Let’s reuse the same NKI kernel implementation defined for PyTorch above: - ```python import nki import nki.language as nl @@ -179,19 +174,15 @@ def tensor_avgpool_kernel(in_tensor, pool_size): return out_tensor ``` - In order to pass `pool_size` as a compile time constant, we pass `pool_size` as kwargs. - ```python out_nki = tensor_avgpool_kernel(in_array, pool_size=POOL_SIZE) ``` - We write a reference JAX implementation of `AveragePool2D` as JAX does not have a primitive for it. - ```python import jax.numpy as jnp @@ -202,12 +193,10 @@ def jax_average_pool_2D(in_tensor, pool_size): return jnp.nanmean(reshaped, axis=(2, 4)) ``` - ### Launching kernel and testing correctness To execute the kernel, we prepare array `in_array` and invoke the kernel caller function `tensor_avgpool_kernel`: - ```python if __name__ == "__main__": POOL_SIZE = 2 @@ -227,22 +216,19 @@ if __name__ == "__main__": print("NKI and JAX differ") ``` - ## Download All Source Code Click the links to download source code of the kernels and the testing code discussed in this tutorial. -* NKI baremetal implementation: [`average_pool2d_nki_kernels.py`](../../downloads/average_pool2d_nki_kernels.py) +- NKI baremetal implementation: [`average_pool2d_nki_kernels.py`](../../downloads/average_pool2d_nki_kernels.py) -* -PyTorch implementation: [`average_pool2d_torch.py`](../../downloads/average_pool2d_torch.py) +- PyTorch implementation: [`average_pool2d_torch.py`](../../downloads/average_pool2d_torch.py) You must also download [`average_pool2d_nki_kernels.py`](../../downloads/average_pool2d_nki_kernels.py) into the same folder to run this PyTorch script. -* -JAX implementation: [`average_pool2d_jax.py`](../../downloads/average_pool2d_jax.py) +- JAX implementation: [`average_pool2d_jax.py`](../../downloads/average_pool2d_jax.py) You must also download [`average_pool2d_nki_kernels.py`](../../downloads/average_pool2d_nki_kernels.py) into the same folder to run this JAX script. @@ -253,23 +239,18 @@ You can also view the source code in the GitHub repository [nki_samples](https:/ Run NKI baremetal implementation: - ```python python3 average_pool2d_nki_kernels.py ``` - Run PyTorch implementation: - ```python python3 average_pool2d_torch.py ``` - Run JAX implementation: - ```python python3 average_pool2d_jax.py -``` \ No newline at end of file +``` diff --git a/skills/neuron-nki-docs/references/programming/tutorials/fused_mamba.md b/skills/neuron-nki-docs/references/programming/tutorials/fused_mamba.md index 4f710d1..babc58d 100644 --- a/skills/neuron-nki-docs/references/programming/tutorials/fused_mamba.md +++ b/skills/neuron-nki-docs/references/programming/tutorials/fused_mamba.md @@ -15,19 +15,18 @@ to improve the scaling efficiency. In this tutorial, we learn about: -* Mapping different vector operations efficiently to NeuronCore compute engines, such as associative scan and element-wise -operations between tensors +- Mapping different vector operations efficiently to NeuronCore compute engines, such as associative scan and element-wise + operations between tensors -* Leveraging data reuse and tiling to reduce excessive data movement and keep compute engines busy +- Leveraging data reuse and tiling to reduce excessive data movement and keep compute engines busy -* Using [neuron-profile](../../optimization/use-neuron-profile.md) to identify performance bottlenecks and opportunities +- Using [neuron-profile](../../optimization/use-neuron-profile.md) to identify performance bottlenecks and opportunities ## PyTorch Reference Implementation Before jumping to NKI, let’s examine the compute definition of a Mamba-v1 layer using the below PyTorch script (`mamba_torch.py`): - ```python import torch import torch_neuronx @@ -125,36 +124,34 @@ if __name__ == "__main__": print(torch_out) ``` - The input tensor shapes are as follows: -* `delta: [batch, channels, seq_len]` +- `delta: [batch, channels, seq_len]` -* `u: [batch, channels, seq_len]` +- `u: [batch, channels, seq_len]` -* `A: [channels, state_size]` +- `A: [channels, state_size]` -* `B: [batch, state_size, seq_len]` +- `B: [batch, state_size, seq_len]` -* `C: [batch, state_size, seq_len]` +- `C: [batch, state_size, seq_len]` The key model parameters are: -* `batch`: batch size of the model. +- `batch`: batch size of the model. -* `seq_len`: sequence length of the model. +- `seq_len`: sequence length of the model. -* `channels`: hidden size of a token. +- `channels`: hidden size of a token. -* `state_size`: number of model states. +- `state_size`: number of model states. We use `[batch=1, seq_len=512, channels = 256, state_size = 16]` as a simple test case for initial performance evaluation. Running the above Python script will compile the `PyTorch` compute graph using Neuron Compiler and generate a Neuron executable file (NEFF) in the same directory. We can then profile the NEFF on a single NeuronCore using [neuron-profiler](../../optimization/use-neuron-profile.md). -Figure below is a screenshot of the profile. We see this initial PyTorch implementation takes **151.83 ms** to execute *on -device*. - +Figure below is a screenshot of the profile. We see this initial PyTorch implementation takes **151.83 ms** to execute _on +device_. > **Figure: mamba torch ref** > @@ -167,6 +164,7 @@ device*. > Near the bottom, a blue annotation box highlights "Easily saturated DMA throughput" pointing to the DMA Throughput track. At the very bottom, a purple annotation box spans nearly the entire timeline with the label "Execution time on device" showing the total execution time of approximately 133.43 ms. > > **Key Elements:** +> > - **Data movement activities (top)**: Dense activity bars showing heavy DMA operations, annotated in red as "noticeably more busy than compute activities" > - **Compute activities (middle)**: Relatively sparse activity bars labeled in orange, indicating underutilization of compute resources > - **DMA Throughput track**: Shows sustained high throughput, annotated as "Easily saturated DMA throughput" @@ -174,13 +172,11 @@ device*. > - **Timeline scale**: Horizontal axis showing time in milliseconds from 0 to over 133 ms > - **Activity tracks**: Multiple rows showing GPSIMD, Tensor Engine, State Buffer, PSUM, and DMA metrics - Fig. 28 Profile of Mamba PyTorch Implementation Zooming into a portion of the profile, we notice the compute activities on different engines (TensorE/VectorE/ScalarE/GpSimdE) are quite sparse compared to data movement activities (the qSyncIO0 and qVectorSpillReload rows): - > **Figure: mamba torch ref zoomed** > > A zoomed-in Neuron profiler timeline view of the Mamba torch reference implementation showing an 11.1 ms window that reveals sparse compute activity and detailed memory usage patterns. @@ -194,6 +190,7 @@ are quite sparse compared to data movement activities (the qSyncIO0 and qVectorS > The lower section displays memory metrics. State Buffer Usage maintains a consistent yellow line near the top, indicating steady buffer utilization. PSUM Usage shows a colorful stacked area chart with multiple colors (green, blue, pink, orange, yellow) representing different PSUM allocation patterns that fluctuate over time. Pending DMA Count displays as a dark blue stepped line showing queue depth variations. DMA Throughput appears as a red line near the bottom, and Estimated MFU is shown at the very bottom of the view. > > **Key Elements:** +> > - **qSyncIO0**: Continuous yellow bar showing constant I/O sync activity at the top > - **SyncE**: Red vertical marks indicating synchronization operations > - **TensorE/TensorMatrixE**: Sparse activity marks showing underutilized tensor engine @@ -203,7 +200,6 @@ are quite sparse compared to data movement activities (the qSyncIO0 and qVectorS > - **Pending DMA Count**: Blue stepped line tracking DMA operation queue depth > - **Timeline scale**: Spans ~11.1 ms window from 24M to 32M cycles - Fig. 29 Profile of Mamba PyTorch Implementation (Zoomed-in) In this seemingly “memory-bound” execution trace, the achieved DMA throughput is also extremely low, hovering around @@ -219,15 +215,15 @@ the importance of choosing appropriate data layouts to achieve good compute effi Recall we have the following input tensor shapes in device memory: -* `delta: [batch_size, channels, seq_len]` +- `delta: [batch_size, channels, seq_len]` -* `u: [batch_size, channels, seq_len]` +- `u: [batch_size, channels, seq_len]` -* `A: [channels, state_size]` +- `A: [channels, state_size]` -* `B: [batch_size, state_size, seq_len]` +- `B: [batch_size, state_size, seq_len]` -* `C: [batch_size, state_size, seq_len]` +- `C: [batch_size, state_size, seq_len]` In fact, the above tensor layout has been chosen carefully based on the computation done in NeuronCore, which we will discuss in more detail below. @@ -241,15 +237,15 @@ is only a partial parallel axis where results from different states will be accu By extracting `batch` and `state_size` dimensions, we get the following input tensor shapes in device memory: -* `delta_i: [channels, seq_len]` +- `delta_i: [channels, seq_len]` -* `u_i:     [channels, seq_len]` +- `u_i:     [channels, seq_len]` -* `A_i:     [channels]` +- `A_i:     [channels]` -* `B_i:     [seq_len]` +- `B_i:     [seq_len]` -* `C_i:     [seq_len]` +- `C_i:     [seq_len]` Next, let’s visualize the data flow and computation using 2D matrices or vectors step-by-step. @@ -257,7 +253,6 @@ Next, let’s visualize the data flow and computation using 2D matrices or vecto We have the following PyTorch reference code for Step 1: - ```python # delta[batch, channels, seq_len] # A [channels, state_size] @@ -269,14 +264,13 @@ delta[:, :, None, :] * A[None, :, :, None] delta_i[:, :] * A_i[:] ``` - After the above transformation, the multiplication between `delta_i` and `A_i` involves a **broadcasting** across the `seq_len` dimension of `delta_i`. In NKI, free-dimension broadcast can often be folded into the actual computation instruction at no additional performance cost, while partition-dim broadcast often requires a separate instruction on TensorE (see TensorE alternative use case in [Trainium/Inferentia2 Architecture Guide](../../architecture/trainium_inferentia2_arch.md#arch-sec-tensor-engine-alternative-use)). As a result, we have two options for executing Step 1. -**Option 1: Map ``seq_len`` to free dimension.** Element-wise multiplication of `delta_i` and `A_i` on NeuronCore can +**Option 1: Map `seq_len` to free dimension.** Element-wise multiplication of `delta_i` and `A_i` on NeuronCore can be done through nisa.tensor_scalar on either VectorE or ScalarE, which automatically broadcast `A_i` along the free dimension to match the `seq_len` dimension in `A_i`. @@ -286,7 +280,6 @@ of 256 in our initial setup, which exceeds the architectural limitation of `nl.t `delta_i` in the `channels` dimension (tiled dimension denoted as `channels_tiled`) and feed one tile into `nisa.tensor_scalar` at a time. Figure below illustrates the computation done for Option 1. - > **Figure: mamba step1 opt1** > > A tensor operation diagram showing the element-wise multiplication of delta_i and A_i tensors to produce deltaA_i, with dimension annotations showing the mapping to NKI partition (P) and free (F) dimensions. @@ -296,16 +289,19 @@ at a time. Figure below illustrates the computation done for Option 1. > Three tensors are displayed from left to right with mathematical operators between them: > > On the left, a large green rectangular tensor labeled "delta_i" has dimensions annotated as: +> > - Width: seq_len (F dim) - the sequence length mapped to the free dimension > - Height: channels_tiled (P dim) - the tiled channels mapped to the partition dimension > > In the center, a narrow blue vertical tensor labeled "A_i" has dimensions: +> > - Width: 1 (F dim) - single element in free dimension > - Height: channels_tiled (P dim) - matching the partition dimension of delta_i > -> A multiplication symbol (*) appears between delta_i and A_i. +> A multiplication symbol (\*) appears between delta_i and A_i. > > On the right, a purple rectangular tensor labeled "deltaA_i" shows the result with dimensions: +> > - Width: seq_len (F dim) - same as delta_i > - Height: channels_tiled (P dim) - same as input tensors > @@ -314,21 +310,20 @@ at a time. Figure below illustrates the computation done for Option 1. > This layout demonstrates broadcasting: the narrow A_i tensor (with free dimension of 1) is broadcast across the seq_len free dimension of delta_i during the element-wise multiplication. > > **Key Elements:** +> > - **delta_i**: Green input tensor with shape [channels_tiled, seq_len] > - **A_i**: Blue input tensor with shape [channels_tiled, 1] > - **deltaA_i**: Purple output tensor with shape [channels_tiled, seq_len] > - **seq_len (F dim)**: Sequence length dimension mapped to NKI free dimension > - **channels_tiled (P dim)**: Tiled channels mapped to NKI partition dimension -> - **Multiplication (*)**: Element-wise multiplication with broadcasting +> - **Multiplication (\*)**: Element-wise multiplication with broadcasting > - **1 (F dim)**: Single-element free dimension in A_i enabling broadcast - Fig. 30 Step 1, Option 1: nisa.tensor_scalar As an example, the associated NKI code for batch `i_batch`, state `i_state` and tile `i_tile_channels` in `channels` is: - ```python # Input shape in device memory matches the computation layout # Device memory layout: @@ -342,18 +337,16 @@ is: deltaA_i = nisa.tensor_scalar(delta_i, op0=nl.multiply, operand0=A_i) ``` - Note, with this compute layout option, the `delta_i` tensor shape `[channels, seq_len]` in device memory can be loaded into SBUF efficiently with `seq_len` as the free dimension and fed into VectorE/ScalarE for computation. No extra transposes are needed. -**Option 2: Map ``seq_len`` to partition dimension.** Alternatively, if we choose a transposed layout for `delta_i` in +**Option 2: Map `seq_len` to partition dimension.** Alternatively, if we choose a transposed layout for `delta_i` in SBUF for computation, we will need a partition-dimension broadcast of `A_i` using a separate instruction on TensorE (`A_i.broadcast_to(...)`) and then a nisa.tensor_tensor operation between `delta_i` and the broadcast `A_i` on VectorE. As a reminder, we need to tile the `seq_len` dimension to meet the tile size constraint `nl.tile_size.pmax=128`. Figure below illustrates the computation done for Option 2. - > **Figure: mamba step1 opt2** > > A tensor operation diagram showing an alternative layout for Mamba Step 1 where A_i_bcast is transposed with explicit p-dim broadcast annotation, multiplying with delta_i to produce deltaA_i. @@ -363,17 +356,20 @@ to meet the tile size constraint `nl.tile_size.pmax=128`. Figure below illustrat > Three tensors are displayed from left to right with mathematical operators between them: > > On the left, a green rectangular tensor labeled "delta_i" has dimensions annotated as: +> > - Width: channels (F dim) - channels mapped to the free dimension > - Height: seq_len_tiled (P dim) - tiled sequence length mapped to the partition dimension > > In the center, a blue horizontal tensor labeled "A_i_bcast" has a different orientation: +> > - Width: channels (F dim) - matching the free dimension of delta_i > - Height: minimal (single row) > - An annotation "p-dim broadcast" with downward-pointing arrows indicates that this tensor will be broadcast across the partition dimension > -> A multiplication symbol (*) appears between delta_i and A_i_bcast. +> A multiplication symbol (\*) appears between delta_i and A_i_bcast. > > On the right, a purple rectangular tensor labeled "deltaA_i" shows the result with dimensions: +> > - Width: channels (F dim) - same as inputs > - Height: seq_len_tiled (P dim) - same as delta_i > @@ -382,6 +378,7 @@ to meet the tile size constraint `nl.tile_size.pmax=128`. Figure below illustrat > This alternative layout transposes the tensor dimensions compared to Option 1, with sequence length now on the partition dimension and channels on the free dimension, requiring broadcast along the partition dimension. > > **Key Elements:** +> > - **delta_i**: Green input tensor with shape [seq_len_tiled, channels] > - **A_i_bcast**: Blue input tensor requiring p-dim broadcast > - **deltaA_i**: Purple output tensor with shape [seq_len_tiled, channels] @@ -390,12 +387,10 @@ to meet the tile size constraint `nl.tile_size.pmax=128`. Figure below illustrat > - **p-dim broadcast**: Explicit annotation showing broadcast direction along partition dimension > - **Downward arrows**: Visual indication of broadcast direction - Fig. 31 Step 1, Option 2: p-dim broadcast + nisa.tensor_tensor The associated NKI code is as follows: - ```python # Input shape in device memory does NOT match the computation layout # Device memory layout: @@ -410,25 +405,24 @@ A_i_bcast = A_i.broadcast_to((nl.tile_size.pmax, channels)) deltaA_i = nisa.tensor_tensor(delta_i, A_i_bcast, op=ml.multiply) ``` - Assuming the same `delta_i` device memory layout `[channels, seq_len]`, before performing the `nisa.tensor_tensor` instruction, we will need to either: -* Do a regular load of `delta_i` into SBUF using nl.load and an explicit transpose on the loaded `delta_i` using -`nl.transpose` to make `seq_len` lie in the free dimension, or +- Do a regular load of `delta_i` into SBUF using nl.load and an explicit transpose on the loaded `delta_i` using + `nl.transpose` to make `seq_len` lie in the free dimension, or -* Do a transposed load of `delta_i` using nl.load_transpose2d, -which is significantly less efficient in memory bandwidth usage compared to `nl.load` +- Do a transposed load of `delta_i` using nl.load_transpose2d, + which is significantly less efficient in memory bandwidth usage compared to `nl.load` If Option2 was chosen as the compute layout, we would have incentives to define the `delta` input tensor shape as `[seq_len, channels]` in device memory instead. From computation perspectives, Option 2 is less efficient than Option 1 because: -* Option 2 needs an extra TensorE instruction performing partition dimension broadcast. +- Option 2 needs an extra TensorE instruction performing partition dimension broadcast. -* `nisa.tensor_tensor` is 2x slower than `nisa.tensor_scalar` for our input data type FP32 (see API doc for instruction -cost estimates). +- `nisa.tensor_tensor` is 2x slower than `nisa.tensor_scalar` for our input data type FP32 (see API doc for instruction + cost estimates). Therefore, for Step 1 only, Option 1 is the winner compared to Option 2. Let’s continue with the rest of the steps to see if we need to revise this selection due to surrounding operator layout preferences. @@ -437,12 +431,10 @@ if we need to revise this selection due to surrounding operator layout preferenc Step 2 is evaluating exponential on `deltaA_i` from the previous step: - ```python torch.exp(...) ``` - In NeuronCore, evaluating an exponential function on a tensor is considered a scalar operation, which runs on ScalarE. This operation can be invoked through nl.exp or nisa.activation. @@ -453,59 +445,58 @@ no additional cost. This functionality is only exposed in the `nisa.activation` chose Option 2 `nisa.tensor_tensor` in Step 1. Figure below illustrates our new execution plan to combine Step 1 and 2 into `nisa.activation` : - > **Figure: mamba step2** > -> A tensor operation diagram showing the exponential function applied to the product of delta_i and A_i tensors, computing exp(delta_i * A_i) = deltaA_i for the Mamba kernel Step 2. +> A tensor operation diagram showing the exponential function applied to the product of delta_i and A_i tensors, computing exp(delta_i \* A_i) = deltaA_i for the Mamba kernel Step 2. > > This diagram illustrates Step 2 of the Mamba kernel implementation, showing the computation of the exponential of the element-wise product. > > The equation is presented visually with "exp(" on the far left, followed by tensor representations, and closing with ")" before the equals sign: > > On the left (inside the exp function), a green rectangular tensor labeled "delta_i" has dimensions: +> > - Width: seq_len (F dim) - sequence length mapped to free dimension > - Height: channels_tiled (P dim) - tiled channels mapped to partition dimension > -> A multiplication symbol (*) follows. +> A multiplication symbol (\*) follows. > > In the center, a narrow blue vertical tensor labeled "A_i" has dimensions: +> > - Width: 1 (F dim) - single element in free dimension > - Height: channels_tiled (P dim) - matching the partition dimension > > The closing parenthesis of exp() is followed by an equals sign (=). > > On the right, a purple rectangular tensor labeled "deltaA_i" shows the final result with dimensions: +> > - Width: seq_len (F dim) - same as delta_i > - Height: channels_tiled (P dim) - same as input tensors > > This step takes the result from Step 1 (the element-wise multiplication with broadcasting) and applies the exponential function element-wise, which is essential for the Mamba state space model computation. > > **Key Elements:** +> > - **exp()**: Exponential function wrapping the multiplication > - **delta_i**: Green input tensor with shape [channels_tiled, seq_len] > - **A_i**: Blue input tensor with shape [channels_tiled, 1] -> - **deltaA_i**: Purple output tensor containing exp(delta_i * A_i) +> - **deltaA_i**: Purple output tensor containing exp(delta_i \* A_i) > - **seq_len (F dim)**: Sequence length on free dimension > - **channels_tiled (P dim)**: Tiled channels on partition dimension > - **1 (F dim)**: Single-element free dimension enabling broadcast in A_i - Fig. 32 Step 1&2: `nisa.activation` The associated NKI code is as follows: - ```python # Input shape in device memory matches the computation layout deltaA_i = nisa.activation(op=nl.exp, data=delta_i, scale=A_i) ``` - ### Step 3: Element-wise multiplication of delta_i, B_i and u_i. PyTorch reference code for Step 3 is: - ```python # delta[batch, channels, seq_len] # B: [batch, state_size, seq_len] @@ -519,18 +510,16 @@ delta[:, :, None, :] * B[:, None, :, :] * u[:, :, None, :] delta_i[:, :] * B_i[None, :] * u_i[:, :] ``` - This step involves similar compute layout and instruction choices as Step 1: -* `channels` is either partition or free dimension for both `delta_i` and `u_i` +- `channels` is either partition or free dimension for both `delta_i` and `u_i` -* multiplication with `B_i` is either through `nisa.tensor_tensor` or `nisa.tensor_scalar` +- multiplication with `B_i` is either through `nisa.tensor_tensor` or `nisa.tensor_scalar` Since we preferred Step 1 to consume `delta_i` using `channels` as the partition dimension in previous steps, it is wise to follow the same layout choice here for `delta_i` to avoid any transposes. Given this layout choice, the multiplication with `B_i` will have to be a `nisa.tensor_tensor`. Figure below visualizes the computation in Step 3: - > **Figure: mamba step3** > > A two-part tensor operation diagram for Mamba Step 3, showing the computation of deltaU_i from delta_i and u_i (top row), followed by computing deltaBu_i from deltaU_i and B_i with p-dim broadcast (bottom row). @@ -538,18 +527,21 @@ with `B_i` will have to be a `nisa.tensor_tensor`. Figure below visualizes the c > This diagram illustrates Step 3 of the Mamba kernel implementation, consisting of two sequential tensor operations displayed in two rows. > > In the top row (first operation): +> > - A green tensor labeled "delta_i" with dimensions seqlen (F dim) width and channels_tiled (P dim) height -> - Multiplied (*) by a yellow tensor labeled "u_i" with matching dimensions seqlen (F dim) width and channels_tiled (P dim) height +> - Multiplied (\*) by a yellow tensor labeled "u_i" with matching dimensions seqlen (F dim) width and channels_tiled (P dim) height > - Equals (=) a purple tensor labeled "deltaU_i" with the same dimensions seqlen (F dim) by channels_tiled (P dim) > > In the bottom row (second operation): +> > - The purple tensor "deltaU_i" from the previous step with dimensions seqlen (F dim) by channels_tiled (P dim) -> - Multiplied (*) by a blue horizontal tensor labeled "B_i" with width seqlen (F dim) but requiring p-dim broadcast (indicated by downward-pointing arrow) +> - Multiplied (\*) by a blue horizontal tensor labeled "B_i" with width seqlen (F dim) but requiring p-dim broadcast (indicated by downward-pointing arrow) > - Equals (=) a pink/salmon tensor labeled "deltaBu_i" with dimensions seqlen (F dim) by channels_tiled (P dim) > > The B_i tensor shows the "p-dim broadcast" annotation with a downward arrow, indicating it is broadcast across the partition dimension to match deltaU_i's shape during the element-wise multiplication. > > **Key Elements:** +> > - **delta_i**: Green input tensor [channels_tiled x seqlen] > - **u_i**: Yellow input tensor [channels_tiled x seqlen] > - **deltaU_i**: Purple intermediate result [channels_tiled x seqlen] @@ -559,12 +551,10 @@ with `B_i` will have to be a `nisa.tensor_tensor`. Figure below visualizes the c > - **channels_tiled (P dim)**: Tiled channels mapped to partition dimension > - **p-dim broadcast**: Broadcast operation along partition dimension for B_i - Fig. 33 Step 3: p-dim broadcast + 2x `nisa.tensor_tensor` The associated NKI code is as follows: - ```python # Input shape in device memory does NOT match the computation layout # Device memory layout: @@ -582,13 +572,11 @@ B_i_bcast = B_i.broadcast_to((nl.tile_size.pmax, seq_len)) deltaBu_i = nisa.tensor_tensor(deltaU_i, B_i_bcast, op=ml.multiply) ``` - ### Step 4: Associative scan between deltaA_i and deltaBu_i In this step, we use an associative scan operator between `deltaA` and `deltaBu` to aggregate information across time sequentially (sequence length, e.g. sequence of tokens), from the past to the present. Here is a PyTorch reference implementation: - ```python # deltaA: [batch_size, channels, state_size, seq_len] # deltaB_u: [batch_size, channels, state_size, seq_len] @@ -602,13 +590,11 @@ for i in range(seq_len): out[..., i] = deltaA[..., i] * prev_state + deltaB_u[..., i] ``` - By holding batch and state_size dimensions constant, we get `deltaA_i` and `deltaBu_i` both with `[channels_tiled, seq_len]`, where `channels_tiled` is the partition dimension. The associative scan between these two tile shapes can be implemented in NKI naively through the following loop: - ```python scan_i = nl.ndarray((channels_tiled, seq_len), ...) @@ -621,14 +607,12 @@ for i in range(seq_len - 1): + deltaBu_i[0:channels_tiled, i+1] ``` - Within the loop, the current implementation invokes one instruction for multiplication and another for addition. Since both instructions are performed among tiles of shape `[channels_tiled, 1]`, we can combine these two instructions using [nisa.tensor_scalar](../api/api-nki-isa-tensor.md#nki-isa-tensor_scalar) which supports two operators in a pipelined fashion within an instruction at the same cost as a single operator. Below is a new implementation that could provide 2x speedup compared to the above: - ```python scan_i = nl.ndarray((channels_tiled, seq_len), dtype=deltaA.dtype, buffer=nl.sbuf) scan_i[0:channels_tiled, 0] = deltaBu[i_p, 0] @@ -642,31 +626,27 @@ for i in range(seq_len - 1): operand1=deltaBu[0:channels_tiled, i+1]) ``` - However, the above loop nest will turn into `seq_len` many instructions with input tiles that have a single element per partition in SBUF. In addition, every `nisa.tensor_scalar` instruction has a data dependency on the output of the previous instruction. As discussed in the [Trainium/Inferentia2 Architecture Guide](../../architecture/trainium_inferentia2_arch.md#arch-sec-vector-engine-perf), -these two traits combined in the instruction sequence is considered extremely *inefficient* on ScalarE/VectorE, where +these two traits combined in the instruction sequence is considered extremely _inefficient_ on ScalarE/VectorE, where the static instruction overhead instead of the useful execution time would be dominating the engine timeline. Conveniently, NKI exposes another instruction [nisa.tensor_tensor_scan](../api/api-nki-isa-tensor.md#nki-isa-tensor_tensor_scan) -on VectorE, which can perform the above loop nest in a *single* instruction by caching the intermediate scan result from +on VectorE, which can perform the above loop nest in a _single_ instruction by caching the intermediate scan result from the previous time step internally in VectorE without going through SBUF. - ```python scan_i = nisa.tensor_tensor_scan(deltaA_i, deltaBu_i, initial=0, op0=np.multiply, op1=np.add) ``` - Note, the shape of `scan_i` is exactly the same as the input `deltaA_i/deltaBu_i`: `[channels_tiled, seq_len]`. ### Step 5: Element-wise multiplication of C_i and scan_i The PyTorch reference implementation is: - ```python # scan_res: [batch_size, channels, state_size, seq_len] # C: [batch_size, state_size, seq_len] @@ -678,11 +658,9 @@ scanC = C[:, None, :, :] * scan_res scanC_i = C_i[None, :] * scan_i[:, :] ``` - You know the drill - Since `channels_tiled` is the partition dimension in `scan_i` from the previous step, we need to perform a partition-dimension broadcast on `C_i` before invoking `nisa.tensor_tensor`: - > **Figure: mamba step5** > > A tensor operation diagram for Mamba Step 5, showing the element-wise multiplication of scan_i with C_i (using p-dim broadcast) to produce scanC_i. @@ -692,17 +670,20 @@ perform a partition-dimension broadcast on `C_i` before invoking `nisa.tensor_te > Three tensors are displayed from left to right: > > On the left, a green rectangular tensor labeled "scan_i" has dimensions: +> > - Width: seqlen (F dim) - sequence length mapped to the free dimension > - Height: channels_tiled (P dim) - tiled channels mapped to the partition dimension > > In the center, a blue horizontal tensor labeled "C_i" has: +> > - Width: seqlen (F dim) - matching the free dimension > - A minimal height requiring broadcast > - An annotation "p-dim broadcast" with a downward-pointing arrow indicates this tensor will be broadcast across the partition dimension > -> A multiplication symbol (*) appears between scan_i and C_i. +> A multiplication symbol (\*) appears between scan_i and C_i. > > On the right, a purple rectangular tensor labeled "scanC_i" shows the result with dimensions: +> > - Width: seqlen (F dim) - same as inputs > - Height: channels_tiled (P dim) - same as scan_i > @@ -711,6 +692,7 @@ perform a partition-dimension broadcast on `C_i` before invoking `nisa.tensor_te > This step computes the element-wise product of the scan output with the C matrix, which is essential for computing the final output in the Mamba selective state space model. > > **Key Elements:** +> > - **scan_i**: Green input tensor from scan operation [channels_tiled x seqlen] > - **C_i**: Blue tensor requiring p-dim broadcast along partition dimension > - **scanC_i**: Purple output tensor [channels_tiled x seqlen] @@ -718,25 +700,21 @@ perform a partition-dimension broadcast on `C_i` before invoking `nisa.tensor_te > - **channels_tiled (P dim)**: Tiled channels mapped to partition dimension > - **p-dim broadcast**: Downward arrow indicating broadcast along partition dimension - Fig. 34 Step 5: p-dim broadcast + `nisa.tensor_tensor` The corresponding NKI code is: - ```python C_i_bcast = C_i.broadcast((nl.tile_size.pmax, seq_len)) scanC_i = nisa.tensor_tensor(scan_i, C_i_bcast, op=ml.multiply) ``` - ### Step 6: Accumulation of scanC_i along `state_size` dimension So far in Step 1-5, all the computation is logically parallel across the `state_size` dimension in a Mamba layer. The next step of computation introduces data dependency along the `state_size` dimension for the first time. The PyTorch reference implementation is: - ```python # scan_res: [batch_size, channels, state_size, seq_len] # C: [batch_size, state_size, seq_len] @@ -748,11 +726,9 @@ scanC.sum(dim=-2) (scanC_i).sum(dim=-2) ``` - In NKI, we can accumulate the `scanC_i` results across states element-wise using `state_size-1` number of `nisa.tensor_tensor` instructions: - > **Figure: mamba step6** > > A tensor summation diagram for Mamba Step 6, showing the reduction of multiple scanC_i tensors (indexed 0 through n-1) to produce a single scanC_i_sum output tensor. @@ -762,11 +738,13 @@ instructions: > Four tensors are displayed from left to right, connected by addition and equals operators: > > The first three tensors (purple/lavender colored) represent individual scanC_i results for different state indices: +> > - "scanC_i[0]" - first state component with dimensions seqlen (F dim) by channels_tiled (P dim) > - "scanC_i[1]" - second state component with same dimensions > - "scanC_i[n-1]" - final state component (with ellipsis "..." between [1] and [n-1] indicating intermediate components) > > Each purple tensor has the same dimensions: +> > - Width: seqlen (F dim) - sequence length on free dimension > - Height: channels_tiled (P dim) - tiled channels on partition dimension > @@ -779,6 +757,7 @@ instructions: > This step performs a reduction across the state dimension (n states), summing all the scanC_i components to produce the final output contribution. > > **Key Elements:** +> > - **scanC_i[0]**: First purple tensor in summation > - **scanC_i[1]**: Second purple tensor in summation > - **scanC_i[n-1]**: Last purple tensor in summation (n-th state) @@ -788,7 +767,6 @@ instructions: > - **Plus signs (+)**: Addition operators showing element-wise summation > - **Ellipsis (...)**: Indicates additional intermediate tensors in the summation - Fig. 35 Step 6: `state_size-1` number of `nisa.tensor_tensor` Since we will be looping over different states, we can also declare an empty accumulation buffer `scanC_accum` of shape @@ -796,7 +774,6 @@ Since we will be looping over different states, we can also declare an empty acc iteration using `+=` operator. The use of a single accumulation buffer avoids allocating memory for `scanC_i` across all states in SBUF. The corresponding NKI code is: - ```python scanC_accum = nl.zeros(...) @@ -805,12 +782,10 @@ for i_state in range(state_size): scanC_accum += scanC_i ``` - ## Initial NKI Kernel Putting all the pieces together from the previous section, we can arrive at the below kernel implementation `mamba_v1`: - ```python import nki import nki.language as nl @@ -895,48 +870,44 @@ def mamba_v1(delta, u, A, B, C): return output ``` - In the above code example, -* -We have three levels of loop nests. From the outer-most to inner-most: +- We have three levels of loop nests. From the outer-most to inner-most: Iterating over `batch`: Different batch samples perform completely different computation. `A` tensor is the only input parameter that is shared among batch samples. -* Iterating over `state_size`: Different states perform parallel computation until Step 6 as discussed in the previous -section. Both `delta` and `u` tensors are shared across different states. +- Iterating over `state_size`: Different states perform parallel computation until Step 6 as discussed in the previous + section. Both `delta` and `u` tensors are shared across different states. -* Iterating over `channels`: This is the most-inner dimension where we tile the input channels dimension into `nl.tile_size.pmax=128` -chunks. Both `B` and `C` tensors are shared across different `channels`. +- Iterating over `channels`: This is the most-inner dimension where we tile the input channels dimension into `nl.tile_size.pmax=128` + chunks. Both `B` and `C` tensors are shared across different `channels`. -* The kernel above assumes channels is a multiple of `nl.tile_size.pmax=128` . We can relax this by adding a `mask` -parameter in all the NKI API call in the kernel. To simplify the code example, we omit this change. -See NKI API Masking for more information. +- The kernel above assumes channels is a multiple of `nl.tile_size.pmax=128` . We can relax this by adding a `mask` + parameter in all the NKI API call in the kernel. To simplify the code example, we omit this change. + See NKI API Masking for more information. -* We declare an empty intermediate tensor `scanC_accum` to hold partial summation from every state. +- We declare an empty intermediate tensor `scanC_accum` to hold partial summation from every state. -* -Within the inner loop, we process data for `nl.tile_size.pmax=128` channels for one batch sample in one state. +- Within the inner loop, we process data for `nl.tile_size.pmax=128` channels for one batch sample in one state. We use the slicing syntax to index a tensor. For example, `delta[i_batch, channel_start:channel_start+channel_psize, 0:seq_len]` grabs data from the input `delta` tensor for the current range of channels at the current batch sample. -* Note, in tensor slicing, the first index dimension from the left with a slicing range will be chosen as the partition -dimension. When loading `B`, since we intend to load only one state’s worth of data into one partition of SBUF (discussed -in Step 3), we need to explicitly slice the state using: `nl.load(B[i_batch, **i_state:i_state+1**, 0:seq_len])`. Otherwise, -`nl.load(B[i_batch, **i_state**, 0:seq_len])` will treat `seq_len` as the partition dimension, which is not what we -planned for in Step 3 and would also trigger a NKI compilation error since `seq_len` exceeds `nl.tile_size.pmax`. +- Note, in tensor slicing, the first index dimension from the left with a slicing range will be chosen as the partition + dimension. When loading `B`, since we intend to load only one state’s worth of data into one partition of SBUF (discussed + in Step 3), we need to explicitly slice the state using: `nl.load(B[i_batch, **i_state:i_state+1**, 0:seq_len])`. Otherwise, + `nl.load(B[i_batch, **i_state**, 0:seq_len])` will treat `seq_len` as the partition dimension, which is not what we + planned for in Step 3 and would also trigger a NKI compilation error since `seq_len` exceeds `nl.tile_size.pmax`. -* We accumulate partial `scanC_i` results into the accumulation buffer using the `+=` operator. This creates a loop-carried -dependency for `scanC_accum` on the `i_state` loop. +- We accumulate partial `scanC_i` results into the accumulation buffer using the `+=` operator. This creates a loop-carried + dependency for `scanC_accum` on the `i_state` loop. ### Performance Check Let’s re-run neuron-profile on the above NKI kernel: - > **Figure: mamba v1 profile** > > A Neuron profiler timeline showing the Mamba v1 optimized implementation with significantly improved execution time of 172.93 microseconds, demonstrating better compute utilization compared to the torch reference. @@ -950,6 +921,7 @@ Let’s re-run neuron-profile on the above NKI kernel: > The memory section shows State Buffer Usage as a stacked colored area chart that ramps up, maintains high usage through the middle portion, then gradually decreases. PSUM Usage appears as a multi-colored stacked area below it. Sem 0 tracks semaphore activity. Pending DMA Count shows a red line with a spike around 20,000 us. DMA Throughput displays as a green line showing data movement rates. The total execution time of 172.93 us is noted at the bottom. > > **Key Elements:** +> > - **Execution time**: 172.93 microseconds total (major improvement from 133.43 ms reference) > - **TensorE/TensorMatrixE**: Active matrix computation marks showing tensor engine utilization > - **VectorE/ScalarE**: Continuous bars at the end indicating optimized vector and scalar operations @@ -959,7 +931,6 @@ Let’s re-run neuron-profile on the above NKI kernel: > - **DMA Throughput**: Green line tracking data movement bandwidth > - **Timeline scale**: 0 to 180,000 us with activity concentrated in first 80,000 us - Fig. 36 Profile of initial Mamba kernel implementation `mamba_v1` Hooray! This NKI kernel implementation now takes `172.93` usec, which is **878x** speedup compared to the reference PyTorch @@ -970,7 +941,6 @@ Therefore, our goal is to keep VectorE as busy as possible throughout execution. start-up and tear-down overhead. We can use the `Selection Summary` feature in `neuron-profile` to find out the percentage of time VectorE is busy during the actual execution period: - > **Figure: mamba v1 profile zoomed** > > A zoomed-in Neuron profiler view of the Mamba v1 implementation with a Selection Summary panel showing that the Vector Engine achieves 98.71% active duration, indicating excellent compute utilization. @@ -984,6 +954,7 @@ of time VectorE is busy during the actual execution period: > Below the timeline, a toolbar shows various blue buttons including Search, Annotations, Edit view settings, Summary, Layer Summary, Selection Summary, NEFF Header, NEFF Nodes, Model Info, DMA Queues Info, and NC Mem. A Selection Summary popup panel is displayed showing detailed metrics for a selected region. > > **Key Elements:** +> > - **VectorE track**: Dense blue bars showing 98.71% active utilization (181 events) > - **TensorMatrixE**: Red marks indicating tensor matrix operations > - **ScalarE/GpSimdE**: Blue bars showing additional compute engine activity @@ -993,7 +964,6 @@ of time VectorE is busy during the actual execution period: > - **DMA Throughput**: Green line tracking data movement bandwidth > - **Timeline scale**: 0 to 150,000 us within 246.2 us total window - Fig. 37 Profile of initial Mamba kernel implementation `mamba_v1` (zoomed in) As indicated by the above profile, VectorE is active over **98.71%** of the time, which is rather impressive. However, @@ -1027,7 +997,7 @@ VectorE idle period, you can also see the exact input tensor name defined in NKI Fig. 40 DMA loading tensor u in `mamba_v1` profile We can find similar VectorE gaps through the execution trace. At this point, we can conclude one of the reasons why we have -a lower VectorE active time percentage is due to *blocking* input tensor loading (`nl.load`) activities in the DMA. +a lower VectorE active time percentage is due to _blocking_ input tensor loading (`nl.load`) activities in the DMA. Next, let’s spend some time analyzing DMA efficiency. Zooming out, we can make several observations. First, we see two orange boxes around the `qSyncIO0` row. Hovering over @@ -1052,43 +1022,42 @@ NKI kernel to avoid it. To understand why delta and u are being reloaded, let’s revisit our input tensor shapes: -* `delta: [batch_size, channels, seq_len]` +- `delta: [batch_size, channels, seq_len]` -* `u:     [batch_size, channels, seq_len]` +- `u:     [batch_size, channels, seq_len]` -* `A:     [channels, state_size]` +- `A:     [channels, state_size]` -* `B:     [batch_size, state_size, seq_len]` +- `B:     [batch_size, state_size, seq_len]` -* `C:     [batch_size, state_size, seq_len]` +- `C:     [batch_size, state_size, seq_len]` Let’s hold `batch_size` constant since the majority of input tensors have completely different slices for different batch samples: -* `delta: [channels, seq_len]` +- `delta: [channels, seq_len]` -* `u:     [channels, seq_len]` +- `u:     [channels, seq_len]` -* `A:     [channels, state_size]` +- `A:     [channels, state_size]` -* `B:     [state_size, seq_len]` +- `B:     [state_size, seq_len]` -* `C:     [state_size, seq_len]` +- `C:     [state_size, seq_len]` `delta` and `u` tensors have the same shape with `channels` as the outer dimensions, while `B` and `C` have the same shape with `state_size` as the outer dimension. All four of these input tensors have `seq_len` as the inner dimension. Therefore, we say `delta/u` is reused across different states, while `B/C` are reused across different channels. Given -this conflicting reuse dimensions, we further say it is more important to **prioritize reuse of ``delta/u``** because +this conflicting reuse dimensions, we further say it is more important to **prioritize reuse of `delta/u`** because the expected size of `channels` is much higher than `state_size`: -* `state_size` is now 16 and typically stay small +- `state_size` is now 16 and typically stay small -* `channels` is now 4096 and typically in the thousands +- `channels` is now 4096 and typically in the thousands In NKI, we can prioritize `delta/u` reuse through loop ordering. Recall in the initial NKI kernel implementation, we have the following inner loops: - ```python ... for i_state in range(state_size): @@ -1097,7 +1066,6 @@ for i_state in range(state_size): ... ``` - Since these two loops are executed serially within a single NeuronCore, the loop instances will be unrolled by Neuron Compiler. With the channel dimension in the fastest dimension, we will need to load `delta/u` across all channels in the first state, and then likely reload them again in the later states due to a large total memory size in `delta` and `u` (16MB in this @@ -1106,7 +1074,6 @@ case). To prioritize reuse of `delta/u`, we should reorder the above loop nests. To further enforce the reuse, we can hoist the `nl.load` calls for `delta/u` outside of the `i_state` inner loop: - ```python ... for i_channel_tile in range(n_channel_tile): @@ -1118,11 +1085,9 @@ for i_channel_tile in range(n_channel_tile): ... ``` - As a side effect of this loop re-ordering, we can also spot a loop fusion opportunity since we have two `i_channel_tile` loop nests at the same level now: - ```python scanC_accum = nl.zeros((n_channel_tile, nl.par_dim(channel_psize), seq_len), ...) ... @@ -1142,12 +1107,10 @@ for i_channel_tile in range(n_channel_tile): ... ``` - By fusing the two `i_channel_tile` loop nests into a single loop nest, we can pull the declaration of `scanC_accum` inside the `i_channel_tile` loop and further reduce the `scanC_accum` size requirement by a factor of `n_channel_tile` : - ```python ... @@ -1166,10 +1129,8 @@ for i_channel_tile in range(n_channel_tile): ... ``` - Let’s modify our initial NKI kernel implementation accordingly to get `mamba_v2`: - ```python @nki.jit def mamba_v2(delta, u, A, B, C): @@ -1245,10 +1206,8 @@ def mamba_v2(delta, u, A, B, C): return output ``` - We recapture the profile for the new kernel implementation: - > **Figure: mamba v2** > > A Neuron profiler timeline showing the Mamba v2 optimized implementation with a Selection Summary panel displaying Vector Engine achieving approximately 99.67% active duration, demonstrating further improved compute utilization. @@ -1264,6 +1223,7 @@ We recapture the profile for the new kernel implementation: > At the bottom, a toolbar with blue buttons is visible, including Search, Annotations, Edit view settings, Summary, Layer Summary, Selection Summary, NEFF Header, NEFF Nodes, Model Info, DMA Queues Info, NC Memory Usage Info, and more. A Selection Summary popup shows detailed metrics for VectorE activity with count of 2086 events and active duration of approximately 99.67%. > > **Key Elements:** +> > - **VectorE track**: Shows 2086 events with ~99.67% active duration (improvement over v1) > - **TensorMatrixE**: Periodic red marks indicating matrix tensor operations > - **State Buffer Usage**: Multi-colored stacked area showing memory utilization @@ -1273,7 +1233,6 @@ We recapture the profile for the new kernel implementation: > - **Selection Summary panel**: Displays Duration, Start/End Time, Event Count, Event Duration Sum, and Event Duration Active metrics > - **Timeline scale**: Spans 0 to ~1,500,000 us showing full kernel execution - Fig. 43 Profile of `mamba_v2` kernel with loop reordering optimization The device execution time is now **1.61 ms**, which is a **31%** reduction in latency compared to our initial kernel implementation. @@ -1315,22 +1274,20 @@ level. We have **three** key considerations when adding this new loop level: -* tile size selection, +- tile size selection, -* loop-carried dependency handling +- loop-carried dependency handling -* loop ordering with other loop nests. +- loop ordering with other loop nests. -**Tile size of ``seq_len``.** Since previously with `seq_len=512` in our toy example, we were able to achieve close to +**Tile size of `seq_len`.** Since previously with `seq_len=512` in our toy example, we were able to achieve close to 100% VectorE utilization, let’s set the tile size `seq_len_fsize` to 512 as a starting point. We can revisit this decision as needed once we obtain a new profile. -**Loop-carried dependency.** Splitting `seq_len` into chunks is straightforward for all computation steps except for Step -4. In the associative scan operation, the next loop iteration requires results from the previous iteration for computation. +**Loop-carried dependency.** Splitting `seq_len` into chunks is straightforward for all computation steps except for Step 4. In the associative scan operation, the next loop iteration requires results from the previous iteration for computation. As a result, we will introduce another loop-carried dependency here with the scan tiles. This dependency can be handled through the `initial` input parameter: - ```python scan_init = nl.zeros((channel_psize, 1), ...) @@ -1340,12 +1297,10 @@ for i_seq_len_tile in range(seq_len // seq_len_fsize): scan_init = scan_i[0:channel_psize, seq_len_fsize-1] ``` - Note the loop-carried dependency: `scan_init` is updated each iteration and used as the initial value in the next. **Loop ordering.** Recall from our latest NKI kernel implementation, we have the following loop nest: - ```python ... for i_batch in range(batch_size): @@ -1372,7 +1327,6 @@ for i_batch in range(batch_size): ... ``` - Let’s denote the above loop ordering as `[batch_size, n_channel_tile, state_size]`, and our key question here is where to insert `seq_len` in this list. @@ -1382,7 +1336,6 @@ loop ordering won’t be tiling `scanC_accum`, `delta_i` and `u_i` tensors. Give these three tensors will occupy 8192*4B*3 = 96 KiB/partition, half of the available SBUF capacity. Let’s go ahead and experiment this loop ordering in a new kernel `mamba_v3`: - > **Figure: mamba v3** > > A Neuron profiler timeline showing the Mamba v3 implementation with an extended execution timeline, displaying Vector Engine utilization at 94.80% active duration across 38,009 events. @@ -1398,6 +1351,7 @@ experiment this loop ordering in a new kernel `mamba_v3`: > The bottom portion shows a toolbar with blue buttons including Search, Annotations, View Settings, Summary, Layer Summary, Selection Summary, NEFF Header, NEFF Nodes, Model Info, DMA Queues Info, NC Memory Usage Info, Summarize, and Help. A Selection Summary popup displays detailed metrics showing Duration, Start Time, End Time, Event Count of 38,009, Event Duration Sum, and Event Duration Active at 94.80%. > > **Key Elements:** +> > - **qSyncIO0**: Continuous orange/yellow bar indicating persistent I/O sync activity > - **VectorE activity**: 38,009 events with 94.80% active duration > - **TensorMatrixE**: Red marks distributed across the timeline for matrix operations @@ -1407,7 +1361,6 @@ experiment this loop ordering in a new kernel `mamba_v3`: > - **DMA Throughput**: Green line showing periodic data transfer activity > - **Execution pattern**: More distributed activity compared to v1/v2 optimizations - Fig. 46 Profile of `mamba_v3` kernel with seq_len tiling optimization With the above profile, the kernel now takes **27.8 ms**, which is **48%** reduction in latency compared to no `seq_len` @@ -1420,18 +1373,17 @@ from `seq_len=512`. We evaluate scaling efficiency using `perfect latency / meas which is a higher the better metric. Finally, to showcase the importance of the last seq_len tiling optimization for scaling seq_len, we also compare scaling efficiency for `mamba_v2` (no seq_len tiling) and `mamba_v3` (seq_len tiling). - | seq_len | Perfect Latency (ms) | mamba_v2 Measured Latency (ms) | mamba_v2 Scaling Efficiency | mamba_v3 Measured Latency (ms) | mamba_v3 Scaling Efficiency | -| --- | --- | --- | --- | --- | --- | -| 512 | N/A | 1.6 | N/A | 1.6 | N/A | -| 1024 | 3.2 | 4.4 | 72.73% | 3.3 | 96.97% | -| 2048 | 6.4 | 8.9 | 71.91% | 6.6 | 96.97% | -| 3072 | 9.6 | 13.1 | 73.28% | 10.1 | 95.05% | -| 4096 | 12.8 | 17.6 | 72.73% | 13.3 | 96.24% | -| 5120 | 16 | 23.7 | 67.51% | 17.3 | 92.49% | -| 6144 | 19.2 | 27.5 | 69.82% | 19.6 | 97.96% | -| 7168 | 22.4 | 41.3 | 54.24% | 24.2 | 92.56% | -| 8192 | 25.6 | 52.2 | 49.04% | 27.8 | 92.09% | +| ------- | -------------------- | ------------------------------ | --------------------------- | ------------------------------ | --------------------------- | +| 512 | N/A | 1.6 | N/A | 1.6 | N/A | +| 1024 | 3.2 | 4.4 | 72.73% | 3.3 | 96.97% | +| 2048 | 6.4 | 8.9 | 71.91% | 6.6 | 96.97% | +| 3072 | 9.6 | 13.1 | 73.28% | 10.1 | 95.05% | +| 4096 | 12.8 | 17.6 | 72.73% | 13.3 | 96.24% | +| 5120 | 16 | 23.7 | 67.51% | 17.3 | 92.49% | +| 6144 | 19.2 | 27.5 | 69.82% | 19.6 | 97.96% | +| 7168 | 22.4 | 41.3 | 54.24% | 24.2 | 92.56% | +| 8192 | 25.6 | 52.2 | 49.04% | 27.8 | 92.09% | The above data shows the last NKI kernel implementation `mamba_v3` can reach 90%+ scaling efficiency up to 8K `seq_len`. To support even larger `seq_len`, we will need more aggressive tiling by pulling the `seq_len` loop level further @@ -1442,9 +1394,9 @@ towards the outer-loop level to tile more input/intermediate tensors to keep spi Click the links to download source code of the kernels and the testing code discussed in this tutorial. -* PyTorch reference implementation: [`mamba_torch.py`](../../downloads/mamba_torch.py) +- PyTorch reference implementation: [`mamba_torch.py`](../../downloads/mamba_torch.py) -* Three versions of NKI kernels: [`mamba_nki_kernels.py`](../../downloads/mamba_nki_kernels.py) +- Three versions of NKI kernels: [`mamba_nki_kernels.py`](../../downloads/mamba_nki_kernels.py) You can also view the source code in the GitHub repository [nki_samples](https://github.com/aws-neuron/nki-samples/tree/main/src/nki_samples/tutorials/fused_mamba/) @@ -1454,33 +1406,26 @@ You can also view the source code in the GitHub repository [nki_samples](https:/ Run PyTorch reference implementation to generate a NEFF for profiling: - ```python python3 mamba_torch.py --mode perf ``` - Check performance numbers of mamba_v1/mamba_v2/mamba_v3: - ```python python3 mamba_nki_kernels.py --mode perf --version v1 v2 v3 --batch 1 --seq_len 2048 --channels 512 --state_size 16 ``` - **Accuracy mode** Check mamba_v1 NKI kernel accuracy against PyTorch implementation: - ```python python3 mamba_torch.py --mode accuracy ``` - Check optimized Mamba kernel (mamba_v2, mamba_v3) accuracy against mamba_v1: - ```python python3 mamba_nki_kernels.py --mode accuracy --version v1 v2 v3 --batch 1 --seq_len 2048 --channels 512 --state_size 16 -``` \ No newline at end of file +``` diff --git a/skills/neuron-nki-docs/references/programming/tutorials/kernel-optimization.md b/skills/neuron-nki-docs/references/programming/tutorials/kernel-optimization.md index 5bc86d1..a0f90d3 100644 --- a/skills/neuron-nki-docs/references/programming/tutorials/kernel-optimization.md +++ b/skills/neuron-nki-docs/references/programming/tutorials/kernel-optimization.md @@ -11,9 +11,9 @@ Developers commonly create NKI kernels to accelerate critical operations in larg This concept is applicable to: -* Improving the performance of critical sections of ML inference or training models. +- Improving the performance of critical sections of ML inference or training models. -* Writing small performant kernels for standalone ML inference or training. +- Writing small performant kernels for standalone ML inference or training. ## When to write a kernel? @@ -27,7 +27,6 @@ The end goal of writing a kernel is to improve the performance of the model, but NKI provides the `nki.isa.nc_matmul` instruction to perform a matrix multiply. This instruction operates over a restricted sized matrix with at most a 128 x 128 “stationary” (weights) matrix and a 128 x 512 “moving” (ifmap) matrix. This allows you to produce a 128 x 512 matrix, at most, as output. The “stationary” matrix must be transposed to get a result that is not transposed. To call the `nki.isa.nc_matmul` instruction, provide to the state buffer (SBUF), and the result will be written into the partial sum buffer (PSUM). If you use a small driver program to invoke the kernel, the arguments will be passed in from the device memory (HBM) and the result will be read from HBM as well. The kernel will move inputs from HBM to SBUF, call the `nki.isa.nc_matmul` instruction, move the result from PSUM to SBUF (you cannot move data directly from PSUM to HBM), and then from SBUF to HBM. - ```python import os import nki @@ -95,10 +94,10 @@ def matrix_multiply_kernel(lhsT, rhs): return result ``` - This small kernel allows you to experiment with the `nki.isa.nc_matmul` instruction and you can test that it works with a simple driver. PyTorchJAX + ```python import numpy as np import torch @@ -134,8 +133,6 @@ else: print(result_torch) ``` - - ```python import numpy as onp import jax.numpy as jnp @@ -165,10 +162,8 @@ else: print(result_jax) ``` - You can validate that you have the correct understanding of the nki.isa.nc_matmul instruction by invoking your test: - ```bash $ python driver.py Kernel computed correct output @@ -181,7 +176,6 @@ tensor([[35.7896, 32.8659, 31.6545, ..., 37.1804, 31.4682, 33.9796], [32.4571, 29.1864, 31.7483, ..., 33.3723, 30.1617, 29.8077]]) ``` - (Note that there will be some additional output, which varies slightly depending on which framework you use. The values will also vary, since the inputs are randomly generated.) As you become more familiar with NKI, you will no longer need to start with quite so simple a variation on the kernel. While this kernel allowed us to validate our understanding of the `nki.isa.nc_matmul` instruction, it will not allow you to pass in matrices larger than a single tile. A more realistic variant of the kernel needs to take matrices larger than the tile size, break down the inputs into single tiles, compute each output tile, then write the result back to HBM. @@ -190,7 +184,6 @@ As you become more familiar with NKI, you will no longer need to start with quit The simple start allowed us to validate our understanding of the `nki.isa.matmul` instruction. The following kernel shows how you can do this with input matrices that are larger than a single tile size. You may recognize the traditional three nested loop structure of matrix multiply, but instead of the inner body computing a scalar value it operates over a full tile. - ```python import os import nki @@ -278,25 +271,23 @@ def matrix_multiply_kernel(lhsT, rhs): return result ``` - The tiled version expects the input and output matrices to be a multiple of the tile sizes. In cases where the matrices you want to multiply do not match that, they can be padded or the implementation could be extended to handle the sub-tile sized edges. The body of the n and m loops allocates a result_tile in the PSUM. The inner-most k loop then loads the tiles from the lhsT and rhs inputs into SBUF from HBM, performs the matrix multiply, accumulating the result into the result_tile. After the k loop completes, the m, n tile has been computed and can be moved from PSUM to SBUF and then written into the correct position in the result HBM. Now that you have a kernel that can handle what you expect the model to need, you can extend the small test driver above to ensure you can keep the kernel functioning correctly as you begin to improve the performance of the kernel. This driver is something you can continue to use with each progressive improvement of the kernel. This is just a variation on the original test that provides input matrices large enough to represent the real workload the kernel will be expected to handle. In this case that just means increasing the size of the input matrices from a single tile at 128x128 x 128x512 to something slightly more realistic at 4096x8192 x 8192x8192. You can update the numpy generation of inputs to set the lhs and rhs to the new dimensions. - ```python lhs = rng.random((4096, 8192), dtype=np.float32) rhs = rng.random((8192, 8192), dtype=np.float32) ``` - It is important to select input sizes that are realistic (or at least representative) of the real work you expect the kernel to handle, because you will use this test not just for correctness, but also to allow you to profile the kernel to guide improvements on the kernel’s performance. In addition to changing the size of the input to the kernel, you will also want to enable profiling of the kernel. You will use the approach described in the [Neuron Explorer user guide](../api/index.md) to profile just the call to the NKI matrix multiply kernel. With this you can surround the call to the kernel with the profiling context. PyTorchJAX + ```python from torch_neuronx.experimental import profiler ... @@ -308,8 +299,6 @@ with profiler.profile(port=9012, result_device = matrix_multiply_kernel(lhsT_device, rhs_device) ``` - - ```python import jax ... @@ -317,7 +306,6 @@ with jax.profiler.trace("./output"): result_jax = matrix_multiply_kernel(lhsT_jax, rhs_jax) ``` - When you run the test driver, in addition to showing that the output matches the numpy result, you will also get both the Neuron Execution File Format (NEFF) file, which is what executes on the accelerator and the Neuron Timing File Format (NTFF) file generated by running the kernel with profiling enabled. You can use these two files with the neuron_profiler to view the results of running the kernel. Looking at this profile for the full kernel run, you can see that the DMA queues which move data from HBM to SBUF and back are quite active. Looking at the Tensor and TensorMatrix lines, it appears there are some gaps within the run as well. The heavy use of DMA and the Tensor Engine (TensorE) is not too surprising, since those are the two things the kernel is primarily doing. The profile also provides some data as an overview of how much each engine is being used. You can zoom in to one of the areas where you see a gap and validate the impression. @@ -328,10 +316,8 @@ You can see that the TensorE is busy from the start of the kernel through the en ![../../../_images/v2-zoom.png](../../../_images/v2-zoom.png) However, there are gaps between matrix multiply operations indicate that the TensorE is waiting on data to be read from the HBM to SBUF for the next operation to take place that we can see when we zoom in.Looking at the original kernel code you can see that you are loading the two tiles before each matrix multiply. Looking at the summary data provided in the profile, you can also see that the DMA engines were active 99.93% of the time while the TensorE was only active 87.28% of the run. - | [ ](../../../_images/v2-dma.png) | [ ](../../../_images/v2-pe.png) | -| --- | --- | - +| -------------------------------- | ------------------------------- | ## Analyzing the kernel @@ -339,7 +325,6 @@ The first step to improving the performance of the kernel is to analyze the perf Looking at this, you might notice two things. First, since the data for each matrix multiply is being loaded just before the multiply, you are always waiting on these loads to complete before you can start the next multiply. If you look at the structure of the iteration, you can also see that you will load the same tile more than once. For instance the m=0, k=0 tile will be loaded N // TILE_N times. One change you could make is to load all of the tiles needed to compute a given output tile before you start the computation. You can accomplish this by moving the loads out into the outer loops, loading all K // TILE_K tiles for a given value of m from the stationary matrix at the start of the m loop, and all K // TILE_K tiles for a given value of n from the stationary matrix at the start of the n loop. - ```python import os import nki @@ -440,7 +425,6 @@ def matrix_multiply_kernel(lhsT, rhs): return result ``` - The test program validates that the new implementation is correct and also provides new NEFF and NTFF. ![../../../_images/v3-full.png](../../../_images/v3-full.png) @@ -449,10 +433,8 @@ At this level the profile does not look too different, however when you zoom in, ![../../../_images/v3-zoom.png](../../../_images/v3-zoom.png) Analyzing the improvement though, you can see that this change has made big strides. The DMA and matrix multiply is better overlapped, the DMA engines are now busy 99.73% of the time, slightly more than before, but the TensorE is busy 99.85% of the time. This is a huge improvement, but the time spent in the kernel is still dominated by DMA. - | [ ](../../../_images/v3-dma.png) | [ ](../../../_images/v3-pe.png) | -| --- | --- | - +| -------------------------------- | ------------------------------- | ## Overlapping data and compute through blocking @@ -460,7 +442,6 @@ The previous refinement of the kernel showed that you can improve the utilizatio Blocking is a technique to help load even larger amounts of data in at a time. Instead of copying single tiles of data from HBM to SBUF, you can load a full block, which is a multiple of the number of tiles. Since matrix multiply still needs to operate tile by tile, you compute all of the tiles in the block before proceeding to the next block. - ```python import os import nki @@ -585,7 +566,6 @@ def matrix_multiply_kernel(lhsT, rhs): return result ``` - Running the test driver ensures the new implementation of the kernel is correct and provides a new NEFF and NTFF that helps us understand the improvements. ![../../../_images/v4-full.png](../../../_images/v4-full.png) @@ -594,16 +574,13 @@ Zooming in on a similarly sized section shows that while the overall time of the ![../../../_images/v4-zoom.png](../../../_images/v4-zoom.png) Again you can see gaps in the matrix multiply. Even though the new implementation of the kernel improves on the overall time of the kernel, the new implementation reduces the number of DMA instructions, because each instruction loads more, but you wait longer for each block to load. In fact, even though the performance improved the TensorE is actually less utilized as a percentage of time, dropping to 99.52% of the time, with the DMA engines hitting 95.70%. This means there is a small amount of time when only the TensorE is being used, but the DMA engine is still active for most of the kernel run, which you should expect could be smaller. - | [ ](../../../_images/v4-dma.png) | [ ](../../../_images/v4-pe.png) | -| --- | --- | - +| -------------------------------- | ------------------------------- | ## Optimizing DMA through blocking the contraction dimension One of the advantages of leaving the K dimension unblocked was that you could rely on the PSUM buffer to hold the final computed value. To block in the K dimension, you will need to store intermediate partial sums in a temporary SBUF array of tiles. The nki.isa.tensor_tensor instruction can be used to add two tensors, allowing you to accumulate into the temporary tile. With this, you can build blocks in all three dimensions. This version of blocking loads the blocks to in BLOCK_K by BLOCK_M and BLOCK_K by BLOCK_N dimensions. - ```python import os import nki @@ -768,7 +745,6 @@ def matrix_multiply_kernel( return result ``` - This version of the kernel is considerably more complicated, but the test driver you created for the simplest version of this kernel means you have a ready test. The sizes of matrices you chose in the original test were forward-looking in that they correspond to the tiling dimensions you selected. However, you expose these as additional arguments (unlike in the previous blocking), so a model calling this kernel can choose block sizes appropriate for the model. The test driver also gives us a new set of NEFF and NTFF files. ![../../../_images/v5-full.png](../../../_images/v5-full.png) @@ -777,10 +753,9 @@ Other than the improved time, this seems similar to the other profile graphs, ho ![../../../_images/v5-zoom.png](../../../_images/v5-zoom.png) Zooming in you can see the gap at the end of the set of matrix multiplies where the results are accumulated into the SBUF temporary results. Looking at the utilization of the DMA engines and TensorE you can see the DMA engines are now active only 21.54% of the time, while the TensorE is now active 99.50%, with the Vector Engine (VectorE) active 10.55% of the time, where it was previously unused. - -| [ ](../../../_images/v5-dma.png) | [ ](../../../_images/v5-pe.png) | -| --- | --- | -| | [ ](../../../_images/v5-vec.png) | +| [ ](../../../_images/v5-dma.png) | [ ](../../../_images/v5-pe.png) | +| -------------------------------- | -------------------------------- | +| | [ ](../../../_images/v5-vec.png) | This final version of the matrix multiply kernel is no longer memory-bound. Instead, as you should expect, it is compute-bound with the TensorE and VectorE engines being the limiting factor on the speed of the kernel. @@ -792,14 +767,14 @@ Once the kernel is ready you use it to replace the section of the model it is in ## Related concepts -* [Tutorial: Matrix multiplication](matrix_multiplication.md) +- [Tutorial: Matrix multiplication](matrix_multiplication.md) -* [Profiling NKI kernels with Neuron Explorer](../../optimization/use-neuron-profile.md) +- [Profiling NKI kernels with Neuron Explorer](../../optimization/use-neuron-profile.md) ## Further reading -* [NKI Language Guide](../nki-language-guide.md) +- [NKI Language Guide](../nki-language-guide.md) -* [NeuronDevice Architecture Guide for NKI](../../architecture/trainium_inferentia2_arch.md) +- [NeuronDevice Architecture Guide for NKI](../../architecture/trainium_inferentia2_arch.md) -* [NKI Performance Guide](../../optimization/nki_perf_guide.md) \ No newline at end of file +- [NKI Performance Guide](../../optimization/nki_perf_guide.md) diff --git a/skills/neuron-nki-docs/references/programming/tutorials/matrix_multiplication.md b/skills/neuron-nki-docs/references/programming/tutorials/matrix_multiplication.md index 138a746..2e46e99 100644 --- a/skills/neuron-nki-docs/references/programming/tutorials/matrix_multiplication.md +++ b/skills/neuron-nki-docs/references/programming/tutorials/matrix_multiplication.md @@ -4,14 +4,15 @@ Matrix multiplication In this tutorial, we will start with a simple NKI matrix multiplication kernel and optimize it step by step. In doing so, we learn about: -* The NKI syntax and programming model. +- The NKI syntax and programming model. -* Layout, tiling, and memory management considerations when performing -matrix multiplication in NKI. +- Layout, tiling, and memory management considerations when performing + matrix multiplication in NKI. ## Basic compute kernel ! + > **Figure: matrix multiplication views** > > A diagram comparing mathematical matrix multiplication view with Tensor Engine view, showing how lhs and rhs matrices map to lhs_T (stationary in Tensor Engine), rhs (moving from SBUF), and output locations. @@ -19,12 +20,14 @@ matrix multiplication in NKI. > This diagram illustrates the mapping between standard matrix multiplication notation and NeuronCore Tensor Engine execution, divided into two parts by a dashed line. > > Part (a) "Mathematical View" (left side) shows: +> > - A blue matrix labeled "rhs" at the top with dimensions N (width) by K (height) > - A green matrix labeled "lhs" at the bottom left with dimensions K (width) by M (height) > - A purple matrix labeled "output" at the bottom right with dimensions N (width) by M (height) -> - This represents the standard lhs * rhs = output matrix multiplication +> - This represents the standard lhs \* rhs = output matrix multiplication > > Part (b) "Tensor Engine View" (right side) shows the hardware mapping: +> > - A green matrix labeled "lhs_T (Tensor Engine)" with dimensions M (lhs_fsize) width by K (lhs_psize) height - the left-hand side transposed and loaded into Tensor Engine as the stationary matrix > - A blue matrix labeled "rhs (SBUF)" with dimensions N (rhs_fsize) width by K (rhs_psize) height - the right-hand side stored in State Buffer as the moving matrix > - A purple matrix labeled "output (PSUM)" with dimensions N (rhs_fsize) width by M (lhs_fsize) height - the output accumulates in Partial Sum buffer @@ -32,13 +35,15 @@ matrix multiplication in NKI. > - A "Copy" arrow shows the PSUM output being copied to "output (SBUF)" with dimensions N width by M height > > Dimension annotations include: +> > - M (lhs_fsize): Free dimension of left-hand side > - N (rhs_fsize): Free dimension of right-hand side > - K (lhs_psize, rhs_psize): Contraction/partition dimension > - PSUM P-dim and SBUF P-dim labels > > **Key Elements:** -> - **Mathematical View (a)**: lhs * rhs = output multiplication +> +> - **Mathematical View (a)**: lhs \* rhs = output multiplication > - **Tensor Engine View (b)**: Hardware implementation view > - **lhs_T (Tensor Engine)**: Transposed left matrix held stationary > - **rhs (SBUF)**: Right matrix streaming from State Buffer @@ -47,7 +52,6 @@ matrix multiplication in NKI. > - **Copy**: Data transfer from PSUM to SBUF > - **Dimension labels**: M, N, K with fsize and psize annotations - Fig. 21 MxKxN Matrix Multiplication Visualization [Fig. 21](#nki-fig-mm-view) illustrates how a simple matrix @@ -61,7 +65,6 @@ The NKI example below implements a compute kernel for a single-tile matrix multiplication. It computes a `64(M) x 128(K) x 512 (N)` matrix multiplication operation. - ```python @nki.jit def nki_matmul_basic_(lhsT, rhs): @@ -107,7 +110,7 @@ def nki_matmul_basic_(lhsT, rhs): # Note: A NKI matmul instruction always writes to PSUM in float32 data-type nisa.nc_matmul(result_psum, lhs_tile, rhs_tile) - # Create a tensor in SBUF and copy the result from PSUM back to SBUF, + # Create a tensor in SBUF and copy the result from PSUM back to SBUF, # and cast to expected output data-type result_sbuf = nl.ndarray(result_psum.shape, dtype=result.dtype, buffer=nl.sbuf) nisa.tensor_copy(dst=result_sbuf, src=result_psum, dtype=result.dtype) @@ -119,37 +122,35 @@ def nki_matmul_basic_(lhsT, rhs): return result ``` - In this example, we define the NKI kernel as `nki_matmul_basic_:` -* We define indices to access the LHS and RHS input tensors. +- We define indices to access the LHS and RHS input tensors. -* To adhere to NKI’s layout considerations, -we map the contraction axis of both LHS and RHS to the P-dimension, -which means we load LHS in transposed form. +- To adhere to NKI’s layout considerations, + we map the contraction axis of both LHS and RHS to the P-dimension, + which means we load LHS in transposed form. -* To adhere to NKI’s tile size considerations, -we limit the matmul instruction arguments to tiles of up to -`[128,128]` for LHS, and `[128,512]` for RHS. +- To adhere to NKI’s tile size considerations, + we limit the matmul instruction arguments to tiles of up to + `[128,128]` for LHS, and `[128,512]` for RHS. -* Using the `nisa.dma_copy` operation, we load the inputs from HBM tensors -to SBUF tiles. +- Using the `nisa.dma_copy` operation, we load the inputs from HBM tensors + to SBUF tiles. -* We then use the `nisa.nc_matmul` operation to perform the matrix -multiplication. Note that we set the LHS argument is transposed. Also note that the *64x128* -dimension here actually under-utilizes the TensorE, but it helps to -distinguish the M, K and N dimensions for education purposes in this first -code example. +- We then use the `nisa.nc_matmul` operation to perform the matrix + multiplication. Note that we set the LHS argument is transposed. Also note that the _64x128_ + dimension here actually under-utilizes the TensorE, but it helps to + distinguish the M, K and N dimensions for education purposes in this first + code example. -* `nisa.nc_matmul` always writes its result to PSUM, and since -`nisa.dma_copy` only moves data from SBUF to HBM, we copy the -multiplication result from PSUM back to SBUF using `nisa.tensor_copy`. +- `nisa.nc_matmul` always writes its result to PSUM, and since + `nisa.dma_copy` only moves data from SBUF to HBM, we copy the + multiplication result from PSUM back to SBUF using `nisa.tensor_copy`. We can then execute the kernel and verify correctness against the torch implementation as follows. Note that we use torch.allclose to tolerate numerical error inherent to floating-point arithmetic. - ```python device = xm.xla_device() cpu = torch.device('cpu') @@ -172,7 +173,6 @@ else: print("NKI and Torch differ") ``` - ## Tiling matrix multiplications So far, we’ve limited our matrix multiplication to the tile sizes @@ -182,7 +182,6 @@ for tiling an `[M,K] @ [K,N]` matrix-multiplication. Note that we assume the left-hand-side matrix (`[M,K]`) is already transposed to LHS_T (`[K,M]`) for optimal performance of the underlying TensorE. - ```python # LHS_T: left-hand-side matmul argument (shape [K,M]) # RHS: right-hand-side matmul argument (shape [K,N]) @@ -202,10 +201,8 @@ for m in range(0, M, 128): RES[m : m+128, n : n+512] = accum ``` - This form of tiling can be achieved in NKI as follows: - ```python @nki.jit def nki_matmul_tiled_(lhsT, rhs): @@ -258,7 +255,7 @@ def nki_matmul_tiled_(lhsT, rhs): nisa.dma_copy(dst=lhsT_tile, src=lhsT[k * TILE_K:(k + 1) * TILE_K, m * TILE_M:(m + 1) * TILE_M]) - nisa.dma_copy(dst=rhs_tile, + nisa.dma_copy(dst=rhs_tile, src=rhs[k * TILE_K:(k + 1) * TILE_K, n * TILE_N:(n + 1) * TILE_N]) @@ -277,10 +274,8 @@ def nki_matmul_tiled_(lhsT, rhs): return result ``` - A few notes about the above code example: - ```python psum_buf = nl.ndarray(..., buffer=nl.psum) @@ -290,7 +285,6 @@ for i in range(N): nisa.nc_matmul(psum_buf, stationary_tile, moving_tile) # or nl.matmul ``` - The use of [PSUM accumulation architecture feature](../../architecture/trainium_inferentia2_arch.md#arch-sec-accumulation-psum) is critical to achieve good performance out of TensorEngine when the contraction dimension of the matmul is greater than 128. @@ -342,6 +336,7 @@ iterations of the inner loop. The following example reduces these redundant loads through hoisting them out of the innermost loop. ! + > **Figure: mm memory pattern after load hoisting** > > A diagram showing the memory access pattern after load hoisting optimization for matrix multiplication, with labeled tiles (LHS tile_00, RHS tile_00, Result tile_00) and specific dimension annotations (128, 512). @@ -351,6 +346,7 @@ loads through hoisting them out of the innermost loop. > The left matrix has dimensions M (height) by K (width), displayed as a 6x5 grid. The upper-left tile is highlighted in solid orange and labeled "LHS tile_00" with dimensions 128 (width) by 128 (height), representing a square tile of the left-hand side operand. > > The middle matrix has dimensions K (height) by N (width), displayed as a 5x6 grid. Two regions are highlighted: +> > - A column labeled "RHS tile_00" in solid orange on the left side with dimensions 512 (width) by 128 (height) > - Adjacent light blue columns showing additional tiles that will be reused > @@ -359,6 +355,7 @@ loads through hoisting them out of the innermost loop. > The right matrix has dimensions M (height) by N (width), displayed as a 6x6 grid. The upper-left tile is highlighted in light blue and labeled "Result tile_00" with dimensions 512 (width) by 128 (height). > > This pattern shows the load hoisting optimization where: +> > - LHS tiles are loaded and reused across multiple RHS tiles > - RHS tiles share the K dimension with LHS > - Result tiles are larger in the N dimension due to accumulating multiple partial products @@ -366,6 +363,7 @@ loads through hoisting them out of the innermost loop. > The specific dimensions (128, 512) suggest typical tile sizes for NeuronCore Tensor Engine operations. > > **Key Elements:** +> > - **LHS tile_00**: Left operand tile (128 x 128) in orange > - **RHS tile_00**: Right operand tile (128 x 512) in orange > - **Result tile_00**: Output tile (128 x 512) in light blue @@ -373,10 +371,8 @@ loads through hoisting them out of the innermost loop. > - **128, 512 dimensions**: Specific tile sizes in elements > - **Light blue columns**: Additional RHS tiles showing reuse pattern - Fig. 23 Memory Pattern After Hoisting Loads Out of the Innermost Loop - ```python @nki.jit def nki_matmul_hoist_load_(lhsT, rhs): @@ -424,7 +420,7 @@ def nki_matmul_hoist_load_(lhsT, rhs): # Allocate space in SBUF for the tile (uninitialized) lhsT_tile = nl.ndarray(shape=(TILE_K, TILE_M), dtype=lhsT.dtype, buffer=nl.sbuf) # Copy the tile from HBM to SBUF - nisa.dma_copy(dst=lhsT_tile, + nisa.dma_copy(dst=lhsT_tile, src=lhsT[k * TILE_K:(k + 1) * TILE_K, m * TILE_M:(m + 1) * TILE_M]) # Append the tile to the list of tiles. @@ -461,7 +457,6 @@ def nki_matmul_hoist_load_(lhsT, rhs): return result ``` - ## Optimization 2: Blocking M and N Dimension While hoisting the load out of the innermost loop eliminates some redundant @@ -480,6 +475,7 @@ small enough for all live blocks remain within SBUF capacity to avoid spilling, after blocking both free dimensions. ! + > **Figure: mm memory pattern after blocking free** > > A diagram showing the memory access pattern after blocking only the free dimension for matrix multiplication, with highlighted rows/columns showing the tiles accessed for each matrix operand and output. @@ -493,6 +489,7 @@ after blocking both free dimensions. > The right matrix has dimensions M (height) by N (width), displayed as a 6x6 grid. A 2x2 block in the upper-middle area is highlighted in solid orange, representing the output tile being computed. > > This pattern shows blocking along the free dimensions (M for left matrix, N for right matrix) while iterating over the full contraction dimension K. The highlighted regions demonstrate: +> > - Full rows of the left matrix are loaded (M blocked, K unblocked) > - Full columns of the right matrix are loaded (K unblocked, N blocked) > - A small tile of output is produced @@ -500,6 +497,7 @@ after blocking both free dimensions. > This approach reduces output memory traffic but requires loading more input data per output tile compared to blocking all dimensions. > > **Key Elements:** +> > - **Left matrix (M x K)**: First operand with top 2 rows fully highlighted > - **Middle matrix (K x N)**: Second operand with 2 middle columns fully highlighted > - **Right matrix (M x N)**: Output matrix with 2x2 tile highlighted @@ -507,10 +505,8 @@ after blocking both free dimensions. > - **Orange highlighting**: Tiles accessed in this computation step > - **Full row/column access**: Shows K dimension not blocked - Fig. 24 Memory Pattern After Blocking Free Dimensions - ```python @nki.jit def nki_matmul_block_free_dimension_(lhsT, rhs): @@ -608,7 +604,7 @@ def nki_matmul_block_free_dimension_(lhsT, rhs): nisa.nc_matmul(dst=result_tile, stationary=lhsT_tiles[bm][k], moving=rhs_tiles[bn][k]) - + # Copy the result from PSUM back to SBUF, and cast to expected # output data-type result_tmp = nl.ndarray(shape=result_tile.shape, @@ -628,7 +624,6 @@ def nki_matmul_block_free_dimension_(lhsT, rhs): return result ``` - ## Optimization 3: Blocking M, N and K Dimension Blocking only free dimension and requiring to load the whole partition dimension (K) will set an upper @@ -636,7 +631,7 @@ limit on block size (M and N) due to limited SBUF capacity. Matrix multiply with shapes [M, K] @ [K, N] = [M, N] requires K multiplies and K additions (or K-1 for accumulation) for each element in resulting [M, N] grid, totaling 2*K*M*N FLOPS. -It has to load M*K + K*N + M*N elements, resulting in arithemtic intensity 2*M*N*K/(2*(M*K + K*N + M*N)) +It has to load M*K + K*N + M*N elements, resulting in arithemtic intensity 2*M*N*K/(2*(M*K + K*N + M\*N)) for 2 byte data type like FP16 or BF16. Since the full K has to fit in memory for optimization 2, it will limit M and N size for a block. Arithmetic intensity will be lower any of the M, N or K is much smaller than the others. @@ -645,6 +640,7 @@ Blocking partition dimension also results in calculating partial matrix multipli be accumulated, resulting in addintional HBM traffic if not handled carefully. ! + > **Figure: mm memory pattern after blocking all** > > A diagram showing the memory access pattern after blocking optimization for matrix multiplication, with three matrices (M x K, K x N, and M x N) where blocked tiles are highlighted in solid orange and dotted orange patterns. @@ -662,6 +658,7 @@ be accumulated, resulting in addintional HBM traffic if not handled carefully. > The highlighted regions show how blocking divides the computation into smaller tiles that fit in on-chip memory, improving data locality and reducing HBM bandwidth requirements. > > **Key Elements:** +> > - **Left matrix (M x K)**: First operand with 2x2 solid orange tile highlighted > - **Middle matrix (K x N)**: Second operand with 2x2 solid orange tile highlighted > - **Right matrix (M x N)**: Output matrix with 3x3 dotted orange tile @@ -670,7 +667,6 @@ be accumulated, resulting in addintional HBM traffic if not handled carefully. > - **Dotted orange tile**: Output tile being accumulated > - **Grid structure**: Shows tiling/blocking boundaries - Fig. 25 Memory Pattern After Blocking All Dimensions With the blocking configuration in the code (16 tiles or 2048 numbers in the @@ -684,7 +680,7 @@ SBUF limit as much as possible. With all matrices in BF16 data type, the `lhsT_tiles` requires 4MB and `rhs_tiles` requires 2MB SBUF memory. The `result_tiles` requires `4 * NUM_BLOCK_M` MB SBUF memory, where `NUM_BLOCK_M` is `M // 2048`. Thus, as long as `M <= 8192`, the required -SBUF memory is under the 24 MB budget (4 + 2 + 4 * (8192 // 2048) == 22 MB). +SBUF memory is under the 24 MB budget (4 + 2 + 4 \* (8192 // 2048) == 22 MB). When the `M` dimension becomes bigger, spilling and reloading of the `result_tiles` will happen, but because the frequency is relatively low, the computation can still be sufficient. @@ -693,7 +689,6 @@ small enough for all live blocks remain within SBUF capacity to avoid spilling, The K blocking loop is hand optimized for our ideal data locality. - ```python @nki.jit def nki_matmul_fully_optimized_( @@ -849,13 +844,11 @@ def nki_matmul_fully_optimized_( return result ``` - ## Testing Correctness and Benchmarking To test the correctness of the kernels, we compare the result with the `torch.matmul` with `torch.allclose`. - ```python # Test the large workload with tiled kernels lhs = torch.rand((4096, 1024), dtype=torch.bfloat16, device=device) @@ -885,10 +878,8 @@ print("Checking correctness of nki_matmul_fully_optimized") check_match(nki_matmul_fully_optimized_) ``` - Output from the test: - ```python Checking correctness of nki_matmul_tiled NKI and Torch match @@ -900,15 +891,14 @@ Checking correctness of nki_matmul_fully_optimized NKI and Torch match ``` - ## Download All Source Code Click the links to download source code of the kernels and the testing code discussed in this tutorial. -* All matrix multiplication NKI kernels: [`matrix_multiplication_nki_kernels.py`](../../downloads/matrix_multiplication_nki_kernels.py) +- All matrix multiplication NKI kernels: [`matrix_multiplication_nki_kernels.py`](../../downloads/matrix_multiplication_nki_kernels.py) -* PyTorch implementation: [`matrix_multiplication_torch.py`](../../downloads/matrix_multiplication_torch.py) +- PyTorch implementation: [`matrix_multiplication_torch.py`](../../downloads/matrix_multiplication_torch.py) You can also view the source code in the GitHub repository [nki_samples](https://github.com/aws-neuron/nki-samples/tree/main/src/nki_samples/tutorials/matrix_multiplication/) @@ -916,16 +906,13 @@ You can also view the source code in the GitHub repository [nki_samples](https:/ Run benchmarking of different NKI kernels: - ```python python3 matrix_multiplication_nki_kernels.py ``` - Run PyTorch implementation to validate the NKI results against the PyTorch implementation: - ```python python3 matrix_multiplication_torch.py -``` \ No newline at end of file +``` diff --git a/skills/neuron-nki-docs/references/programming/tutorials/spmd_multiple_nc_tensor_addition.md b/skills/neuron-nki-docs/references/programming/tutorials/spmd_multiple_nc_tensor_addition.md index b9646ac..59311b1 100644 --- a/skills/neuron-nki-docs/references/programming/tutorials/spmd_multiple_nc_tensor_addition.md +++ b/skills/neuron-nki-docs/references/programming/tutorials/spmd_multiple_nc_tensor_addition.md @@ -6,9 +6,9 @@ but directly control how our kernels and tensors are distributed across multiple Doing so, we expand our knowledge about: -* The NKI syntax and the [Logical Neuron Cores (LNC)](../lnc.md#nki-about-lnc). +- The NKI syntax and the [Logical Neuron Cores (LNC)](../lnc.md#nki-about-lnc). -* nki.language.spmd_dim() and nki.language.nc() +- nki.language.spmd_dim() and nki.language.nc() ## PyTorch @@ -20,7 +20,6 @@ The partition dimension tile size is chosen according to the tile size restrictions (nki.language.tile_size.pmax), while the free dimension tile size is chosen arbitrarily (`512`). - ```python def nki_tensor_add_nc2(a_input, b_input): """NKI kernel caller to compute element-wise addition of two input tensors using multiple Neuron cores. @@ -38,7 +37,7 @@ def nki_tensor_add_nc2(a_input, b_input): # The SPMD launch grid denotes the number of kernel instances. # In this case, we use a 2D grid where the size of each invocation is 128x512 - # Since we're sharding across neuron cores on the 1st dimension we want to do our slicing at + # Since we're sharding across neuron cores on the 1st dimension we want to do our slicing at # 128 per core * 2 cores = 256 grid_x = a_input.shape[0] // (128 * 2) grid_y = a_input.shape[1] // 512 @@ -53,42 +52,41 @@ def nki_tensor_add_nc2(a_input, b_input): return nki_tensor_add_kernel_[nl.spmd_dim(grid_x, nl.nc(2)), grid_y](a_input, b_input) ``` - In this example: -* We reuse the NKI kernel in `nki_tensor_add_kernel_` which is decorated with the -nki.jit decorator to call the nki compiler to compile the kernel. +- We reuse the NKI kernel in `nki_tensor_add_kernel_` which is decorated with the + nki.jit decorator to call the nki compiler to compile the kernel. -* Recall this kernel defines offsets into the tensors based on the ID of -the worker executing the code (`nl.program_id`), and generates tile -indices using these offsets with `nl.arange`. +- Recall this kernel defines offsets into the tensors based on the ID of + the worker executing the code (`nl.program_id`), and generates tile + indices using these offsets with `nl.arange`. -* Using SPMD execution as discussed in [Logical Neuron Cores (LNC)](../lnc.md#nki-about-lnc), -note that each worker only operates on a (sub-tensor) tile of the -input/output tensors. By accessing its own `program_id`, each -worker can calculate the offsets it needs to access the correct -tiles. +- Using SPMD execution as discussed in [Logical Neuron Cores (LNC)](../lnc.md#nki-about-lnc), + note that each worker only operates on a (sub-tensor) tile of the + input/output tensors. By accessing its own `program_id`, each + worker can calculate the offsets it needs to access the correct + tiles. -* When multiple Neuron Cores are specified in the SPMD launch grid, these tensors are further -sharded across available cores. On Trainium 2, we have 2 local cores that have shared HBM. +- When multiple Neuron Cores are specified in the SPMD launch grid, these tensors are further + sharded across available cores. On Trainium 2, we have 2 local cores that have shared HBM. -* As before, the first axis of the tensor (mapped to the partition-dimension) is -tiled into blocks of 128, based on hardware restrictions (see [Tile -Size Considerations](../tiling-overview.md#nki-about-tiling)). -The second axis (mapped to the free-dimension) is tiled into blocks of 512 (no tile-size constraint, -since the addition operation is performed on the Vector engine, the only restriction is on-chip memory capacity). +- As before, the first axis of the tensor (mapped to the partition-dimension) is + tiled into blocks of 128, based on hardware restrictions (see [Tile + Size Considerations](../tiling-overview.md#nki-about-tiling)). + The second axis (mapped to the free-dimension) is tiled into blocks of 512 (no tile-size constraint, + since the addition operation is performed on the Vector engine, the only restriction is on-chip memory capacity). -* `nl.store` for kernels running on both cores will write to an `c_output` in -shared HBM, dramatically increasing the throughput of the computation. +- `nl.store` for kernels running on both cores will write to an `c_output` in + shared HBM, dramatically increasing the throughput of the computation. ### SPMD execution -* We want to shard the workload across 2 cores, so for every `nl.nc(2)` we determine our initial `axis=0` to be -`128` from the expected slice size in the kernel `*` the number of cores `= 256`. +- We want to shard the workload across 2 cores, so for every `nl.nc(2)` we determine our initial `axis=0` to be + `128` from the expected slice size in the kernel `*` the number of cores `= 256`. -* Thus we alter our previous sample and change `grid_x` to `a_input.shape[0] // (128 * 2)` to account for this. +- Thus we alter our previous sample and change `grid_x` to `a_input.shape[0] // (128 * 2)` to account for this. -* Launch the kernel with launch grid `[nl.spmd_dim(grid_x, nl.nc(2)), grid_y]` +- Launch the kernel with launch grid `[nl.spmd_dim(grid_x, nl.nc(2)), grid_y]` As before, we are using a two-dimensional grid where the first dimension of the tensor is tiled in the X dimension of the grid while the second @@ -99,20 +97,17 @@ so we do not need to handle partial tiles. However, this time we also directly specify how each instance of our kernel will be distributed across multiple local Neuron Cores such that: - ```python # Physical NC [0]: kernel[n, m] where n is 0 or even # Physical NC [1]: kernel[n, m] where n is odd ``` - ### Launching kernel and testing correctness To execute the kernel, we prepare tensors `a` and `b`, and call the `nki_tensor_add_nc2` helper function. We also verify the correctness of the NKI kernel against, torch by comparing the outputs of both, using `torch.allclose`: - ```python import torch from torch_xla.core import xla_model as xm @@ -138,10 +133,8 @@ if __name__ == "__main__": assert allclose ``` - Output: - ```python 2023-12-29 15:18:00.000558: 14283 INFO ||NEURON_CACHE||: Compile cache path: /var/tmp/neuron-compile-cache 2023-12-29 15:18:00.000559: 14283 INFO ||NEURON_CC_WRAPPER||: Call compiler with cmd: ['neuronx-cc', '--target=trn1', 'compile', '--framework', 'XLA', '/tmp/neuroncc_compile_workdir/49f554a2-2c55-4a88-8054-cc9f20824a46/model.MODULE_5007921933048625946+d41d8cd9.hlo.pb', '--output', '/tmp/neuroncc_compile_workdir/49f554a2-2c55-4a88-8054-cc9f20824a46/model.MODULE_5007921933048625946+d41d8cd9.neff', '--verbose=35'] @@ -174,7 +167,6 @@ Compiler status PASS NKI and Torch match ``` - Note that the tensor values you see will differ from what’s printed above, since this example uses torch.rand to initialize the inputs. @@ -185,7 +177,6 @@ above, since this example uses torch.rand to initialize the inputs. We can reuse the same NKI compute kernel defined for PyTorch above and declare a helper function to launch the compute-kernel with appropriate grid/block sizes, to perform the computation: - ```python def nki_tensor_add_nc2(a_input, b_input): """NKI kernel caller to compute element-wise addition of two input tensors using multiple Neuron cores. @@ -203,7 +194,7 @@ def nki_tensor_add_nc2(a_input, b_input): # The SPMD launch grid denotes the number of kernel instances. # In this case, we use a 2D grid where the size of each invocation is 128x512 - # Since we're sharding across neuron cores on the 1st dimension we want to do our slicing at + # Since we're sharding across neuron cores on the 1st dimension we want to do our slicing at # 128 per core * 2 cores = 256 grid_x = a_input.shape[0] // (128 * 2) grid_y = a_input.shape[1] // 512 @@ -218,7 +209,6 @@ def nki_tensor_add_nc2(a_input, b_input): return nki_tensor_add_kernel_[nl.spmd_dim(grid_x, nl.nc(2)), grid_y](a_input, b_input) ``` - As before, we are using a two-dimensional grid where the first dimension of the tensor is tiled in the X dimension of the grid, while the second dimension is tiled in the Y dimension of the grid. We similarly @@ -228,20 +218,17 @@ so we do not need to handle partial tiles. However, this time we also directly specify how each instance of our kernel will be distributed across multiple local Neuron Cores such that: - ```python # Physical NC [0]: kernel[n, m] where n is 0 or even # Physical NC [1]: kernel[n, m] where n is odd ``` - ### Launching kernel and testing correctness To execute the kernel, we prepare arrays `a` and `b`, and call the `nki_tensor_add_nc2` helper function. We also verify the correctness of the NKI kernel against, JAX by comparing the outputs of both, using `jax.numpy.allclose`: - ```python import jax import jax.numpy as jnp @@ -267,10 +254,8 @@ if __name__ == "__main__": assert allclose ``` - Output: - ```python . Compiler status PASS @@ -299,7 +284,6 @@ Compiler status PASS NKI and JAX match ``` - Note that the array values you see will differ from what’s printed above, since this example uses jax.random.uniform to initialize the inputs. @@ -308,21 +292,18 @@ above, since this example uses jax.random.uniform to initialize the inputs. Click the links to download source code of the kernels and the testing code discussed in this tutorial. -* -NKI baremetal implementation: [`spmd_multiple_nc_tensor_addition_nki_kernels.py`](../../downloads/spmd_multiple_nc_tensor_addition_nki_kernels.py) +- NKI baremetal implementation: [`spmd_multiple_nc_tensor_addition_nki_kernels.py`](../../downloads/spmd_multiple_nc_tensor_addition_nki_kernels.py) You must also download [`spmd_tensor_addition_nki_kernels.py`](../../downloads/spmd_tensor_addition_nki_kernels.py) into the same folder to run this script. -* -PyTorch implementation: [`spmd_multiple_nc_tensor_addition_torch.py`](../../downloads/spmd_multiple_nc_tensor_addition_torch.py) +- PyTorch implementation: [`spmd_multiple_nc_tensor_addition_torch.py`](../../downloads/spmd_multiple_nc_tensor_addition_torch.py) You must also download [`spmd_multiple_nc_tensor_addition_nki_kernels.py`](../../downloads/spmd_multiple_nc_tensor_addition_nki_kernels.py) and [`spmd_tensor_addition_nki_kernels.py`](../../downloads/spmd_tensor_addition_nki_kernels.py) into the same folder to run this PyTorch script. -* -JAX implementation: [`spmd_multiple_nc_tensor_addition_jax.py`](../../downloads/spmd_multiple_nc_tensor_addition_jax.py) +- JAX implementation: [`spmd_multiple_nc_tensor_addition_jax.py`](../../downloads/spmd_multiple_nc_tensor_addition_jax.py) You must also download [`spmd_multiple_nc_tensor_addition_nki_kernels.py`](../../downloads/spmd_multiple_nc_tensor_addition_nki_kernels.py) and [`spmd_tensor_addition_nki_kernels.py`](../../downloads/spmd_tensor_addition_nki_kernels.py) @@ -334,23 +315,18 @@ You can also view the source code in the GitHub repository [nki_samples](https:/ Run NKI baremetal implementation: - ```python python3 spmd_multiple_nc_tensor_addition_nki_kernels.py ``` - Run PyTorch implementation: - ```python python3 spmd_multiple_nc_tensor_addition_torch.py ``` - Run JAX implementation: - ```python python3 spmd_multiple_nc_tensor_addition_jax.py -``` \ No newline at end of file +``` diff --git a/skills/neuron-nki-docs/references/programming/tutorials/spmd_tensor_addition.md b/skills/neuron-nki-docs/references/programming/tutorials/spmd_tensor_addition.md index 3998658..37f757a 100644 --- a/skills/neuron-nki-docs/references/programming/tutorials/spmd_tensor_addition.md +++ b/skills/neuron-nki-docs/references/programming/tutorials/spmd_tensor_addition.md @@ -4,10 +4,10 @@ Single Program, Multiple Data (SPMD) Tensor Addition In this tutorial we write a simple tensor addition kernel using NKI in PyTorch and JAX. In doing so, we learn about: -* The NKI syntax and [Logical Neuron Cores (LNC)](../lnc.md#nki-about-lnc). +- The NKI syntax and [Logical Neuron Cores (LNC)](../lnc.md#nki-about-lnc). -* Best practices for validating and benchmarking your custom kernel -against a reference native PyTorch or JAX implementation. +- Best practices for validating and benchmarking your custom kernel + against a reference native PyTorch or JAX implementation. ## PyTorch @@ -19,7 +19,6 @@ The partition dimension tile size is chosen according to the tile size restrictions (nki.language.tile_size.pmax), while the free dimension tile size is chosen arbitrarily (`512`). - ```python import nki import nki.isa as nisa @@ -65,51 +64,49 @@ def nki_tensor_add_kernel_(a_input, b_input): return c_output ``` - In this example: -* We define the NKI kernel in `nki_tensor_add_kernel_`, decorate it with the -nki.jit decorator to call the nki compiler to compile the kernel. +- We define the NKI kernel in `nki_tensor_add_kernel_`, decorate it with the + nki.jit decorator to call the nki compiler to compile the kernel. -* Inside, we first allocate tensor `c_output` as the result of the kernel +- Inside, we first allocate tensor `c_output` as the result of the kernel -* Next, we define offsets into the tensors, based on the ID of -the worker executing the code (`nl.program_id`). We allocate tiles -in on-chip memory (SBUF) using `nl.ndarray` and use direct slicing -to load data. See NKI Programming Model for more information on -different tensor indexing modes. +- Next, we define offsets into the tensors, based on the ID of + the worker executing the code (`nl.program_id`). We allocate tiles + in on-chip memory (SBUF) using `nl.ndarray` and use direct slicing + to load data. See NKI Programming Model for more information on + different tensor indexing modes. -* We use `nl.program_id` to enable SPMD execution (single-program, -multiple-data, see [Logical Neuron Cores (LNC)](../lnc.md#nki-about-lnc)), -where each worker only operates on a (sub-tensor) tile of the -input/output tensors. By accessing its own `program_id`, each -worker can calculate the offsets it needs to access the correct -tiles. +- We use `nl.program_id` to enable SPMD execution (single-program, + multiple-data, see [Logical Neuron Cores (LNC)](../lnc.md#nki-about-lnc)), + where each worker only operates on a (sub-tensor) tile of the + input/output tensors. By accessing its own `program_id`, each + worker can calculate the offsets it needs to access the correct + tiles. -* The first axis of the tensor (mapped to the partition-dimension) is -tiled into blocks of 128, based on hardware restrictions (see [Tile -Size Considerations](../tiling-overview.md#nki-tile-size)). -The second axis (mapped to the free-dimension) is tiled into blocks of 512 (no tile-size constraint, -since the addition operation is performed on the Vector engine, the only restriction is on-chip memory capacity). +- The first axis of the tensor (mapped to the partition-dimension) is + tiled into blocks of 128, based on hardware restrictions (see [Tile + Size Considerations](../tiling-overview.md#nki-tile-size)). + The second axis (mapped to the free-dimension) is tiled into blocks of 512 (no tile-size constraint, + since the addition operation is performed on the Vector engine, the only restriction is on-chip memory capacity). -* We then load sub-tensors data from tensors `a_input` and -`b_input` using `nisa.dma_copy`, to place the tiles `a_tile` and -`b_tile` in the on-chip memory (SBUF) +- We then load sub-tensors data from tensors `a_input` and + `b_input` using `nisa.dma_copy`, to place the tiles `a_tile` and + `b_tile` in the on-chip memory (SBUF) -* We sum them using `nisa.tensor_tensor` to compute `c_tile`, and store it back to DRAM in the -relevant portion of the `c_output` tensor, using `nisa.dma_copy`. -Since both inputs and output are the same shape, we can use the same -set of indices to access all three tensors. +- We sum them using `nisa.tensor_tensor` to compute `c_tile`, and store it back to DRAM in the + relevant portion of the `c_output` tensor, using `nisa.dma_copy`. + Since both inputs and output are the same shape, we can use the same + set of indices to access all three tensors. -* At the end, we use `return` statement to transfer the ownership of -tensor `c_output` to the caller of the kernel. +- At the end, we use `return` statement to transfer the ownership of + tensor `c_output` to the caller of the kernel. ### SPMD execution We declare a helper function, to launch the compute-kernel with appropriate grid/block sizes, to perform the computation over the whole input tensors. - ```python def nki_tensor_add(a_input, b_input): """NKI kernel caller to compute element-wise addition of two input tensors @@ -132,7 +129,6 @@ def nki_tensor_add(a_input, b_input): return nki_tensor_add_kernel_[grid_x, grid_y](a_input, b_input) ``` - We are using a two-dimensional grid, where the first dimension of the tensor is tiled in the X dimension of the grid, while the second dimension is tiled in the Y dimension of the grid. In this scenario we @@ -145,7 +141,6 @@ To execute the kernel, we prepare tensors `a` and `b`, and call the `nki_tensor_add` helper function. We also verify the correctness of the NKI kernel against, torch by comparing the outputs of both, using `torch.allclose`: - ```python import torch from torch_xla.core import xla_model as xm @@ -171,10 +166,8 @@ if __name__ == "__main__": assert allclose ``` - Output: - ```python 2023-12-29 15:18:00.000558: 14283 INFO ||NEURON_CACHE||: Compile cache path: /var/tmp/neuron-compile-cache 2023-12-29 15:18:00.000559: 14283 INFO ||NEURON_CC_WRAPPER||: Call compiler with cmd: ['neuronx-cc', '--target=trn1', 'compile', '--framework', 'XLA', '/tmp/neuroncc_compile_workdir/49f554a2-2c55-4a88-8054-cc9f20824a46/model.MODULE_5007921933048625946+d41d8cd9.hlo.pb', '--output', '/tmp/neuroncc_compile_workdir/49f554a2-2c55-4a88-8054-cc9f20824a46/model.MODULE_5007921933048625946+d41d8cd9.neff', '--verbose=35'] @@ -207,7 +200,6 @@ Compiler status PASS NKI and Torch match ``` - Note that the tensor values you see will differ from what’s printed above, since this example uses torch.rand to initialize the inputs. @@ -217,7 +209,6 @@ above, since this example uses torch.rand to initialize the inputs. We can reuse the same NKI compute kernel defined for PyTorch above. - ```python import nki import nki.isa as nisa @@ -263,13 +254,11 @@ def nki_tensor_add_kernel_(a_input, b_input): return c_output ``` - ### SPMD execution Now we can also declare a helper function, to launch the compute-kernel with appropriate grid/block sizes, to perform the computation: - ```python def nki_tensor_add(a_input, b_input): """NKI kernel caller to compute element-wise addition of two input tensors @@ -292,7 +281,6 @@ def nki_tensor_add(a_input, b_input): return nki_tensor_add_kernel_[grid_x, grid_y](a_input, b_input) ``` - We are using a two-dimensional grid, where the first dimension of the tensor is tiled in the X dimension of the grid, while the second dimension is tiled in the Y dimension of the grid. In this scenario we @@ -305,7 +293,6 @@ To execute the kernel, we prepare arrays `a` and `b`, and call the `nki_tensor_add` helper function. We also verify the correctness of the NKI kernel against, JAX by comparing the outputs of both, using `jax.numpy.allclose`: - ```python import jax import jax.numpy as jnp @@ -331,10 +318,8 @@ if __name__ == "__main__": assert allclose ``` - Output: - ```python . Compiler status PASS @@ -363,7 +348,6 @@ Compiler status PASS NKI and JAX match ``` - Note that the array values you see will differ from what’s printed above, since this example uses jax.random.uniform to initialize the inputs. @@ -372,16 +356,14 @@ above, since this example uses jax.random.uniform to initialize the inputs. Click the links to download source code of the kernels and the testing code discussed in this tutorial. -* NKI baremetal implementation: [`spmd_tensor_addition_nki_kernels.py`](../../downloads/spmd_tensor_addition_nki_kernels.py) +- NKI baremetal implementation: [`spmd_tensor_addition_nki_kernels.py`](../../downloads/spmd_tensor_addition_nki_kernels.py) -* -PyTorch implementation: [`spmd_tensor_addition_torch.py`](../../downloads/spmd_tensor_addition_torch.py) +- PyTorch implementation: [`spmd_tensor_addition_torch.py`](../../downloads/spmd_tensor_addition_torch.py) You must also download [`spmd_tensor_addition_nki_kernels.py`](../../downloads/spmd_tensor_addition_nki_kernels.py) into the same folder to run this PyTorch script. -* -JAX implementation: [`spmd_tensor_addition_jax.py`](../../downloads/spmd_tensor_addition_jax.py) +- JAX implementation: [`spmd_tensor_addition_jax.py`](../../downloads/spmd_tensor_addition_jax.py) You must also download [`spmd_tensor_addition_nki_kernels.py`](../../downloads/spmd_tensor_addition_nki_kernels.py) into the same folder to run this PyTorch script. @@ -392,23 +374,18 @@ You can also view the source code in the GitHub repository [nki_samples](https:/ Run NKI baremetal implementation: - ```python python3 spmd_tensor_addition_nki_kernels.py ``` - Run PyTorch implementation: - ```python python3 spmd_tensor_addition_torch.py ``` - Run JAX implementation: - ```python python3 spmd_tensor_addition_jax.py -``` \ No newline at end of file +``` diff --git a/skills/neuron-nki-docs/references/programming/tutorials/transpose2d.md b/skills/neuron-nki-docs/references/programming/tutorials/transpose2d.md index 427dc1c..101e988 100644 --- a/skills/neuron-nki-docs/references/programming/tutorials/transpose2d.md +++ b/skills/neuron-nki-docs/references/programming/tutorials/transpose2d.md @@ -4,19 +4,19 @@ Transpose2D In this tutorial, we transpose a tensor along two of its axes using NKI. In doing so, we learn about: -* The NKI syntax and programming model. +- The NKI syntax and programming model. -* Multi-dimensional memory address patterns in NKI. +- Multi-dimensional memory address patterns in NKI. As background, there are two main types of transposition in NKI: -* Transposition between the partition-dimension axis and one of the -free-dimension axes, which is achieved via the -`nki.isa.nc_transpose` instruction. +- Transposition between the partition-dimension axis and one of the + free-dimension axes, which is achieved via the + `nki.isa.nc_transpose` instruction. -* Transposition between two axes on the free-dimension, which is achieved -via a `nki.language.copy` instruction, with indexing manipulation -in the free axis to re-arrange the data. +- Transposition between two axes on the free-dimension, which is achieved + via a `nki.language.copy` instruction, with indexing manipulation + in the free axis to re-arrange the data. In this example, we’ll focus on the second case: consider a three-dimensional input tensor `[P, F1, F2]`, where the `P` axis is mapped @@ -35,7 +35,6 @@ Fig. 27 Tensor F1:F2 Transpose ### Compute kernel - ```python import nki import nki.isa as nisa @@ -102,12 +101,10 @@ def tensor_transpose2D_kernel_(in_tensor, shape2D): return out_tensor ``` - ### Launching kernel and testing correctness To execute the kernel, we prepare tensors `a` and call `tensor_transpose2D_kernel_`: - ```python import torch from torch_xla.core import xla_model as xm @@ -134,14 +131,12 @@ if __name__ == "__main__": assert allclose ``` - ## JAX ### Compute kernel We can reuse the same NKI compute kernel defined for PyTorch above. - ```python import nki import nki.isa as nisa @@ -208,12 +203,10 @@ def tensor_transpose2D_kernel_(in_tensor, shape2D): return out_tensor ``` - ### Launching kernel and testing correctness To execute the kernel, we prepare array `a` and call `tensor_transpose2D_kernel_`: - ```python import jax import jax.numpy as jnp @@ -235,12 +228,10 @@ if __name__ == "__main__": assert allclose ``` - > **Note** > > Note -> -> +> > We pass `shape2D` as kwargs to pass the shape as a compile-time constant > to the kernel function. @@ -249,16 +240,14 @@ if __name__ == "__main__": Click the links to download source code of the kernels and the testing code discussed in this tutorial. -* NKI baremetal implementation: [`transpose2d_nki_kernels.py`](../../downloads/transpose2d_nki_kernels.py) +- NKI baremetal implementation: [`transpose2d_nki_kernels.py`](../../downloads/transpose2d_nki_kernels.py) -* -PyTorch implementation: [`transpose2d_torch.py`](../../downloads/transpose2d_torch.py) +- PyTorch implementation: [`transpose2d_torch.py`](../../downloads/transpose2d_torch.py) You must also download [`transpose2d_nki_kernels.py`](../../downloads/transpose2d_nki_kernels.py) into the same folder to run this PyTorch script. -* -JAX implementation: [`transpose2d_jax.py`](../../downloads/transpose2d_jax.py) +- JAX implementation: [`transpose2d_jax.py`](../../downloads/transpose2d_jax.py) You must also download [`transpose2d_nki_kernels.py`](../../downloads/transpose2d_nki_kernels.py) into the same folder to run this JAX script. @@ -269,23 +258,18 @@ You can also view the source code in the GitHub repository [nki_samples](https:/ Run NKI baremetal implementation: - ```python python3 transpose2d_nki_kernels.py ``` - Run PyTorch implementation: - ```python python3 transpose2d_torch.py ``` - Run JAX implementation: - ```python python3 transpose2d_jax.py -``` \ No newline at end of file +``` diff --git a/skills/neuron-nki-docs/references/programming/tutorials/tutorials.md b/skills/neuron-nki-docs/references/programming/tutorials/tutorials.md index 30271bd..7722897 100644 --- a/skills/neuron-nki-docs/references/programming/tutorials/tutorials.md +++ b/skills/neuron-nki-docs/references/programming/tutorials/tutorials.md @@ -22,4 +22,3 @@ Create custom 2D average pooling kernels for computer vision workloads [Fused Mamba](fused_mamba.md) Implement fused Mamba state space model kernels - diff --git a/skills/neuron-nki-docs/references/reference/library/attention-cte.md b/skills/neuron-nki-docs/references/reference/library/attention-cte.md index 7f90ef7..9dbf743 100644 --- a/skills/neuron-nki-docs/references/reference/library/attention-cte.md +++ b/skills/neuron-nki-docs/references/reference/library/attention-cte.md @@ -5,21 +5,21 @@ This topic provides the API reference for the `Attention CTE` kernel. The kernel The kernel supports: -* Efficient attention computation for long sequence lengths +- Efficient attention computation for long sequence lengths -* Causal masking +- Causal masking -* Sliding window attention +- Sliding window attention -* Context parallelism for distributed computation +- Context parallelism for distributed computation -* Prefix caching for efficient inference +- Prefix caching for efficient inference -* Sink tokens for streaming attention +- Sink tokens for streaming attention -* Native Grouped Query Attention (GQA) support +- Native Grouped Query Attention (GQA) support -* Softmax caching for training +- Softmax caching for training ## Background @@ -33,46 +33,46 @@ The kernel employs efficient tiling strategies and memory access patterns to max ### attention_cte -nkilib.core.attention_cte.attention_cte(*q*, *k*, *v*, *scale=1.0*, *causal_mask=True*, *k_prior=None*, *v_prior=None*, *prior_used_len=None*, *sink=None*, *sliding_window=None*, *tp_q=True*, *tp_k=False*, *tp_out=False*, *cache_softmax=False*, *softmax_dtype=nl.float32*, *cp_offset=None*, *global_cp_deg=None*) +nkilib.core.attention*cte.attention_cte(\_q*, _k_, _v_, _scale=1.0_, _causal_mask=True_, _k_prior=None_, _v_prior=None_, _prior_used_len=None_, _sink=None_, _sliding_window=None_, _tp_q=True_, _tp_k=False_, _tp_out=False_, _cache_softmax=False_, _softmax_dtype=nl.float32_, _cp_offset=None_, _global_cp_deg=None_) Entrypoint NKI kernel that supports multiple attention variants. The kernel can be invoked with 1D SPMD grid for LNC2 or without grid. Parameters: -* **q** (`nl.ndarray`) – Query tensor with layout dependent on `tp_q` parameter +- **q** (`nl.ndarray`) – Query tensor with layout dependent on `tp_q` parameter -* **k** (`nl.ndarray`) – Key tensor with layout dependent on `tp_k` parameter +- **k** (`nl.ndarray`) – Key tensor with layout dependent on `tp_k` parameter -* **v** (`nl.ndarray`) – Value tensor with shape `(batch_size_kv, seqlen, d)` +- **v** (`nl.ndarray`) – Value tensor with shape `(batch_size_kv, seqlen, d)` -* **scale** (`float`, optional) – Scaling factor for attention scores. Must be 1.0 when using sliding window, context parallel, or prefix caching. +- **scale** (`float`, optional) – Scaling factor for attention scores. Must be 1.0 when using sliding window, context parallel, or prefix caching. -* **causal_mask** (`bool`, optional) – Whether to use causal mask +- **causal_mask** (`bool`, optional) – Whether to use causal mask -* **k_prior** (`nl.ndarray`, optional) – (Prefix caching) Prior key tensor with layout dependent on `tp_k` parameter +- **k_prior** (`nl.ndarray`, optional) – (Prefix caching) Prior key tensor with layout dependent on `tp_k` parameter -* **v_prior** (`nl.ndarray`, optional) – (Prefix caching) Prior value tensor with shape `(batch_size_kv, seqlen_prior, d)` +- **v_prior** (`nl.ndarray`, optional) – (Prefix caching) Prior value tensor with shape `(batch_size_kv, seqlen_prior, d)` -* **prior_used_len** (`nl.ndarray`, optional) – (Prefix caching) Actual used length in prior with shape `(1,)` +- **prior_used_len** (`nl.ndarray`, optional) – (Prefix caching) Actual used length in prior with shape `(1,)` -* **sink** (`nl.ndarray`, optional) – Sink token tensor +- **sink** (`nl.ndarray`, optional) – Sink token tensor -* **sliding_window** (`int`, optional) – Sliding window size for attention, `None` or `0` denotes no sliding window mask +- **sliding_window** (`int`, optional) – Sliding window size for attention, `None` or `0` denotes no sliding window mask -* **tp_q** (`bool`, optional) – Query tensor transpose flag +- **tp_q** (`bool`, optional) – Query tensor transpose flag -* **tp_k** (`bool`, optional) – Key tensor transpose flag +- **tp_k** (`bool`, optional) – Key tensor transpose flag -* **tp_out** (`bool`, optional) – Output tensor transpose flag +- **tp_out** (`bool`, optional) – Output tensor transpose flag -* **cache_softmax** (`bool`, optional) – Whether to cache softmax intermediate values +- **cache_softmax** (`bool`, optional) – Whether to cache softmax intermediate values -* **softmax_dtype** (`nl.dtype`, optional) – Data type for softmax computations +- **softmax_dtype** (`nl.dtype`, optional) – Data type for softmax computations -* **cp_offset** (`nl.ndarray`, optional) – Context parallel offset tensor +- **cp_offset** (`nl.ndarray`, optional) – Context parallel offset tensor -* **global_cp_deg** (`int`, optional) – Global context parallel degree +- **global_cp_deg** (`int`, optional) – Global context parallel degree Returns: Output tensor with attention results. Shape depends on `tp_out` parameter. If `cache_softmax` is `True`, returns tuple of `(output, out_neg_max, out_sum_recip)`. @@ -82,77 +82,77 @@ Return type: **IO Shapes**: -* q: -`(batch_size, seqlen_q, d)` when `tp_q` is `True` -`(batch_size, d, seqlen_q)` when `tp_q` is `False` +- q: + `(batch_size, seqlen_q, d)` when `tp_q` is `True` + `(batch_size, d, seqlen_q)` when `tp_q` is `False` -* k: -`(batch_size_kv, seqlen_kv, d)` when `tp_k` is `True` -`(batch_size_kv, d, seqlen_kv)` when `tp_k` is `False` +- k: + `(batch_size_kv, seqlen_kv, d)` when `tp_k` is `True` + `(batch_size_kv, d, seqlen_kv)` when `tp_k` is `False` -* v: `(batch_size_kv, seqlen_kv, d)` +- v: `(batch_size_kv, seqlen_kv, d)` -* returns: -`(batch_size, d, seqlen_q)` if `tp_out` is `True` -`(batch_size, seqlen_q, d)` if `tp_out` is `False` +- returns: + `(batch_size, d, seqlen_q)` if `tp_out` is `True` + `(batch_size, seqlen_q, d)` if `tp_out` is `False` **Constraints**: -* Head dimension (`d`) must be <= 128 +- Head dimension (`d`) must be <= 128 -* `scale` must be 1.0 when using sliding window, context parallel, or prefix caching +- `scale` must be 1.0 when using sliding window, context parallel, or prefix caching -* Context parallelism currently only supports causal attention +- Context parallelism currently only supports causal attention -* Sliding window attention currently only supports causal attention +- Sliding window attention currently only supports causal attention ## Features -* **Causal Masking (causal_mask=True)**: +- **Causal Masking (causal_mask=True)**: Masks upper triangle of attention scores: `S[i,j] = -inf` when `i < j` -* Enables compute skipping: skip MM1/MM2 for upper triangle tiles +- Enables compute skipping: skip MM1/MM2 for upper triangle tiles -* **Sliding Window Attention (SWA, when sliding_window > 0)**: +- **Sliding Window Attention (SWA, when sliding_window > 0)**: Local attention: each query only attends to nearby keys within a window -* Masks attention scores: `S[i,j] = -inf` when `|i - j| > sliding_window` +- Masks attention scores: `S[i,j] = -inf` when `|i - j| > sliding_window` -* Currently only works with causal: masks both upper triangle AND positions outside window +- Currently only works with causal: masks both upper triangle AND positions outside window -* When used with CP: loads only required KV slice to save memory +- When used with CP: loads only required KV slice to save memory -* **Context Parallelism (CP, global_cp_deg > 1, cp_offset != None)**: +- **Context Parallelism (CP, global_cp_deg > 1, cp_offset != None)**: Distributes long sequence computation across multiple devices/ranks -* Each rank (kernel call) processes a slice of Q sequence with full K/V +- Each rank (kernel call) processes a slice of Q sequence with full K/V -* `cp_offset` indicates which Q slice this rank handles (runtime value) +- `cp_offset` indicates which Q slice this rank handles (runtime value) -* Requires dynamic masking since offset unknown at compile time +- Requires dynamic masking since offset unknown at compile time -* Currently only supports causal attention +- Currently only supports causal attention -* **Prefix Caching (k_prior/v_prior provided)**: +- **Prefix Caching (k_prior/v_prior provided)**: K/V split into two parts: prior (cached) and active (current) -* `prior_used_len` specifies how much of prior to use (dynamic mask) +- `prior_used_len` specifies how much of prior to use (dynamic mask) -* Causal mask not required for prior portion (although SWA still applies if enabled) +- Causal mask not required for prior portion (although SWA still applies if enabled) -* **Sink Tokens (sink provided)**: +- **Sink Tokens (sink provided)**: Add additional sink token to softmax denominator -* **Grouped Query Attention (GQA, batch_size_kv < batch_size)**: +- **Grouped Query Attention (GQA, batch_size_kv < batch_size)**: Kernel handles GQA natively without explicit K/V replication -* **Support for training**: +- **Support for training**: Kernel can optionally return maximum attention score and softmax denominator (per row) for backpropagation @@ -160,52 +160,52 @@ Kernel can optionally return maximum attention score and softmax denominator (pe The kernel implementation includes several key optimizations: -* **LNC2 Sharding**: Shards computation across 2 NeuronCores with primary sharding on batch dimension and secondary sharding on sequence length for odd batch sizes. +- **LNC2 Sharding**: Shards computation across 2 NeuronCores with primary sharding on batch dimension and secondary sharding on sequence length for odd batch sizes. -* **Flash Attention**: For K/V length > 10K tokens, divides into 8K-token sections and processes one section at a time to fit in SBUF memory. +- **Flash Attention**: For K/V length > 10K tokens, divides into 8K-token sections and processes one section at a time to fit in SBUF memory. -* **Software Pipelining**: Overlaps operations across Q groups (`i`, `i+1`, `i+2`) for efficient hardware utilization: +- **Software Pipelining**: Overlaps operations across Q groups (`i`, `i+1`, `i+2`) for efficient hardware utilization: Group `i`: PV computation, writeback -* Group `i+1`: Exp computation +- Group `i+1`: Exp computation -* Group `i+2`: Q load, QK computation +- Group `i+2`: Q load, QK computation -* **Modular Allocation**: Uses efficient buffer reuse with modular allocation for intermediate tensors. +- **Modular Allocation**: Uses efficient buffer reuse with modular allocation for intermediate tensors. -* **Dynamic Masking**: Implements efficient masking strategies for causal, sliding window, and context parallel scenarios. +- **Dynamic Masking**: Implements efficient masking strategies for causal, sliding window, and context parallel scenarios. -* **Optimized Memory Access**: Employs careful memory access patterns to optimize data movement between HBM and SBUF. +- **Optimized Memory Access**: Employs careful memory access patterns to optimize data movement between HBM and SBUF. ## Algorithm The kernel goes through the following steps: -* **Setup**: Initialize intermediate buffers, mask, and debug tensors. +- **Setup**: Initialize intermediate buffers, mask, and debug tensors. -* **Loop over K/V sections**: For long sequences, divide K/V into sections of 8K tokens. +- **Loop over K/V sections**: For long sequences, divide K/V into sections of 8K tokens. -* **For each section**: +- **For each section**: Load K and V to SBUF -* Loop over Q (groups) - each group has seqlen 128 +- Loop over Q (groups) - each group has seqlen 128 -* Within each group: +- Within each group: Load Q -* Compute QK^T (MM1) and max +- Compute QK^T (MM1) and max -* Compute exponential and transpose +- Compute exponential and transpose -* Compute PV (MM2) +- Compute PV (MM2) -* Write to output +- Write to output -* **Flash Attention**: Maintain running statistics (max, sum) across sections and use these to update the output using flash attention rescaling. +- **Flash Attention**: Maintain running statistics (max, sum) across sections and use these to update the output using flash attention rescaling. ## See Also -* [Attention TKG Kernel API Reference](attention-tkg.md) \ No newline at end of file +- [Attention TKG Kernel API Reference](attention-tkg.md) diff --git a/skills/neuron-nki-docs/references/reference/library/attention-tkg.md b/skills/neuron-nki-docs/references/reference/library/attention-tkg.md index b737f03..ac8c0db 100644 --- a/skills/neuron-nki-docs/references/reference/library/attention-tkg.md +++ b/skills/neuron-nki-docs/references/reference/library/attention-tkg.md @@ -5,21 +5,21 @@ This topic provides the API reference for the `Attention TKG` kernel. The kernel The kernel supports: -* Efficient attention computation for small active sequence lengths +- Efficient attention computation for small active sequence lengths -* Flexible tensor placement in SBUF or HBM +- Flexible tensor placement in SBUF or HBM -* Adaptive LNC2 sharding strategies +- Adaptive LNC2 sharding strategies -* In-kernel mask generation +- In-kernel mask generation -* Fused RoPE (Rotary Position Embedding) +- Fused RoPE (Rotary Position Embedding) -* Block KV cache for efficient long-context inference +- Block KV cache for efficient long-context inference -* Attention sink for streaming attention +- Attention sink for streaming attention -* GPSIMD optimizations for inter-core communication +- GPSIMD optimizations for inter-core communication ## Background @@ -39,87 +39,87 @@ Configuration for token-generation attention kernel. This dataclass contains shape parameters and performance optimization flags for the attention_tkg kernel, which is optimized for small active sequence lengths. -bs*: [int](https://docs.python.org/3/library/functions.html#int)** = 0* +bs*: [int](https://docs.python.org/3/library/functions.html#int)\*\* = 0* Batch size -q_head*: [int](https://docs.python.org/3/library/functions.html#int)** = 0* +q_head*: [int](https://docs.python.org/3/library/functions.html#int)\*\* = 0* Number of query heads -s_active*: [int](https://docs.python.org/3/library/functions.html#int)** = 0* +s_active*: [int](https://docs.python.org/3/library/functions.html#int)\*\* = 0* Active sequence length (>1 means speculative decoding) -curr_sprior*: [int](https://docs.python.org/3/library/functions.html#int)** = 0* +curr_sprior*: [int](https://docs.python.org/3/library/functions.html#int)\*\* = 0* Current prior sequence length (KV cache length for this execution) -full_sprior*: [int](https://docs.python.org/3/library/functions.html#int)** = 0* +full_sprior*: [int](https://docs.python.org/3/library/functions.html#int)\*\* = 0* Full prior sequence length (maximum KV cache capacity) -d_head*: [int](https://docs.python.org/3/library/functions.html#int)** = 0* +d_head*: [int](https://docs.python.org/3/library/functions.html#int)\*\* = 0* Head dimension (embedding size per head) -block_len*: [int](https://docs.python.org/3/library/functions.html#int)** = 0* +block_len*: [int](https://docs.python.org/3/library/functions.html#int)\*\* = 0* Block length for block KV cache (0 if not using block KV) -tp_k_prior*: [bool](https://docs.python.org/3/library/functions.html#bool)** = False* +tp_k_prior*: [bool](https://docs.python.org/3/library/functions.html#bool)\*\* = False* Specifies that k_prior is transposed (shape `[B, 1, d, s_prior]` instead of `[B, 1, s_prior, d]`) -strided_mm1*: [bool](https://docs.python.org/3/library/functions.html#bool)** = True* +strided_mm1*: [bool](https://docs.python.org/3/library/functions.html#bool)\*\* = True* Use strided memory access for first matmul to improve cache locality -use_pos_id*: [bool](https://docs.python.org/3/library/functions.html#bool)** = False* +use_pos_id*: [bool](https://docs.python.org/3/library/functions.html#bool)\*\* = False* Generate attention mask from position IDs in-kernel instead of loading pre-generated mask -fuse_rope*: [bool](https://docs.python.org/3/library/functions.html#bool)** = False* +fuse_rope*: [bool](https://docs.python.org/3/library/functions.html#bool)\*\* = False* Fuse RoPE (Rotary Position Embedding) computation into the kernel -use_gpsimd_sb2sb*: [bool](https://docs.python.org/3/library/functions.html#bool)** = True* +use_gpsimd_sb2sb*: [bool](https://docs.python.org/3/library/functions.html#bool)\*\* = True* Use GPSIMD instructions for SBUF-to-SBUF data transfers (LNC2 sharding) -qk_in_sb*: [bool](https://docs.python.org/3/library/functions.html#bool)** = False* +qk_in_sb*: [bool](https://docs.python.org/3/library/functions.html#bool)\*\* = False* Query and key tensors are already in SBUF instead of HBM -k_out_in_sb*: [bool](https://docs.python.org/3/library/functions.html#bool)** = False* +k_out_in_sb*: [bool](https://docs.python.org/3/library/functions.html#bool)\*\* = False* Output key tensor after RoPE should be stored in SBUF instead of HBM -out_in_sb*: [bool](https://docs.python.org/3/library/functions.html#bool)** = False* +out_in_sb*: [bool](https://docs.python.org/3/library/functions.html#bool)\*\* = False* Output tensor should be stored in SBUF instead of HBM ### attention_tkg -nkilib.core.attention_tkg.attention_tkg(*q*, *k_active*, *v_active*, *k_prior*, *v_prior*, *mask*, *out*, *cfg*, *sbm*, *inv_freqs=None*, *rope_pos_ids=None*, *sink=None*, *active_blocks_table=None*, *k_out=None*, *DBG_TENSORS=None*) +nkilib.core.attention*tkg.attention_tkg(\_q*, _k_active_, _v_active_, _k_prior_, _v_prior_, _mask_, _out_, _cfg_, _sbm_, _inv_freqs=None_, _rope_pos_ids=None_, _sink=None_, _active_blocks_table=None_, _k_out=None_, _DBG_TENSORS=None_) Attention specifically optimized for token-gen (where s_active is small). Can optionally fuse RoPE at the start. Parameters: -* **q** (`nl.ndarray`) – Query tensor. Shape depends on `cfg.qk_in_sb`: If `True`: `[d, B * H * s_active]`, else: `[B, d, H, s_active]` +- **q** (`nl.ndarray`) – Query tensor. Shape depends on `cfg.qk_in_sb`: If `True`: `[d, B * H * s_active]`, else: `[B, d, H, s_active]` -* **k_active** (`nl.ndarray`) – Active key tensor. Shape depends on `cfg.qk_in_sb`: If `True`: `[d, B * s_active]`, else: `[B, d, s_active]` +- **k_active** (`nl.ndarray`) – Active key tensor. Shape depends on `cfg.qk_in_sb`: If `True`: `[d, B * s_active]`, else: `[B, d, s_active]` -* **v_active** (`nl.ndarray`) – Active value tensor. Shape: `[B, 1, s_active, d]` +- **v_active** (`nl.ndarray`) – Active value tensor. Shape: `[B, 1, s_active, d]` -* **k_prior** (`nl.ndarray`) – Prior key tensor from KV cache. Shape: `[B+, 1, s_prior, d]` if `cfg.tp_k_prior` else `[B+, 1, d, s_prior]`. For block KV cache, shape is `[B+ * block_count, block_len, d]` +- **k_prior** (`nl.ndarray`) – Prior key tensor from KV cache. Shape: `[B+, 1, s_prior, d]` if `cfg.tp_k_prior` else `[B+, 1, d, s_prior]`. For block KV cache, shape is `[B+ * block_count, block_len, d]` -* **v_prior** (`nl.ndarray`) – Prior value tensor from KV cache. Shape: `[B+, 1, s_prior, d]`. For block KV cache, shape is `[B+ * block_count, block_len, d]` +- **v_prior** (`nl.ndarray`) – Prior value tensor from KV cache. Shape: `[B+, 1, s_prior, d]`. For block KV cache, shape is `[B+ * block_count, block_len, d]` -* **mask** (`nl.ndarray`) – Attention mask. Shape: `[s_active, B, H, s_active]` if `cfg.use_pos_id` else `[s_prior, B, H, s_active]` +- **mask** (`nl.ndarray`) – Attention mask. Shape: `[s_active, B, H, s_active]` if `cfg.use_pos_id` else `[s_prior, B, H, s_active]` -* **out** (`nl.ndarray`) – Output tensor. Shape depends on `cfg.out_in_sb`: If `True`: `[d, B * H * s_active]`, else: `[B, H, d, s_active]` +- **out** (`nl.ndarray`) – Output tensor. Shape depends on `cfg.out_in_sb`: If `True`: `[d, B * H * s_active]`, else: `[B, H, d, s_active]` -* **cfg** (`AttnTKGConfig`) – Kernel configuration with shapes and performance flags +- **cfg** (`AttnTKGConfig`) – Kernel configuration with shapes and performance flags -* **sbm** (`SbufManager`) – SBUF memory manager for allocating temporary buffers +- **sbm** (`SbufManager`) – SBUF memory manager for allocating temporary buffers -* **inv_freqs** (`nl.ndarray`, optional) – Inverse frequencies for RoPE. Shape: `[d // 2, 1]`. Required when `cfg.fuse_rope` is `True` +- **inv_freqs** (`nl.ndarray`, optional) – Inverse frequencies for RoPE. Shape: `[d // 2, 1]`. Required when `cfg.fuse_rope` is `True` -* **rope_pos_ids** (`nl.ndarray`, optional) – Position IDs for RoPE. Shape: `[B, s_active]`. Required when `cfg.fuse_rope` or `cfg.use_pos_id` is `True` +- **rope_pos_ids** (`nl.ndarray`, optional) – Position IDs for RoPE. Shape: `[B, s_active]`. Required when `cfg.fuse_rope` or `cfg.use_pos_id` is `True` -* **sink** (`nl.ndarray`, optional) – Sink attention tokens. Shape: `[H, 1]` for streaming attention sink tokens +- **sink** (`nl.ndarray`, optional) – Sink attention tokens. Shape: `[H, 1]` for streaming attention sink tokens -* **active_blocks_table** (`nl.ndarray`, optional) – Table of active blocks for block KV cache. Shape: `[B, num_blocks]`. Required when using block KV cache +- **active_blocks_table** (`nl.ndarray`, optional) – Table of active blocks for block KV cache. Shape: `[B, num_blocks]`. Required when using block KV cache -* **k_out** (`nl.ndarray`, optional) – Output key tensor after RoPE. Shape depends on `cfg.k_out_in_sb`: If `True`: `[d, B * s_active]`, else: `[B, 1, d, s_active]` +- **k_out** (`nl.ndarray`, optional) – Output key tensor after RoPE. Shape depends on `cfg.k_out_in_sb`: If `True`: `[d, B * s_active]`, else: `[B, 1, d, s_active]` -* **DBG_TENSORS** (`tuple`, optional) – Optional tuple of 4-5 debug tensors with shared HBM type for intermediate value inspection +- **DBG_TENSORS** (`tuple`, optional) – Optional tuple of 4-5 debug tensors with shared HBM type for intermediate value inspection Returns: Tuple of `(out, k_out)` where `out` is the attention output tensor and `k_out` is the key output tensor (if `cfg.fuse_rope` is `True`) @@ -129,170 +129,170 @@ Return type: **Constraints**: -* Optimized for `s_active <= 7` and `d_head <= 128` +- Optimized for `s_active <= 7` and `d_head <= 128` -* `cfg.qk_in_sb=True` is required when skipping fused RoPE +- `cfg.qk_in_sb=True` is required when skipping fused RoPE -* Block KV cache requires `cfg.qk_in_sb=True` +- Block KV cache requires `cfg.qk_in_sb=True` -* In-kernel mask generation (`cfg.use_pos_id=True`) is not supported with batch sharding or block KV cache +- In-kernel mask generation (`cfg.use_pos_id=True`) is not supported with batch sharding or block KV cache ## Features -* **Flexible Tensor Placement**: +- **Flexible Tensor Placement**: `q`, `k`, `k_out`, and `out` tensors can be placed in either SBUF or HBM -* When `qk_in_sb=True`, q and k tensors are pre-loaded in SBUF (required for block KV cache) +- When `qk_in_sb=True`, q and k tensors are pre-loaded in SBUF (required for block KV cache) -* `out_in_sb` and `k_out_in_sb` flags control output tensor placement for reduced memory transfers +- `out_in_sb` and `k_out_in_sb` flags control output tensor placement for reduced memory transfers -* Use this feature for performance improvement when integrating this kernel into a larger kernel +- Use this feature for performance improvement when integrating this kernel into a larger kernel -* **Adaptive LNC2 Sharding**: +- **Adaptive LNC2 Sharding**: Automatically selects sharding strategy based on tensor dimensions -* Batch sharding: Used when batch is even AND (`s_prior < 256` OR `b*q_head*s_active > 128`) +- Batch sharding: Used when batch is even AND (`s_prior < 256` OR `b*q_head*s_active > 128`) -* Sequence sharding: Used when `s_prior >= 256` and batch sharding criteria not met +- Sequence sharding: Used when `s_prior >= 256` and batch sharding criteria not met -* Balances computation across 2 NeuronCores for improved throughput +- Balances computation across 2 NeuronCores for improved throughput -* **Mask Generation**: +- **Mask Generation**: `use_pos_id=False`: Pre-generated mask loaded from HBM -* `use_pos_id=True`: Mask generated in-kernel from position IDs +- `use_pos_id=True`: Mask generated in-kernel from position IDs -* In-kernel generation reduces memory bandwidth but requires position ID input +- In-kernel generation reduces memory bandwidth but requires position ID input -* **Fused RoPE (Rotary Position Embedding)**: +- **Fused RoPE (Rotary Position Embedding)**: `fuse_rope` integrates RoPE computation directly into the attention kernel -* Applies rotary embeddings to Q and K tensors, scaling Q by `1/sqrt(d_head)` +- Applies rotary embeddings to Q and K tensors, scaling Q by `1/sqrt(d_head)` -* Reduces memory traffic by avoiding separate RoPE passes +- Reduces memory traffic by avoiding separate RoPE passes -* **Block KV Cache**: +- **Block KV Cache**: Supports block-sparse KV cache with configurable `block_len` -* Uses `active_blocks_table` to track which cache blocks are active per batch +- Uses `active_blocks_table` to track which cache blocks are active per batch -* Enables efficient long-context inference with sparse memory access patterns +- Enables efficient long-context inference with sparse memory access patterns -* **K_prior Transpose Handling**: +- **K_prior Transpose Handling**: `tp_k_prior` flag indicates whether K_prior is pre-transposed in memory -* Optimizes memory layout: `[B, 1, d, s_prior]` when `tp_k_prior=True` vs `[B, 1, s_prior, d]` when False +- Optimizes memory layout: `[B, 1, d, s_prior]` when `tp_k_prior=True` vs `[B, 1, s_prior, d]` when False -* Reduces transpose operations during computation and improves interoperability with other kernels +- Reduces transpose operations during computation and improves interoperability with other kernels -* **Strided Memory Access (strided_mm1)**: +- **Strided Memory Access (strided_mm1)**: Enables strided read patterns for K in first matmul -* When enabled, allows MM2 to use sequential V reads for better DMA throughput +- When enabled, allows MM2 to use sequential V reads for better DMA throughput -* Trades off MM1 memory access for MM2 optimization +- Trades off MM1 memory access for MM2 optimization -* **Attention Sink**: +- **Attention Sink**: -* Supports streaming attention with sink tokens for infinite context +- Supports streaming attention with sink tokens for infinite context -* Sink tokens maintain fixed attention scores across all positions +- Sink tokens maintain fixed attention scores across all positions -* Integrated into softmax reduction for minimal overhead +- Integrated into softmax reduction for minimal overhead -* **GPSIMD SBUF-to-SBUF Transfers**: +- **GPSIMD SBUF-to-SBUF Transfers**: -* `use_gpsimd_sb2sb` enables high-performance GPSIMD instructions for inter-core communication +- `use_gpsimd_sb2sb` enables high-performance GPSIMD instructions for inter-core communication -* Optimizes LNC2 sharding by using extended instructions for SBUF-to-SBUF data transfers +- Optimizes LNC2 sharding by using extended instructions for SBUF-to-SBUF data transfers -* **Context Length Management**: +- **Context Length Management**: `curr_sprior`: Current prior sequence length (actual KV cache content for this invocation) -* `full_sprior`: Full prior sequence length (maximum KV cache capacity allocated) +- `full_sprior`: Full prior sequence length (maximum KV cache capacity allocated) -* Allows progressive filling of KV cache during autoregressive generation +- Allows progressive filling of KV cache during autoregressive generation ## Implementation Details The kernel implementation includes several key optimizations: -* **Efficient Tiling Strategy**: Uses carefully chosen tile sizes for processing batches, sequences, and heads to maximize hardware utilization. +- **Efficient Tiling Strategy**: Uses carefully chosen tile sizes for processing batches, sequences, and heads to maximize hardware utilization. -* **Cascaded Reduction**: Implements cascaded max and sum reduction operations for softmax computation to maintain numerical stability. +- **Cascaded Reduction**: Implements cascaded max and sum reduction operations for softmax computation to maintain numerical stability. -* **Memory Access Optimization**: Employs careful memory access patterns to optimize data movement between HBM and SBUF. +- **Memory Access Optimization**: Employs careful memory access patterns to optimize data movement between HBM and SBUF. -* **Block KV Cache Support**: Implements efficient block-sparse KV cache with dynamic block size adjustment to ensure optimal hardware utilization. +- **Block KV Cache Support**: Implements efficient block-sparse KV cache with dynamic block size adjustment to ensure optimal hardware utilization. -* **Attention Sink Integration**: Efficiently integrates attention sink tokens into the softmax computation for streaming attention. +- **Attention Sink Integration**: Efficiently integrates attention sink tokens into the softmax computation for streaming attention. -* **Fused RoPE Implementation**: Implements efficient rotary position embeddings with optimized trigonometric computations. +- **Fused RoPE Implementation**: Implements efficient rotary position embeddings with optimized trigonometric computations. -* **Adaptive Sharding**: Dynamically selects between batch and sequence sharding based on tensor dimensions to optimize performance. +- **Adaptive Sharding**: Dynamically selects between batch and sequence sharding based on tensor dimensions to optimize performance. -* **GPSIMD Optimization**: Uses GPSIMD instructions for high-performance SBUF-to-SBUF data transfers in LNC2 sharding. +- **GPSIMD Optimization**: Uses GPSIMD instructions for high-performance SBUF-to-SBUF data transfers in LNC2 sharding. -* **Debug Support**: Provides comprehensive debug tensor support for intermediate value inspection. +- **Debug Support**: Provides comprehensive debug tensor support for intermediate value inspection. -* **Stack-based SBUF Allocation**: Uses SbufManager for efficient on-chip memory management with hierarchical scoping. +- **Stack-based SBUF Allocation**: Uses SbufManager for efficient on-chip memory management with hierarchical scoping. ## Algorithm The kernel goes through the following steps: -* **Setup**: Initialize intermediate buffers, mask, block KV, and debug tensors. +- **Setup**: Initialize intermediate buffers, mask, block KV, and debug tensors. -* **Optional RoPE**: If `fuse_rope` is enabled, apply rotary position embeddings to Q and K tensors. +- **Optional RoPE**: If `fuse_rope` is enabled, apply rotary position embeddings to Q and K tensors. -* **KQ^T Computation**: Perform the first matrix multiplication to compute attention scores. +- **KQ^T Computation**: Perform the first matrix multiplication to compute attention scores. Loop over each batch -* Load the current chunk of K based on configuration (block KV, transpose, etc.) +- Load the current chunk of K based on configuration (block KV, transpose, etc.) -* Tile over the multiplication of K and Q in groups of 4k size +- Tile over the multiplication of K and Q in groups of 4k size -* **Max Reduction**: Compute the max reduction of KQ^T for softmax stability. +- **Max Reduction**: Compute the max reduction of KQ^T for softmax stability. Compute the max in tiles of size 128 over `bs * q_head * s_active` -* Prepare the sink if used +- Prepare the sink if used -* Transpose and broadcast along the partition dimension +- Transpose and broadcast along the partition dimension -* **Exp(KQ^T - max(KQ^T))**: Apply the exponentiation for softmax computation. +- **Exp(KQ^T - max(KQ^T))**: Apply the exponentiation for softmax computation. Add/subtract the max based on whether it was negated -* Apply the exponentiation activation +- Apply the exponentiation activation -* **Sum Reduction**: Compute sum reduction of the exponentiation result. +- **Sum Reduction**: Compute sum reduction of the exponentiation result. Compute the sum in tiles of size 128 over `bs * q_head * s_active` -* Perform additional reductions based on sink or other optimization flags +- Perform additional reductions based on sink or other optimization flags -* Compute the reciprocal with the same tiling scheme, and then broadcast +- Compute the reciprocal with the same tiling scheme, and then broadcast -* **Final Matrix Multiplication**: Compute the product of the softmax output and V and store the result +- **Final Matrix Multiplication**: Compute the product of the softmax output and V and store the result Loop over each batch -* Load the current chunk of V based on configuration +- Load the current chunk of V based on configuration -* Perform the matmul over sprior tiles +- Perform the matmul over sprior tiles -* If needed, copy information over core boundaries or to HBM +- If needed, copy information over core boundaries or to HBM ## See Also -* [Output Projection TKG Kernel API Reference](output-projection-tkg.md) \ No newline at end of file +- [Output Projection TKG Kernel API Reference](output-projection-tkg.md) diff --git a/skills/neuron-nki-docs/references/reference/library/design-rmsnorm-quant.md b/skills/neuron-nki-docs/references/reference/library/design-rmsnorm-quant.md index d3c4b95..4000476 100644 --- a/skills/neuron-nki-docs/references/reference/library/design-rmsnorm-quant.md +++ b/skills/neuron-nki-docs/references/reference/library/design-rmsnorm-quant.md @@ -7,7 +7,7 @@ For details on how to use this kernel, see the [RMSNorm-Quant Kernel API Referen ## Background -This kernel performs *optional* [RMS normalization](https://arxiv.org/abs/1910.07467) followed by quantization to `fp8`. +This kernel performs _optional_ [RMS normalization](https://arxiv.org/abs/1910.07467) followed by quantization to `fp8`. ### Motivation @@ -36,7 +36,7 @@ The equation is: \[\mathrm{RMSNorm}(x_i)=\frac{x_i}{\mathrm{RMS}(x)} \gamma_i \quad \text{for } i = 1 \dots H\] where: -\[\begin{split}\mathrm{RMS}(x)=\sqrt{(\frac{1}{H} \sum_{i=1}^{H} x_i^2) + \epsilon} \\ +\[\begin{split}\mathrm{RMS}(x)=\sqrt{(\frac{1}{H} \sum\_{i=1}^{H} x_i^2) + \epsilon} \\ x = \text{each [B,S] with shape [H]} \\ \gamma \text{ = gamma with shape [H]} \\ \epsilon = \text{ small positive value for numerical stability}\end{split}\] @@ -60,17 +60,17 @@ Quantization is independently performed on each [B,S]. The equation is: -\[\begin{split}M = \max_{i=1}^{H} |x_i| \\ +\[\begin{split}M = \max\_{i=1}^{H} |x_i| \\ D = \frac{M}{240} \\ Q = \frac{1}{D} \\ -\mathbf{x}_q = xQ\end{split}\] +\mathbf{x}\_q = xQ\end{split}\] or equivalently -\[x_{q,i} = x_iQ \quad \text{for } i = 1, \dots, H\] +\[x\_{q,i} = x_iQ \quad \text{for } i = 1, \dots, H\] where \[\begin{split}x = \text{each [B,S] with shape [H]} \\ -\mathbf{x}_q = \text{quantized } \mathbf{x} \\ +\mathbf{x}\_q = \text{quantized } \mathbf{x} \\ D = \text{de-quantization scale} \\ Q = \text{quantization scale}\end{split}\] The above equation omits clipping/flooring details which are instead included later in this document. @@ -130,7 +130,6 @@ The commented code and the above sections should together deliver a good underst The following is a simple Python equivalent to the kernel which can be another useful way of understanding the kernel’s behaviour. - ```python def rmsnorm_quant_ref(inp: np.ndarray, gamma: np.ndarray, eps: float = 1e-6) -> Tuple[np.ndarray, np.ndarray]: """RMSNorm + Quantization reference impl. @@ -161,25 +160,20 @@ def rmsnorm_quant_ref(inp: np.ndarray, gamma: np.ndarray, eps: float = 1e-6) -> return norm_quant, dequant_scale ``` - ### Kernel Code Details rms_normalize_tile() contains a loop to tile across the processing dimension. This loop contains the following directive: - ```python directives=ncc.multi_buffer(constants.num_hw_psum_banks) ``` - This enables the compiler to replicate the gamma PSUM allocation (into which the gamma-broadcast matmul result is stored), improving pipeline parallelism by enabling each loop iteration to write into a separate PSUM bank. - ```python skip_middle_end_transformations ``` - The compiler middle-end-transformation passes contain heuristic-driven optimizations, including loop-reordering and loop-fusion. While these passes could help improve performance, in some cases, they are not predictable. Kernels are generally hand-tuned to achieve optimal performance, so we turn them off. ## Kernel API @@ -192,20 +186,17 @@ The section includes some example performance targets for real world model confi **Llama3.3 70B** - | Target Latency (us) | Batch Count | Sequence Length | Hidden | -| --- | --- | --- | --- | -| 458.2 | 1 | 2K | 8192 | -| 6,287.0 | 1 | 32K | 8192 | +| ------------------- | ----------- | --------------- | ------ | +| 458.2 | 1 | 2K | 8192 | +| 6,287.0 | 1 | 32K | 8192 | **Llama3.1 405B** - | Target Latency (us) | Batch Count | Sequence Length | Hidden | -| --- | --- | --- | --- | -| 866.81 | 1 | 2K | 16384 | -| 13,214.40 | 1 | 32K | 16384 | - +| ------------------- | ----------- | --------------- | ------ | +| 866.81 | 1 | 2K | 16384 | +| 13,214.40 | 1 | 32K | 16384 | ## Performance Analysis @@ -213,19 +204,19 @@ Here we demonstrate a sample execution of this kernel and break it down in the P **Test Parameters:** -* LNC: 2 ( Note, two pairs of instructions in nc0, and nc1 in captured figures ) +- LNC: 2 ( Note, two pairs of instructions in nc0, and nc1 in captured figures ) -* Batch Size: 1 +- Batch Size: 1 -* Sequence Length: 160 +- Sequence Length: 160 -* Hidden Size: 16,384 +- Hidden Size: 16,384 -* Data Type: dt.bfloat16 +- Data Type: dt.bfloat16 -* Quantization Data Type: dt.float8_e4m3 +- Quantization Data Type: dt.float8_e4m3 -* Quantization Only: False +- Quantization Only: False The following picture shows the overall execution. @@ -235,33 +226,33 @@ The following picture shows the overall execution. This phase involves two DMA load operations: one for the hidden tensor and one for the gamma tensor. -* **Hidden Tensor**: The DMA buffer size is calculated as hidden_size * sizeof(dtype). +- **Hidden Tensor**: The DMA buffer size is calculated as hidden_size \* sizeof(dtype). -* **Gamma Tensor**: The code intends to load the entire [1, H] tensor in a single operation. However, it should be noted that the compiler performs optimizations for trivial dimensions, which can result in several small (e.g., 4-byte) DMA buffer loads. +- **Gamma Tensor**: The code intends to load the entire [1, H] tensor in a single operation. However, it should be noted that the compiler performs optimizations for trivial dimensions, which can result in several small (e.g., 4-byte) DMA buffer loads. ### Phase 2: RMSNorm ![../../../_images/profile_phase_2.png](../../../_images/profile_phase_2.png) -* Compute Inverse RMS scale +- Compute Inverse RMS scale This step involves two ACT (activation) instructions: activation_reduce: Squares each element of the hidden tensor and performs a reduction (sum) across the hidden dimension. -* activation: Adds a small constant eps for numerical stability, applies a scaling factor (1 / H), and then computes the reciprocal square root of the result. +- activation: Adds a small constant eps for numerical stability, applies a scaling factor (1 / H), and then computes the reciprocal square root of the result. -* Broadcast Gamma – Part 1 / Part 2 +- Broadcast Gamma – Part 1 / Part 2 As previously mentioned, a multi-buffer strategy is used for PSUM. Assuming there are N PSUM banks, Part 1 of the broadcast operation replicates the gamma values of shape [1, 512] to [128, 512] tiles, repeating this process N times. -* The size 512 corresponds to the **free dimension limit** of the TensorEngine, meaning we must slice the H dimension (processing dimension) into chunks of 512. +- The size 512 corresponds to the **free dimension limit** of the TensorEngine, meaning we must slice the H dimension (processing dimension) into chunks of 512. -* The broadcast is divided into Part 1 and Part 2 because the inverse RMS scale value is needed before evicting data from the PSUM buffers after Part 1. The PSUM data is not evicted to the SBUF immediately; instead, it remains in place to be consumed by the scalar_tensor_tensor operation once inverse_rms_scale is ready. This behavior is intentional, as there is limited performance benefit in evicting PSUMs early. Part 2 of the gamma broadcast is fully pipelined with the subsequent scalar_tensor_tensor instruction, making early eviction unnecessary. +- The broadcast is divided into Part 1 and Part 2 because the inverse RMS scale value is needed before evicting data from the PSUM buffers after Part 1. The PSUM data is not evicted to the SBUF immediately; instead, it remains in place to be consumed by the scalar_tensor_tensor operation once inverse_rms_scale is ready. This behavior is intentional, as there is limited performance benefit in evicting PSUMs early. Part 2 of the gamma broadcast is fully pipelined with the subsequent scalar_tensor_tensor instruction, making early eviction unnecessary. -* Apply gamma and inverse RMS scale +- Apply gamma and inverse RMS scale -This step is performed using the scalar_tensor_tensor instruction, with a free dimension size of 512, matching the limit of the TensorEngine. This allows the operation to be *efficiently pipelined* with the TensorEngine activity. +This step is performed using the scalar*tensor_tensor instruction, with a free dimension size of 512, matching the limit of the TensorEngine. This allows the operation to be \_efficiently pipelined* with the TensorEngine activity. ### Phase 3: Quantization @@ -269,26 +260,26 @@ This step is performed using the scalar_tensor_tensor instruction, with a free d The overall quantization process involves heavy use of the VectorEngine, primarily due to the max function. These instructions are executed **sequentially with no parallelism**, as each step depends on the result of the previous one. -* Compute absolute maximum +- Compute absolute maximum -* Compute dequantization scale +- Compute dequantization scale -activation: The dequantization scale is derived by dividing the absolute max by _FP8_RANGE +activation: The dequantization scale is derived by dividing the absolute max by \_FP8_RANGE -* Compute quantized output +- Compute quantized output -tensor_scalar: clamp to _MIN_DEQUANT_SCALE_VAL for numerical stability +tensor_scalar: clamp to \_MIN_DEQUANT_SCALE_VAL for numerical stability -* reciprocal: compute the reciprocal to get the quantization scale +- reciprocal: compute the reciprocal to get the quantization scale -* tensor_scalar: Apply quantization scale to produce the quantized result +- tensor_scalar: Apply quantization scale to produce the quantized result ### Phase 4: Store output Store quantized value with dequantizing scale -* **Hidden Tensor**: -The DMA buffer size is calculated as hidden_size * sizeof(quant_dtype). +- **Hidden Tensor**: + The DMA buffer size is calculated as hidden_size \* sizeof(quant_dtype). -* **Dequantization Scale:** -The DMA buffer size is calculated as 4* sizeof(quant_dtype). \ No newline at end of file +- **Dequantization Scale:** + The DMA buffer size is calculated as 4\* sizeof(quant_dtype). diff --git a/skills/neuron-nki-docs/references/reference/library/mlp.md b/skills/neuron-nki-docs/references/reference/library/mlp.md index 8f6afff..8105094 100644 --- a/skills/neuron-nki-docs/references/reference/library/mlp.md +++ b/skills/neuron-nki-docs/references/reference/library/mlp.md @@ -5,25 +5,25 @@ This topic provides the API reference for the `MLP` kernel. The kernel implement The kernel supports: -* Both context encoding (CTE) and token generation (TKG) modes +- Both context encoding (CTE) and token generation (TKG) modes -* Optional normalization fusion (RMSNorm, LayerNorm) +- Optional normalization fusion (RMSNorm, LayerNorm) -* Various activation functions +- Various activation functions -* Residual connections via fused addition +- Residual connections via fused addition -* Flexible tensor layouts and column tiling optimizations +- Flexible tensor layouts and column tiling optimizations -* Bias addition for all projections and normalization +- Bias addition for all projections and normalization -* FP8 quantization (static and row-wise, TKG mode only) +- FP8 quantization (static and row-wise, TKG mode only) -* Gate and up projection result clamping +- Gate and up projection result clamping -* Optional gate projection skipping +- Optional gate projection skipping -* SBUF output for kernel fusion +- SBUF output for kernel fusion ## Background @@ -32,8 +32,7 @@ The `MLP` kernel is a critical component in transformer architectures, responsib > **Note** > > Note -> -> +> > This kernel automatically selects between TKG (Token Generation) and CTE (Context Encoding) implementations based on the batch size × sequence length threshold (currently 96, planned to increase to 128), ensuring optimal performance across different use cases. ## API Reference @@ -42,7 +41,7 @@ The `MLP` kernel is a critical component in transformer architectures, responsib ### mlp -nkilib.core.mlp.mlp(*hidden_tensor: nl.ndarray*, *gate_proj_weights_tensor: nl.ndarray*, *up_proj_weights_tensor: nl.ndarray*, *down_proj_weights_tensor: nl.ndarray*, *normalization_weights_tensor: Optional[nl.ndarray] = None*, *gate_proj_bias_tensor: Optional[nl.ndarray] = None*, *up_proj_bias_tensor: Optional[nl.ndarray] = None*, *down_proj_bias_tensor: Optional[nl.ndarray] = None*, *normalization_bias_tensor: Optional[nl.ndarray] = None*, *fused_add_tensor: Optional[nl.ndarray] = None*, *store_fused_add_result: [bool](https://docs.python.org/3/library/functions.html#bool) = False*, *activation_fn: ActFnType = ActFnType.SiLU*, *normalization_type: NormType = NormType.NO_NORM*, *quantization_type: QuantizationType = QuantizationType.NONE*, *gate_w_scale: Optional[nl.ndarray] = None*, *up_w_scale: Optional[nl.ndarray] = None*, *down_w_scale: Optional[nl.ndarray] = None*, *gate_up_in_scale: Optional[nl.ndarray] = None*, *down_in_scale: Optional[nl.ndarray] = None*, *output_dtype=None*, *store_output_in_sbuf: [bool](https://docs.python.org/3/library/functions.html#bool) = False*, *eps: [float](https://docs.python.org/3/library/functions.html#float) = 1e-6*, *skip_gate_proj: [bool](https://docs.python.org/3/library/functions.html#bool) = False*, *use_tkg_gate_up_proj_column_tiling: [bool](https://docs.python.org/3/library/functions.html#bool) = True*, *use_tkg_down_proj_column_tiling: [bool](https://docs.python.org/3/library/functions.html#bool) = True*, *use_tkg_down_proj_optimized_layout: [bool](https://docs.python.org/3/library/functions.html#bool) = False*, *gate_clamp_upper_limit: Optional[[float](https://docs.python.org/3/library/functions.html#float)] = None*, *gate_clamp_lower_limit: Optional[[float](https://docs.python.org/3/library/functions.html#float)] = None*, *up_clamp_upper_limit: Optional[[float](https://docs.python.org/3/library/functions.html#float)] = None*, *up_clamp_lower_limit: Optional[[float](https://docs.python.org/3/library/functions.html#float)] = None*, *force_cte_mode: [bool](https://docs.python.org/3/library/functions.html#bool) = False*) +nkilib.core.mlp.mlp(_hidden_tensor: nl.ndarray_, _gate_proj_weights_tensor: nl.ndarray_, _up_proj_weights_tensor: nl.ndarray_, _down_proj_weights_tensor: nl.ndarray_, _normalization_weights_tensor: Optional[nl.ndarray] = None_, _gate_proj_bias_tensor: Optional[nl.ndarray] = None_, _up_proj_bias_tensor: Optional[nl.ndarray] = None_, _down_proj_bias_tensor: Optional[nl.ndarray] = None_, _normalization_bias_tensor: Optional[nl.ndarray] = None_, _fused_add_tensor: Optional[nl.ndarray] = None_, _store_fused_add_result: [bool](https://docs.python.org/3/library/functions.html#bool) = False_, _activation_fn: ActFnType = ActFnType.SiLU_, _normalization_type: NormType = NormType.NO_NORM_, _quantization_type: QuantizationType = QuantizationType.NONE_, _gate_w_scale: Optional[nl.ndarray] = None_, _up_w_scale: Optional[nl.ndarray] = None_, _down_w_scale: Optional[nl.ndarray] = None_, _gate_up_in_scale: Optional[nl.ndarray] = None_, _down_in_scale: Optional[nl.ndarray] = None_, _output_dtype=None_, _store_output_in_sbuf: [bool](https://docs.python.org/3/library/functions.html#bool) = False_, _eps: [float](https://docs.python.org/3/library/functions.html#float) = 1e-6_, _skip_gate_proj: [bool](https://docs.python.org/3/library/functions.html#bool) = False_, _use_tkg_gate_up_proj_column_tiling: [bool](https://docs.python.org/3/library/functions.html#bool) = True_, _use_tkg_down_proj_column_tiling: [bool](https://docs.python.org/3/library/functions.html#bool) = True_, _use_tkg_down_proj_optimized_layout: [bool](https://docs.python.org/3/library/functions.html#bool) = False_, _gate_clamp_upper_limit: Optional[[float](https://docs.python.org/3/library/functions.html#float)] = None_, _gate_clamp_lower_limit: Optional[[float](https://docs.python.org/3/library/functions.html#float)] = None_, _up_clamp_upper_limit: Optional[[float](https://docs.python.org/3/library/functions.html#float)] = None_, _up_clamp_lower_limit: Optional[[float](https://docs.python.org/3/library/functions.html#float)] = None_, _force_cte_mode: [bool](https://docs.python.org/3/library/functions.html#bool) = False_) MLP(Multi-Layer Perceptron) Kernel implementation. Performs the standard MLP computation with support for both context encoding (CTE) and @@ -51,116 +50,116 @@ on input dimensions and supports various optimizations. Parameters: -* **hidden_tensor** (`nl.ndarray`) – Input hidden states tensor with shape [B, S, H] or SBUF layout. +- **hidden_tensor** (`nl.ndarray`) – Input hidden states tensor with shape [B, S, H] or SBUF layout. -* **gate_proj_weights_tensor** (`nl.ndarray`) – Gate projection weight matrix with shape [H, I]. +- **gate_proj_weights_tensor** (`nl.ndarray`) – Gate projection weight matrix with shape [H, I]. -* **up_proj_weights_tensor** (`nl.ndarray`) – Up projection weight matrix with shape [H, I]. +- **up_proj_weights_tensor** (`nl.ndarray`) – Up projection weight matrix with shape [H, I]. -* **down_proj_weights_tensor** (`nl.ndarray`) – Down projection weight matrix with shape [I, H]. +- **down_proj_weights_tensor** (`nl.ndarray`) – Down projection weight matrix with shape [I, H]. -* **normalization_weights_tensor** (`nl.ndarray`, optional) – Normalization weights with shape [1, H]. +- **normalization_weights_tensor** (`nl.ndarray`, optional) – Normalization weights with shape [1, H]. -* **gate_proj_bias_tensor** (`nl.ndarray`, optional) – Bias tensor for gate projection with shape [1, I]. +- **gate_proj_bias_tensor** (`nl.ndarray`, optional) – Bias tensor for gate projection with shape [1, I]. -* **up_proj_bias_tensor** (`nl.ndarray`, optional) – Bias tensor for up projection with shape [1, I]. +- **up_proj_bias_tensor** (`nl.ndarray`, optional) – Bias tensor for up projection with shape [1, I]. -* **down_proj_bias_tensor** (`nl.ndarray`, optional) – Bias tensor for down projection with shape [1, H]. +- **down_proj_bias_tensor** (`nl.ndarray`, optional) – Bias tensor for down projection with shape [1, H]. -* **normalization_bias_tensor** (`nl.ndarray`, optional) – Bias tensor for normalization with shape [1, H]. Only applicable for layer normalization. +- **normalization_bias_tensor** (`nl.ndarray`, optional) – Bias tensor for normalization with shape [1, H]. Only applicable for layer normalization. -* **fused_add_tensor** (`nl.ndarray`, optional) – Tensor to fuse for the residual connection. +- **fused_add_tensor** (`nl.ndarray`, optional) – Tensor to fuse for the residual connection. -* **store_fused_add_result** (`bool`) – If True, stores the fused_add output to HBM, and the kernel returns both the fused_add output and the MLP output. Default: False. +- **store_fused_add_result** (`bool`) – If True, stores the fused_add output to HBM, and the kernel returns both the fused_add output and the MLP output. Default: False. -* **activation_fn** (`ActFnType`) – Activation function type. +- **activation_fn** (`ActFnType`) – Activation function type. -* **normalization_type** (`NormType`) – Type of normalization. +- **normalization_type** (`NormType`) – Type of normalization. -* **quantization_type** (`QuantizationType`) – Quantization type to use (default: QuantizationType.NONE). Supported values are QuantizationType.STATIC and QuantizationType.ROW. Quantization is only supported in TKG mode. +- **quantization_type** (`QuantizationType`) – Quantization type to use (default: QuantizationType.NONE). Supported values are QuantizationType.STATIC and QuantizationType.ROW. Quantization is only supported in TKG mode. -* **gate_w_scale** (`nl.ndarray`, optional) – FP8 dequantization scales for gate weights. Shape is [128, I] for row-wise quantization, [128, 1] for static quantization. Defaults to None. +- **gate_w_scale** (`nl.ndarray`, optional) – FP8 dequantization scales for gate weights. Shape is [128, I] for row-wise quantization, [128, 1] for static quantization. Defaults to None. -* **up_w_scale** (`nl.ndarray`, optional) – FP8 dequantization scales for up weights. Shape is [128, I] for row-wise quantization, [128, 1] for static quantization. Defaults to None. +- **up_w_scale** (`nl.ndarray`, optional) – FP8 dequantization scales for up weights. Shape is [128, I] for row-wise quantization, [128, 1] for static quantization. Defaults to None. -* **down_w_scale** (`nl.ndarray`, optional) – FP8 dequantization scales for down weights. Shape is [128, I] for row-wise quantization, [128, 1] for static quantization. Defaults to None. +- **down_w_scale** (`nl.ndarray`, optional) – FP8 dequantization scales for down weights. Shape is [128, I] for row-wise quantization, [128, 1] for static quantization. Defaults to None. -* **gate_up_in_scale** (`nl.ndarray`, optional) – FP8 dequantization scales for gate and up input. Used for static quantization with shape [128, 1]. Defaults to None. +- **gate_up_in_scale** (`nl.ndarray`, optional) – FP8 dequantization scales for gate and up input. Used for static quantization with shape [128, 1]. Defaults to None. -* **down_in_scale** (`nl.ndarray`, optional) – FP8 dequantization scales for down input. Used for static quantization with shape [128, 1]. Defaults to None. +- **down_in_scale** (`nl.ndarray`, optional) – FP8 dequantization scales for down input. Used for static quantization with shape [128, 1]. Defaults to None. -* **output_dtype** (`nki.dtype`) – Output tensor data type. Defaults to None; if None, the hidden tensor’s `dtype` is used. +- **output_dtype** (`nki.dtype`) – Output tensor data type. Defaults to None; if None, the hidden tensor’s `dtype` is used. -* **store_output_in_sbuf** (`bool`) – If True, stores the output in SBUF instead of HBM, allowing the next layer to read it directly without an additional load operation. This option is only available in TKG mode where output tensor is small enough to fit in SBUF. Default: False. +- **store_output_in_sbuf** (`bool`) – If True, stores the output in SBUF instead of HBM, allowing the next layer to read it directly without an additional load operation. This option is only available in TKG mode where output tensor is small enough to fit in SBUF. Default: False. -* **eps** (`float`) – Epsilon value for numerical stability. +- **eps** (`float`) – Epsilon value for numerical stability. -* **skip_gate_proj** (`bool`) – Skip gate projection. +- **skip_gate_proj** (`bool`) – Skip gate projection. -* **use_tkg_gate_up_proj_column_tiling** (`bool`) – If True, uses column tiling for the gate and up projection in TKG mode. Default: True. +- **use_tkg_gate_up_proj_column_tiling** (`bool`) – If True, uses column tiling for the gate and up projection in TKG mode. Default: True. -* **use_tkg_down_proj_column_tiling** (`bool`) – If True, uses column tiling for the down projection in TKG mode. Default: True. +- **use_tkg_down_proj_column_tiling** (`bool`) – If True, uses column tiling for the down projection in TKG mode. Default: True. -* **use_tkg_down_proj_optimized_layout** (`bool`) – If True, the standard down_weight tensor (`shape [I, H]`) is reinterpreted as `[I, lnc, 128, H // (128 * lnc)]`, then transposed to `[I, lnc, H // (128 * lnc), 128]`. This layout provides unit-stride weight loading, reducing the matrix multiplication initiation interval. Only applied when `use_tkg_down_proj_column_tiling` is False. Default: False. +- **use_tkg_down_proj_optimized_layout** (`bool`) – If True, the standard down_weight tensor (`shape [I, H]`) is reinterpreted as `[I, lnc, 128, H // (128 * lnc)]`, then transposed to `[I, lnc, H // (128 * lnc), 128]`. This layout provides unit-stride weight loading, reducing the matrix multiplication initiation interval. Only applied when `use_tkg_down_proj_column_tiling` is False. Default: False. -* **gate_clamp_upper_limit** (`float`, optional) – Upper bound value to clamp on gate projection results, does not perform clamping if the value is set to None. +- **gate_clamp_upper_limit** (`float`, optional) – Upper bound value to clamp on gate projection results, does not perform clamping if the value is set to None. -* **gate_clamp_lower_limit** (`float`, optional) – Lower bound value to clamp on gate projection results, does not perform clamping if the value is set to None. +- **gate_clamp_lower_limit** (`float`, optional) – Lower bound value to clamp on gate projection results, does not perform clamping if the value is set to None. -* **up_clamp_upper_limit** (`float`, optional) – Upper bound value to clamp on up projection results, does not perform clamping if the value is set to None. +- **up_clamp_upper_limit** (`float`, optional) – Upper bound value to clamp on up projection results, does not perform clamping if the value is set to None. -* **up_clamp_lower_limit** (`float`, optional) – Lower bound value to clamp on up projection results, does not perform clamping if the value is set to None. +- **up_clamp_lower_limit** (`float`, optional) – Lower bound value to clamp on up projection results, does not perform clamping if the value is set to None. -* **force_cte_mode** (`bool`) – If True, forces the use of CTE mode. Default: False. +- **force_cte_mode** (`bool`) – If True, forces the use of CTE mode. Default: False. Returns: -The MLP output tensor(s). HBM output: Tensor with shape [B, S, H]. SBUF output: Shape depends on the mode setting. CTE: Not applicable. TKG when `use_tkg_down_proj_column_tiling` is `True = [BxS, H]`. TKG when `use_tkg_down_proj_column_tiling` is `False = [128(p_max), H/128, BxS`]``. If `store_fused_add_result` is `True`, returns a list containing both the output and the stored fused output. +The MLP output tensor(s). HBM output: Tensor with shape [B, S, H]. SBUF output: Shape depends on the mode setting. CTE: Not applicable. TKG when `use_tkg_down_proj_column_tiling` is `True = [BxS, H]`. TKG when `use_tkg_down_proj_column_tiling` is `False = [128(p_max), H/128, BxS`]``. If `store_fused_add_result`is`True`, returns a list containing both the output and the stored fused output. Return type: `list[nl.ndarray]` **Notes**: -* Automatically dispatches to either CTE or TKG implementation based on batch size and sequence length. +- Automatically dispatches to either CTE or TKG implementation based on batch size and sequence length. -* Token generation mode (TKG) is used for small batch/sequence dimensions (`batch_size × sequence_length ≤ 96`), while context encoding (CTE) handles larger inputs. +- Token generation mode (TKG) is used for small batch/sequence dimensions (`batch_size × sequence_length ≤ 96`), while context encoding (CTE) handles larger inputs. -* Column tiling and tensor layout optimization (`use_tkg_down_proj_optimized_layout`) are valid only in TKG mode. +- Column tiling and tensor layout optimization (`use_tkg_down_proj_optimized_layout`) are valid only in TKG mode. -* FP8 quantization support is available only in TKG mode. +- FP8 quantization support is available only in TKG mode. -* Supported input data types: `nl.bfloat16`, `nl.float16`, `nl.float32` +- Supported input data types: `nl.bfloat16`, `nl.float16`, `nl.float32` ## Implementation Details The kernel implementation includes several key optimizations: -* **Dual Implementation Strategy**: Automatically selects between CTE (Context Encoding) and TKG (Token Generation) implementations based on `batch_size × sequence_length` threshold (currently 96, planned to increase to 128). +- **Dual Implementation Strategy**: Automatically selects between CTE (Context Encoding) and TKG (Token Generation) implementations based on `batch_size × sequence_length` threshold (currently 96, planned to increase to 128). -* **Normalization Fusion**: Optionally fuses RMSNorm or LayerNorm operations with the MLP computation for improved performance. +- **Normalization Fusion**: Optionally fuses RMSNorm or LayerNorm operations with the MLP computation for improved performance. -* **FP8 Quantization**: Supports FP8 quantization with both static and row-wise dequantization scales. Available only in TKG mode for weights and activations. +- **FP8 Quantization**: Supports FP8 quantization with both static and row-wise dequantization scales. Available only in TKG mode for weights and activations. -* **Flexible Tensor Layouts**: Supports column tiling optimizations and tensor layout optimizations in TKG mode to improve memory access patterns. +- **Flexible Tensor Layouts**: Supports column tiling optimizations and tensor layout optimizations in TKG mode to improve memory access patterns. -* **Activation Function Options**: Supports multiple activation functions, including SiLU (Swish), GELU, and ReLU. +- **Activation Function Options**: Supports multiple activation functions, including SiLU (Swish), GELU, and ReLU. -* **Result Clamping**: Provides optional clamping of gate and up projection results with configurable upper and lower bounds. +- **Result Clamping**: Provides optional clamping of gate and up projection results with configurable upper and lower bounds. -* **Gate Projection Skipping**: Allows skipping the gate projection computation when `skip_gate_proj` is enabled. +- **Gate Projection Skipping**: Allows skipping the gate projection computation when `skip_gate_proj` is enabled. -* **Residual Connection Fusion**: Can incorporate residual connections through fused_add_tensor for improved performance. +- **Residual Connection Fusion**: Can incorporate residual connections through fused_add_tensor for improved performance. -* **SBUF Output Option**: Provides the option to keep output in SBUF for fusion with subsequent operations (TKG mode only). +- **SBUF Output Option**: Provides the option to keep output in SBUF for fusion with subsequent operations (TKG mode only). -* **Bias Addition**: Supports optional bias addition for gate, up, and down projections, as well as for normalization. +- **Bias Addition**: Supports optional bias addition for gate, up, and down projections, as well as for normalization. -* **Optimized Weight Loading**: In TKG mode, `use_tkg_down_proj_optimized_layout` enables unit-stride weight loading to reduce matrix multiplication initiation interval. +- **Optimized Weight Loading**: In TKG mode, `use_tkg_down_proj_optimized_layout` enables unit-stride weight loading to reduce matrix multiplication initiation interval. -* **Multi-Precision Support**: Supports `bfloat16`, `float16`, and `float32` input data types for flexible precision requirements. +- **Multi-Precision Support**: Supports `bfloat16`, `float16`, and `float32` input data types for flexible precision requirements. ## See Also -* [QKV Kernel API Reference](qkv.md) +- [QKV Kernel API Reference](qkv.md) -* [RMSNorm-Quant Kernel API Reference](rmsnorm-quant.md) \ No newline at end of file +- [RMSNorm-Quant Kernel API Reference](rmsnorm-quant.md) diff --git a/skills/neuron-nki-docs/references/reference/library/output-projection-cte.md b/skills/neuron-nki-docs/references/reference/library/output-projection-cte.md index 9f027d2..b26d959 100644 --- a/skills/neuron-nki-docs/references/reference/library/output-projection-cte.md +++ b/skills/neuron-nki-docs/references/reference/library/output-projection-cte.md @@ -5,15 +5,15 @@ This topic provides the API reference for the `Output Projection CTE` kernel. Th The kernel supports: -* Efficient projection of attention outputs +- Efficient projection of attention outputs -* Optional bias addition +- Optional bias addition -* LNC sharding for distributed computation +- LNC sharding for distributed computation -* Optimized memory access patterns +- Optimized memory access patterns -* Head dimension packing for improved performance +- Head dimension packing for improved performance ## Background @@ -27,7 +27,7 @@ The kernel employs efficient tiling strategies and memory access patterns to max ### output_projection_cte -nkilib.core.output_projection.output_projection_cte.output_projection_cte(*attention*, *weight*, *bias=None*) +nkilib.core.output*projection.output_projection_cte.output_projection_cte(\_attention*, _weight_, _bias=None_) Output Projection Kernel optimized for Context Encoding (Prefill) use cases. This kernel computes `out = attention @ weight + bias`, typically used to project the output scores after an attention block in transformer models. @@ -38,11 +38,11 @@ This kernel uses a layout also used by other Context Encoding kernels to avoid n Parameters: -* **attention** (`nl.ndarray`) – Input tensor in HBM, typically the scores output from an attention block. Shape: `[B, N, D, S]`, where `B` is batch size, `N` is number of heads, `D` is head dimension, and `S` is sequence length. Indexing: `[b, n, d, s]`. +- **attention** (`nl.ndarray`) – Input tensor in HBM, typically the scores output from an attention block. Shape: `[B, N, D, S]`, where `B` is batch size, `N` is number of heads, `D` is head dimension, and `S` is sequence length. Indexing: `[b, n, d, s]`. -* **weight** (`nl.ndarray`) – Weight tensor in HBM. Shape: `[N*D, H]`, where `H` is hidden dimension size. Indexing: `[n * D + d, h]`. +- **weight** (`nl.ndarray`) – Weight tensor in HBM. Shape: `[N*D, H]`, where `H` is hidden dimension size. Indexing: `[n * D + d, h]`. -* **bias** (`nl.ndarray`, optional) – Optional bias tensor in HBM. Shape: `[1, H]`. Indexing: `[1, h]`. +- **bias** (`nl.ndarray`, optional) – Optional bias tensor in HBM. Shape: `[1, H]`. Indexing: `[1, h]`. Returns: Output tensor in HBM. Shape: `[B, S, H]`. Indexing: `[b, s, h]`. @@ -55,52 +55,52 @@ However, for `nl.float32`, large inputs may not fit in SBUF. **Dimensions**: -* `B`: Batch size +- `B`: Batch size -* `N`: Number of heads +- `N`: Number of heads -* `S`: Sequence length +- `S`: Sequence length -* `H`: Hidden dimension size +- `H`: Hidden dimension size -* `D`: Head dimension size +- `D`: Head dimension size **Restrictions**: -* The contract dimension of input and weight tensors must match (`N*D == weight.shape[0]`) +- The contract dimension of input and weight tensors must match (`N*D == weight.shape[0]`) -* Output projection kernel currently only supports `H` to be no more than 32768 +- Output projection kernel currently only supports `H` to be no more than 32768 -* Hidden dimension (`H`) needs to be divisible by LNC size since LNC sharding is on the weight hidden dimension +- Hidden dimension (`H`) needs to be divisible by LNC size since LNC sharding is on the weight hidden dimension -* Head dimension (`D`) must be <= 128 +- Head dimension (`D`) must be <= 128 -* Maximum validated `H` size is 20705 +- Maximum validated `H` size is 20705 -* Maximum validated `B*S` size is 131072 +- Maximum validated `B*S` size is 131072 -* Maximum validated `N` size is 17 +- Maximum validated `N` size is 17 ## Implementation Details The kernel implementation includes several key optimizations: -* **Dimension Packing**: Optimizes the contraction dimension by folding `N` (number of heads) into `D` (head dimension) when beneficial, improving computational efficiency. +- **Dimension Packing**: Optimizes the contraction dimension by folding `N` (number of heads) into `D` (head dimension) when beneficial, improving computational efficiency. -* **Efficient Tiling Strategy**: Uses carefully chosen tile sizes for processing batches and sequences to maximize hardware utilization. +- **Efficient Tiling Strategy**: Uses carefully chosen tile sizes for processing batches and sequences to maximize hardware utilization. -* **LNC Sharding**: Supports sharding across multiple Logical Neuron Cores (LNCs) by dividing the hidden dimension, enabling processing of larger models. +- **LNC Sharding**: Supports sharding across multiple Logical Neuron Cores (LNCs) by dividing the hidden dimension, enabling processing of larger models. -* **Memory Access Optimization**: Employs optimized memory access patterns to maximize bandwidth utilization and minimize data movement. +- **Memory Access Optimization**: Employs optimized memory access patterns to maximize bandwidth utilization and minimize data movement. -* **PSUM Bank Utilization**: Efficiently utilizes PSUM banks for accumulating partial results during matrix multiplication operations. +- **PSUM Bank Utilization**: Efficiently utilizes PSUM banks for accumulating partial results during matrix multiplication operations. -* **Stream Shuffle Broadcast**: Uses stream shuffle broadcast for bias tensors to efficiently distribute them across processing elements. +- **Stream Shuffle Broadcast**: Uses stream shuffle broadcast for bias tensors to efficiently distribute them across processing elements. -* **Specialized Engine Selection**: Alternates between scalar and vector engines for tensor copy operations to balance workload and improve performance. +- **Specialized Engine Selection**: Alternates between scalar and vector engines for tensor copy operations to balance workload and improve performance. ## See Also -* [Output Projection TKG Kernel API Reference](output-projection-tkg.md) +- [Output Projection TKG Kernel API Reference](output-projection-tkg.md) -* [QKV Kernel API Reference](qkv.md) \ No newline at end of file +- [QKV Kernel API Reference](qkv.md) diff --git a/skills/neuron-nki-docs/references/reference/library/output-projection-tkg.md b/skills/neuron-nki-docs/references/reference/library/output-projection-tkg.md index 3382de9..485e3a4 100644 --- a/skills/neuron-nki-docs/references/reference/library/output-projection-tkg.md +++ b/skills/neuron-nki-docs/references/reference/library/output-projection-tkg.md @@ -5,19 +5,19 @@ This topic provides the API reference for the `Output Projection TKG` kernel. Th The kernel supports: -* Efficient projection of attention outputs +- Efficient projection of attention outputs -* Optional bias addition +- Optional bias addition -* LNC sharding for distributed computation +- LNC sharding for distributed computation -* Optimized memory access patterns +- Optimized memory access patterns -* Head dimension packing for improved performance +- Head dimension packing for improved performance -* Flexible output tensor layouts +- Flexible output tensor layouts -* SBUF output option for kernel fusion +- SBUF output option for kernel fusion ## Background @@ -33,7 +33,7 @@ The input layouts expected for this kernel are different from those for the CTE ### output_projection_tkg -nkilib.core.output_projection.output_projection_tkg.output_projection_tkg(*attention*, *weight*, *bias*, *TRANSPOSE_OUT=False*, *OUT_IN_SB=False*) +nkilib.core.output*projection.output_projection_tkg.output_projection_tkg(\_attention*, _weight_, _bias_, _TRANSPOSE_OUT=False_, _OUT_IN_SB=False_) Output Projection Kernel optimized for Token Generation (Decode) use cases. This kernel computes `out = attention @ weight + bias`, typically used to project the output scores after an attention block in transformer models. @@ -42,15 +42,15 @@ This kernel is optimized for Token Generation (aka Decode) use cases where seque Parameters: -* **attention** (`nl.ndarray`) – Input tensor in HBM or SBUF, typically the scores output from an attention block. Shape: `[D, B, N, S]`, where `D` is head dimension, `B` is batch size, `N` is number of heads, and `S` is sequence length. Indexing: `[d, b, n, s]`. +- **attention** (`nl.ndarray`) – Input tensor in HBM or SBUF, typically the scores output from an attention block. Shape: `[D, B, N, S]`, where `D` is head dimension, `B` is batch size, `N` is number of heads, and `S` is sequence length. Indexing: `[d, b, n, s]`. -* **weight** (`nl.ndarray`) – Weight tensor in HBM. Shape: `[N*D, H]`, where `H` is hidden dimension size. Indexing: `[n * D + d, h]`. +- **weight** (`nl.ndarray`) – Weight tensor in HBM. Shape: `[N*D, H]`, where `H` is hidden dimension size. Indexing: `[n * D + d, h]`. -* **bias** (`nl.ndarray`) – Optional bias tensor in HBM. Shape: `[1, H]`. Indexing: `[1, h]`. +- **bias** (`nl.ndarray`) – Optional bias tensor in HBM. Shape: `[1, H]`. Indexing: `[1, h]`. -* **TRANSPOSE_OUT** (`bool`) – Whether to store the output in transposed shape. If `False`, output shape is `[B*S, H]` with indexing `[b*S+s, h]`. If `True`, output shape is `[H_1, H_0, H_2, B*S]` with indexing `[h_1, h_0, h_2, b*S+s]`, where `H_0 = logical core size (LNC)`, `H_1 = 128`, `H_2 = H/(H_0*H_1)`, such that `h = h_0*H_1*H_2 + h_1*H_2 + h_2`. +- **TRANSPOSE_OUT** (`bool`) – Whether to store the output in transposed shape. If `False`, output shape is `[B*S, H]` with indexing `[b*S+s, h]`. If `True`, output shape is `[H_1, H_0, H_2, B*S]` with indexing `[h_1, h_0, h_2, b*S+s]`, where `H_0 = logical core size (LNC)`, `H_1 = 128`, `H_2 = H/(H_0*H_1)`, such that `h = h_0*H_1*H_2 + h_1*H_2 + h_2`. -* **OUT_IN_SB** (`bool`) – If `True`, output is in SBUF. Else, it is written out to HBM. +- **OUT_IN_SB** (`bool`) – If `True`, output is in SBUF. Else, it is written out to HBM. Returns: Output tensor in HBM or SBUF. Shape depends on `TRANSPOSE_OUT` parameter. @@ -63,58 +63,58 @@ However, for `nl.float32`, large inputs may not fit in SBUF. **Dimensions**: -* `B`: Batch size +- `B`: Batch size -* `N`: Number of heads +- `N`: Number of heads -* `S`: Sequence length +- `S`: Sequence length -* `H`: Hidden dimension size +- `H`: Hidden dimension size -* `D`: Head dimension size +- `D`: Head dimension size **Restrictions**: -* The contract dimension of input and weight tensors must match (`N*D == weight.shape[0]`) +- The contract dimension of input and weight tensors must match (`N*D == weight.shape[0]`) -* Hidden dimension (`H`) needs to be divisible by LNC size since LNC sharding is on the weight hidden dimension +- Hidden dimension (`H`) needs to be divisible by LNC size since LNC sharding is on the weight hidden dimension -* `B*S` must be <= 128 +- `B*S` must be <= 128 -* Head dimension (`D`) must be <= 128 +- Head dimension (`D`) must be <= 128 -* When `TRANSPOSE_OUT` is `False`, `H` must be a multiple of `512*LNC` +- When `TRANSPOSE_OUT` is `False`, `H` must be a multiple of `512*LNC` -* When `TRANSPOSE_OUT` is `True`, `H` must be a multiple of `128*LNC` +- When `TRANSPOSE_OUT` is `True`, `H` must be a multiple of `128*LNC` -* When `TRANSPOSE_OUT` is `True` and using 32-bit floats, `N*H` must be <= 81920 +- When `TRANSPOSE_OUT` is `True` and using 32-bit floats, `N*H` must be <= 81920 -* When `TRANSPOSE_OUT` is `True` and using 16-bit floats, `N*H` must be <= 163840 +- When `TRANSPOSE_OUT` is `True` and using 16-bit floats, `N*H` must be <= 163840 ## Implementation Details The kernel implementation includes several key optimizations: -* **Dimension Packing**: Optimizes the contraction dimension by folding `N` (number of heads) into `D` (head dimension) when beneficial, improving computational efficiency. +- **Dimension Packing**: Optimizes the contraction dimension by folding `N` (number of heads) into `D` (head dimension) when beneficial, improving computational efficiency. -* **Efficient Tiling Strategy**: Uses carefully chosen tile sizes for processing batches and sequences to maximize hardware utilization. +- **Efficient Tiling Strategy**: Uses carefully chosen tile sizes for processing batches and sequences to maximize hardware utilization. -* **LNC Sharding**: Supports sharding across multiple Logical Neuron Cores (LNCs) by dividing the hidden dimension, enabling processing of larger models. +- **LNC Sharding**: Supports sharding across multiple Logical Neuron Cores (LNCs) by dividing the hidden dimension, enabling processing of larger models. -* **Memory Access Optimization**: Employs optimized memory access patterns to maximize bandwidth utilization and minimize data movement. +- **Memory Access Optimization**: Employs optimized memory access patterns to maximize bandwidth utilization and minimize data movement. -* **PSUM Bank Utilization**: Efficiently utilizes PSUM banks for accumulating partial results during matrix multiplication operations. +- **PSUM Bank Utilization**: Efficiently utilizes PSUM banks for accumulating partial results during matrix multiplication operations. -* **Stream Shuffle Broadcast**: Uses stream shuffle broadcast for bias tensors to efficiently distribute them across processing elements. +- **Stream Shuffle Broadcast**: Uses stream shuffle broadcast for bias tensors to efficiently distribute them across processing elements. -* **Flexible Output Layouts**: Supports both standard and transposed output layouts to accommodate different downstream kernel requirements. +- **Flexible Output Layouts**: Supports both standard and transposed output layouts to accommodate different downstream kernel requirements. -* **SBUF Output Option**: Provides the option to keep output in SBUF for fusion with subsequent operations. +- **SBUF Output Option**: Provides the option to keep output in SBUF for fusion with subsequent operations. -* **Block-based Weight Loading**: Uses block-based loading of weights to encourage prefetching and improve memory access patterns. +- **Block-based Weight Loading**: Uses block-based loading of weights to encourage prefetching and improve memory access patterns. ## See Also -* [Output Projection CTE Kernel API Reference](output-projection-cte.md) +- [Output Projection CTE Kernel API Reference](output-projection-cte.md) -* [QKV Kernel API Reference](qkv.md) \ No newline at end of file +- [QKV Kernel API Reference](qkv.md) diff --git a/skills/neuron-nki-docs/references/reference/library/qkv.md b/skills/neuron-nki-docs/references/reference/library/qkv.md index ae8fdb8..77c7229 100644 --- a/skills/neuron-nki-docs/references/reference/library/qkv.md +++ b/skills/neuron-nki-docs/references/reference/library/qkv.md @@ -5,15 +5,15 @@ This topic provides the API reference for the `QKV` kernel. The kernel performs The kernel supports: -* Optional RMSNorm/LayerNorm fusion +- Optional RMSNorm/LayerNorm fusion -* Multiple output tensor layouts +- Multiple output tensor layouts -* Residual connections from previous MLP and attention outputs +- Residual connections from previous MLP and attention outputs -* Automatic selection between TKG and CTE implementations based on batch_size * seqlen threshold +- Automatic selection between TKG and CTE implementations based on batch_size \* seqlen threshold -* Optional RoPE (Rotary Position Embedding) fusion +- Optional RoPE (Rotary Position Embedding) fusion ## Background @@ -22,8 +22,7 @@ The `QKV` kernel is a critical component in transformer architectures, responsib > **Note** > > Note -> -> +> > This kernel automatically selects between TKG (Token Generation) and CTE (Context Encoding) implementations based on sequence length threshold (currently 96), ensuring optimal performance across different use cases. CTE is used for longer sequences, while TKG is optimized for shorter sequences. ## API Reference @@ -32,7 +31,7 @@ The `QKV` kernel is a critical component in transformer architectures, responsib ### qkv -nkilib.core.qkv.qkv(*input: nl.ndarray*, *fused_qkv_weights: nl.ndarray*, *output_layout: QKVOutputLayout = QKVOutputLayout.BSD*, *bias: Optional[nl.ndarray] = None*, *fused_residual_add: Optional[[bool](https://docs.python.org/3/library/functions.html#bool)] = False*, *mlp_prev: Optional[nl.ndarray] = None*, *attention_prev: Optional[nl.ndarray] = None*, *fused_norm_type: NormType = NormType.NO_NORM*, *gamma_norm_weights: Optional[nl.ndarray] = None*, *layer_norm_bias: Optional[nl.ndarray] = None*, *norm_eps: Optional[[float](https://docs.python.org/3/library/functions.html#float)] = 1e-6*, *hidden_actual: Optional[[int](https://docs.python.org/3/library/functions.html#int)] = None*, *fused_rope: Optional[[bool](https://docs.python.org/3/library/functions.html#bool)] = False*, *cos_cache: Optional[nl.ndarray] = None*, *sin_cache: Optional[nl.ndarray] = None*, *d_head: Optional[[int](https://docs.python.org/3/library/functions.html#int)] = None*, *num_q_heads: Optional[[int](https://docs.python.org/3/library/functions.html#int)] = None*, *num_kv_heads: Optional[[int](https://docs.python.org/3/library/functions.html#int)] = None*, *store_output_in_sbuf: [bool](https://docs.python.org/3/library/functions.html#bool) = False*, *sbm: Optional[SbufManager] = None*, *use_auto_allocation: [bool](https://docs.python.org/3/library/functions.html#bool) = False*, *load_input_with_DMA_transpose: [bool](https://docs.python.org/3/library/functions.html#bool) = True*) +nkilib.core.qkv.qkv(_input: nl.ndarray_, _fused_qkv_weights: nl.ndarray_, _output_layout: QKVOutputLayout = QKVOutputLayout.BSD_, _bias: Optional[nl.ndarray] = None_, _fused_residual_add: Optional[[bool](https://docs.python.org/3/library/functions.html#bool)] = False_, _mlp_prev: Optional[nl.ndarray] = None_, _attention_prev: Optional[nl.ndarray] = None_, _fused_norm_type: NormType = NormType.NO_NORM_, _gamma_norm_weights: Optional[nl.ndarray] = None_, _layer_norm_bias: Optional[nl.ndarray] = None_, _norm_eps: Optional[[float](https://docs.python.org/3/library/functions.html#float)] = 1e-6_, _hidden_actual: Optional[[int](https://docs.python.org/3/library/functions.html#int)] = None_, _fused_rope: Optional[[bool](https://docs.python.org/3/library/functions.html#bool)] = False_, _cos_cache: Optional[nl.ndarray] = None_, _sin_cache: Optional[nl.ndarray] = None_, _d_head: Optional[[int](https://docs.python.org/3/library/functions.html#int)] = None_, _num_q_heads: Optional[[int](https://docs.python.org/3/library/functions.html#int)] = None_, _num_kv_heads: Optional[[int](https://docs.python.org/3/library/functions.html#int)] = None_, _store_output_in_sbuf: [bool](https://docs.python.org/3/library/functions.html#bool) = False_, _sbm: Optional[SbufManager] = None_, _use_auto_allocation: [bool](https://docs.python.org/3/library/functions.html#bool) = False_, _load_input_with_DMA_transpose: [bool](https://docs.python.org/3/library/functions.html#bool) = True_) QKV (Query, Key, Value) projection kernel with multiple optional fused operations. Performs matrix multiplication between hidden states and fused QKV weights matrix with optional @@ -41,49 +40,49 @@ Automatically selects between TKG and CTE implementations based on sequence leng Parameters: -* **input** (`nl.ndarray`) – Input hidden states tensor. Shape: [B, S, H] where B=batch, S=sequence_length, H=hidden_dim. +- **input** (`nl.ndarray`) – Input hidden states tensor. Shape: [B, S, H] where B=batch, S=sequence_length, H=hidden_dim. -* **fused_qkv_weights** (`nl.ndarray`) – Fused QKV weight matrix. Shape: [H, I] where I=fused_qkv_dim=(num_q_heads + 2*num_kv_heads)*d_head. +- **fused_qkv_weights** (`nl.ndarray`) – Fused QKV weight matrix. Shape: [H, I] where I=fused_qkv_dim=(num_q_heads + 2*num_kv_heads)*d_head. -* **output_layout** (`QKVOutputLayout`) – Output tensor layout. QKVOutputLayout.BSD=[B, S, I] or QKVOutputLayout.NBSd=[num_heads, B, S, d_head]. Default: QKVOutputLayout.BSD. +- **output_layout** (`QKVOutputLayout`) – Output tensor layout. QKVOutputLayout.BSD=[B, S, I] or QKVOutputLayout.NBSd=[num_heads, B, S, d_head]. Default: QKVOutputLayout.BSD. -* **bias** (`nl.ndarray`, optional) – Bias tensor to add to QKV projection output. Shape: [1, I]. +- **bias** (`nl.ndarray`, optional) – Bias tensor to add to QKV projection output. Shape: [1, I]. -* **fused_residual_add** (`bool`, optional) – Whether to perform residual addition: input = input + mlp_prev + attention_prev. Default: False. +- **fused_residual_add** (`bool`, optional) – Whether to perform residual addition: input = input + mlp_prev + attention_prev. Default: False. -* **mlp_prev** (`nl.ndarray`, optional) – Previous MLP output tensor for residual addition. Shape: [B, S, H]. +- **mlp_prev** (`nl.ndarray`, optional) – Previous MLP output tensor for residual addition. Shape: [B, S, H]. -* **attention_prev** (`nl.ndarray`, optional) – Previous attention output tensor for residual addition. Shape: [B, S, H]. +- **attention_prev** (`nl.ndarray`, optional) – Previous attention output tensor for residual addition. Shape: [B, S, H]. -* **fused_norm_type** (`NormType`) – Type of normalization (NO_NORM, RMS_NORM, RMS_NORM_SKIP_GAMMA, LAYER_NORM). Default: NormType.NO_NORM. +- **fused_norm_type** (`NormType`) – Type of normalization (NO_NORM, RMS_NORM, RMS_NORM_SKIP_GAMMA, LAYER_NORM). Default: NormType.NO_NORM. -* **gamma_norm_weights** (`nl.ndarray`, optional) – Normalization gamma/scale weights. Shape: [1, H]. Required for RMS_NORM and LAYER_NORM. +- **gamma_norm_weights** (`nl.ndarray`, optional) – Normalization gamma/scale weights. Shape: [1, H]. Required for RMS_NORM and LAYER_NORM. -* **layer_norm_bias** (`nl.ndarray`, optional) – Layer normalization beta/bias weights. Shape: [1, H]. Only for LAYER_NORM. +- **layer_norm_bias** (`nl.ndarray`, optional) – Layer normalization beta/bias weights. Shape: [1, H]. Only for LAYER_NORM. -* **norm_eps** (`float`, optional) – Epsilon value for numerical stability in normalization. Default: 1e-6. +- **norm_eps** (`float`, optional) – Epsilon value for numerical stability in normalization. Default: 1e-6. -* **hidden_actual** (`int`, optional) – Actual hidden dimension for padded tensors (if H contains padding). +- **hidden_actual** (`int`, optional) – Actual hidden dimension for padded tensors (if H contains padding). -* **fused_rope** (`bool`, optional) – Whether to apply RoPE rotation to Query and Key heads after QKV projection. Default: False. +- **fused_rope** (`bool`, optional) – Whether to apply RoPE rotation to Query and Key heads after QKV projection. Default: False. -* **cos_cache** (`nl.ndarray`, optional) – Cosine cache for RoPE. Shape: [B, S, d_head]. Required if fused_rope=True. +- **cos_cache** (`nl.ndarray`, optional) – Cosine cache for RoPE. Shape: [B, S, d_head]. Required if fused_rope=True. -* **sin_cache** (`nl.ndarray`, optional) – Sine cache for RoPE. Shape: [B, S, d_head]. Required if fused_rope=True. +- **sin_cache** (`nl.ndarray`, optional) – Sine cache for RoPE. Shape: [B, S, d_head]. Required if fused_rope=True. -* **d_head** (`int`, optional) – Dimension per attention head. Required for QKVOutputLayout.NBSd and RoPE. +- **d_head** (`int`, optional) – Dimension per attention head. Required for QKVOutputLayout.NBSd and RoPE. -* **num_q_heads** (`int`, optional) – Number of query heads. Required for RoPE. +- **num_q_heads** (`int`, optional) – Number of query heads. Required for RoPE. -* **num_kv_heads** (`int`, optional) – Number of key/value heads. Required for RoPE. +- **num_kv_heads** (`int`, optional) – Number of key/value heads. Required for RoPE. -* **store_output_in_sbuf** (`bool`) – Whether to store output in SBUF (currently unsupported, must be False). Default: False. +- **store_output_in_sbuf** (`bool`) – Whether to store output in SBUF (currently unsupported, must be False). Default: False. -* **sbm** (`SbufManager`, optional) – Optional SBUF manager for memory allocation control with pre-specified bounds for SBUF usage. +- **sbm** (`SbufManager`, optional) – Optional SBUF manager for memory allocation control with pre-specified bounds for SBUF usage. -* **use_auto_allocation** (`bool`) – Whether to use automatic SBUF allocation. Default: False. +- **use_auto_allocation** (`bool`) – Whether to use automatic SBUF allocation. Default: False. -* **load_input_with_DMA_transpose** (`bool`) – Whether to use DMA transpose optimization. Default: True. +- **load_input_with_DMA_transpose** (`bool`) – Whether to use DMA transpose optimization. Default: True. Returns: QKV projection output tensor with shape determined by output_layout. @@ -93,42 +92,42 @@ Return type: **Raises**: -* **ValueError** – Raised when contract dimension mismatch occurs between `input` and `fused_qkv_weights`. +- **ValueError** – Raised when contract dimension mismatch occurs between `input` and `fused_qkv_weights`. -* **AssertionError** – Raised when required parameters for fused operations are missing or have incorrect shapes. +- **AssertionError** – Raised when required parameters for fused operations are missing or have incorrect shapes. ## Implementation Details The kernel implementation includes several key optimizations: -* **Automatic Implementation Selection**: The kernel automatically selects between TKG (Token Generation) and CTE (Context Encoding) implementations based on sequence length threshold (currently 96). Some features like RoPE fusion and loading input with DMA transpose are only available in CTE mode. TKG mode only supports automatic allocation at the moment. +- **Automatic Implementation Selection**: The kernel automatically selects between TKG (Token Generation) and CTE (Context Encoding) implementations based on sequence length threshold (currently 96). Some features like RoPE fusion and loading input with DMA transpose are only available in CTE mode. TKG mode only supports automatic allocation at the moment. -* **Fused Operations Support**: +- **Fused Operations Support**: **Residual Addition**: Fuses `input` + `mlp_prev` + `attention_prev` -* **Normalization**: Supports RMSNorm, LayerNorm, and `RMS_NORM_SKIP_GAMMA` +- **Normalization**: Supports RMSNorm, LayerNorm, and `RMS_NORM_SKIP_GAMMA` -* **Bias Addition**: Adds bias to QKV projection output +- **Bias Addition**: Adds bias to QKV projection output -* **RoPE Fusion**: Applies Rotary Position Embedding to Query and Key heads +- **RoPE Fusion**: Applies Rotary Position Embedding to Query and Key heads -* **Flexible Output Layouts**: Supports BSD (`[B, S, I]`) and NBSd (`[num_heads, B, S, d_head`]) output tensor layouts. +- **Flexible Output Layouts**: Supports BSD (`[B, S, I]`) and NBSd (`[num_heads, B, S, d_head`]) output tensor layouts. -* **Memory Management**: +- **Memory Management**: Optional SBUF manager for controlled memory allocation -* DMA transpose optimization for weight loading +- DMA transpose optimization for weight loading -* Automatic or manual SBUF allocation modes +- Automatic or manual SBUF allocation modes -* **Hardware Compatibility**: Supports bf16, fp16, and fp32 data types (fp32 inputs are internally converted to bf16). +- **Hardware Compatibility**: Supports bf16, fp16, and fp32 data types (fp32 inputs are internally converted to bf16). -* **Constraints**: +- **Constraints**: H must be ≤ 24576 and divisible by 128 -* I must be ≤ 4096 +- I must be ≤ 4096 -* For NBSd output: d_head must equal 128 \ No newline at end of file +- For NBSd output: d_head must equal 128 diff --git a/skills/neuron-nki-docs/references/reference/library/rmsnorm-quant.md b/skills/neuron-nki-docs/references/reference/library/rmsnorm-quant.md index 579ac83..bd71494 100644 --- a/skills/neuron-nki-docs/references/reference/library/rmsnorm-quant.md +++ b/skills/neuron-nki-docs/references/reference/library/rmsnorm-quant.md @@ -5,17 +5,17 @@ This topic provides the API reference for the `RMSNorm-Quant` kernel. The kernel The kernel supports: -* Optional RMS normalization before quantization +- Optional RMS normalization before quantization -* 8-bit quantization along the last dimension of the input tensor +- 8-bit quantization along the last dimension of the input tensor -* Single program multiple data (SPMD) sharding for distributed computation +- Single program multiple data (SPMD) sharding for distributed computation -* Flexible input tensor shapes (minimum 2 dimensions) +- Flexible input tensor shapes (minimum 2 dimensions) -* Input validation with configurable dimension limits +- Input validation with configurable dimension limits -* Lower bound clipping for numerical stability +- Lower bound clipping for numerical stability ## Background @@ -35,10 +35,10 @@ RMS Norm Quantization Kernel arguments. lower_bound*: [float](https://docs.python.org/3/library/functions.html#float)* Non-negative float used for clipping input values and scale. -norm_type*: NormType** = NormType.RMS_NORM* +norm_type*: NormType\*\* = NormType.RMS_NORM* Normalization type to use [`RMS_NORM`, `NO_NORM`] -eps*: [float](https://docs.python.org/3/library/functions.html#float)** = 1e-6* +eps*: [float](https://docs.python.org/3/library/functions.html#float)\*\* = 1e-6* Epsilon value for numerical stability, model hyperparameter needs_rms_normalization() → [bool](https://docs.python.org/3/library/functions.html#bool) @@ -49,26 +49,26 @@ Returns True if a positive lower bound is specified, False otherwise. **Raises**: -* **AssertionError** – Raised when unsupported normalization types are used, negative bounds are provided, or invalid epsilon values are specified. +- **AssertionError** – Raised when unsupported normalization types are used, negative bounds are provided, or invalid epsilon values are specified. ### rmsnorm_quant_kernel -nkilib.core.rmsnorm_quant.rmsnorm_quant.rmsnorm_quant_kernel(*hidden: nl.ndarray*, *ln_w: nl.ndarray*, *kargs: [RmsNormQuantKernelArgs](#nkilib.core.rmsnorm_quant.rmsnorm_quant.RmsNormQuantKernelArgs)*) +nkilib.core.rmsnorm*quant.rmsnorm_quant.rmsnorm_quant_kernel(\_hidden: nl.ndarray*, _ln_w: nl.ndarray_, _kargs: [RmsNormQuantKernelArgs](#nkilib.core.rmsnorm_quant.rmsnorm_quant.RmsNormQuantKernelArgs)_) Entrypoint NKI kernel that performs one of the following: -* Perform RMSNorm and quantize the normalized hidden over the hidden dimension (`H`, or `axis=-1`). +- Perform RMSNorm and quantize the normalized hidden over the hidden dimension (`H`, or `axis=-1`). -* Quantize hidden over dimension `H`. +- Quantize hidden over dimension `H`. The kernel supports no specialization, or specialization along 1 dimension (1D SPMD grid). Parameters: -* **hidden** (`nl.ndarray`) – Input hidden states tensor with minimum 2 dimensions. For 3D inputs, expected layout is `[B, S, H]`. For 2D inputs, layout is `[outer_dim, processing_dim]` where outer_dim is the product of all major dimensions. +- **hidden** (`nl.ndarray`) – Input hidden states tensor with minimum 2 dimensions. For 3D inputs, expected layout is `[B, S, H]`. For 2D inputs, layout is `[outer_dim, processing_dim]` where outer_dim is the product of all major dimensions. -* **ln_w** (`nl.ndarray`) – Gamma multiplicative bias vector with `[H]` or `[1, H]` layout. Required when RMS normalization is enabled. +- **ln_w** (`nl.ndarray`) – Gamma multiplicative bias vector with `[H]` or `[1, H]` layout. Required when RMS normalization is enabled. -* **kargs** (`RmsNormQuantKernelArgs`) – Kernel arguments specifying normalization type, bounds, and epsilon values. See [`RmsNormQuantKernelArgs`](#nkilib.core.rmsnorm_quant.rmsnorm_quant.RmsNormQuantKernelArgs) for details. +- **kargs** (`RmsNormQuantKernelArgs`) – Kernel arguments specifying normalization type, bounds, and epsilon values. See [`RmsNormQuantKernelArgs`](#nkilib.core.rmsnorm_quant.rmsnorm_quant.RmsNormQuantKernelArgs) for details. Returns: Output tensor with shape `[..., H + 4]` on HBM where the last dimension is extended by 4 elements. The first H elements store the possibly normalized and quantized tensor, while the last 4 elements store fp8 floats that can be reinterpreted as fp32 dequantization scales. @@ -78,45 +78,44 @@ Return type: **Constraints**: -* Input tensor must have at least 2 dimensions +- Input tensor must have at least 2 dimensions -* For 3D inputs: batch dimension ≤ MAX_B, sequence length ≤ MAX_S, hidden dimension ≤ MAX_H +- For 3D inputs: batch dimension ≤ MAX_B, sequence length ≤ MAX_S, hidden dimension ≤ MAX_H -* For 2D inputs: processing dimension ≤ MAX_H, outer dimension ≤ MAX_B × MAX_S +- For 2D inputs: processing dimension ≤ MAX_H, outer dimension ≤ MAX_B × MAX_S -* When RMS normalization is enabled, ln_w must have shape [H] or [1, H] where H matches the processing dimension +- When RMS normalization is enabled, ln_w must have shape [H] or [1, H] where H matches the processing dimension -* Supports 1D SPMD grid or no specialization +- Supports 1D SPMD grid or no specialization > **Note** > > Note -> -> +> > The autocast argument may NOT be respected properly. The kernel automatically handles dimension validation and provides detailed error messages for constraint violations. ## Implementation Details The kernel implementation includes several key optimizations: -* **Input Tensor Outer Dimension Collapse**: All major dimensions are collapsed into one for simplification, allowing the kernel to process along the minor dimension efficiently. +- **Input Tensor Outer Dimension Collapse**: All major dimensions are collapsed into one for simplification, allowing the kernel to process along the minor dimension efficiently. -* **Tiling**: The kernel is tiled on the major dimension by a size equal to the hardware’s maximum partition dimension, ensuring full utilization of the hardware engines’ input width. +- **Tiling**: The kernel is tiled on the major dimension by a size equal to the hardware’s maximum partition dimension, ensuring full utilization of the hardware engines’ input width. -* **SBUF/PSUM Allocation**: Uses Stack Allocator for consistent and deterministic memory allocations within the kernel scope. +- **SBUF/PSUM Allocation**: Uses Stack Allocator for consistent and deterministic memory allocations within the kernel scope. -* **SPMD Sharding**: Supports splitting computation across the constituent cores of a Logical Neuron Core by sharding on the outer-most dimension with automatic load balancing for non-divisible dimensions. +- **SPMD Sharding**: Supports splitting computation across the constituent cores of a Logical Neuron Core by sharding on the outer-most dimension with automatic load balancing for non-divisible dimensions. -* **Gamma Broadcast**: Improves pipeline parallelism by distributing work to the TensorEngine through matrix multiplication against a vector of ones. +- **Gamma Broadcast**: Improves pipeline parallelism by distributing work to the TensorEngine through matrix multiplication against a vector of ones. -* **Activation Reduce**: Uses specialized instructions to perform reduce-add operations efficiently along with square operations. +- **Activation Reduce**: Uses specialized instructions to perform reduce-add operations efficiently along with square operations. -* **Optimized Batch Processing**: Processes tiles in batches of 8 for improved efficiency, with remainder handling for non-divisible cases. +- **Optimized Batch Processing**: Processes tiles in batches of 8 for improved efficiency, with remainder handling for non-divisible cases. -* **Input Validation**: Comprehensive validation of tensor dimensions against hardware limits (MAX_B, MAX_S, MAX_H) with detailed error messages. +- **Input Validation**: Comprehensive validation of tensor dimensions against hardware limits (MAX_B, MAX_S, MAX_H) with detailed error messages. -* **Numerical Stability**: Implements lower bound clipping and minimum dequantization scale clamping to prevent numerical instabilities. +- **Numerical Stability**: Implements lower bound clipping and minimum dequantization scale clamping to prevent numerical instabilities. ## See Also -* [RMSNorm-Quant Kernel Design Specification](design-rmsnorm-quant.md) \ No newline at end of file +- [RMSNorm-Quant Kernel Design Specification](design-rmsnorm-quant.md) diff --git a/skills/neuron-nki-docs/references/reference/migration/nki-030-update-guide.md b/skills/neuron-nki-docs/references/reference/migration/nki-030-update-guide.md index 79b0d58..092d884 100644 --- a/skills/neuron-nki-docs/references/reference/migration/nki-030-update-guide.md +++ b/skills/neuron-nki-docs/references/reference/migration/nki-030-update-guide.md @@ -25,11 +25,13 @@ NKI 0.3.0 introduces `nki.simulate(kernel)`, which executes NKI kernels entirely The simulator can be invoked in two ways: 1. Set the environment variable `NKI_SIMULATOR=1` to run existing kernels without code changes: + ``` NKI_SIMULATOR=1 python my_script.py ``` 2. Wrap the kernel call with `nki.simulate`: + ```python import nki import numpy as np @@ -317,12 +319,12 @@ buf = nl.ndarray((128, 512), dtype=nl.float16, buffer=nl.sbuf) **Buffer type mapping:** | NKI 0.2.0 (string) | NKI 0.3.0 (object) | -|---------------------|---------------------| -| "sbuf" | nl.sbuf | -| "psum" | nl.psum | -| "hbm" | nl.hbm | -| "private_hbm" | nl.private_hbm | -| "shared_hbm" | nl.shared_hbm | +| ------------------ | ------------------ | +| "sbuf" | nl.sbuf | +| "psum" | nl.psum | +| "hbm" | nl.hbm | +| "private_hbm" | nl.private_hbm | +| "shared_hbm" | nl.shared_hbm | ### nki.isa.dma_engine Alias Repurposed @@ -332,7 +334,7 @@ The NKI 0.2.0 `nki.isa.dma_engine` module-level alias was unused and did not map The NKI 0.3.0 compiler has stricter validation. The following patterns require changes for NKI 0.3.0. -### Remove Keyword-Only Argument Separator (*) +### Remove Keyword-Only Argument Separator (\*) The NKI 0.3.0 compiler does not support the `*` separator in kernel function signatures. Move all parameters with defaults to the end of the signature. diff --git a/skills/neuron-nki-docs/references/reference/migration/nki-060-dynamic-loop-migration-guide.md b/skills/neuron-nki-docs/references/reference/migration/nki-060-dynamic-loop-migration-guide.md index 9e026e4..6a24448 100644 --- a/skills/neuron-nki-docs/references/reference/migration/nki-060-dynamic-loop-migration-guide.md +++ b/skills/neuron-nki-docs/references/reference/migration/nki-060-dynamic-loop-migration-guide.md @@ -7,10 +7,10 @@ Migrate on-device dynamic loops from `for i in nl.dynamic_range(...)` and bare The NKI frontend is moving from **Parsing** to **Tracing**: -| Frontend | Status | -|----------|--------| -| Parsing (legacy) | current default | -| Tracing | available in NKI **0.6.0**, default in **0.7.0**, parser removed in **0.8.0** | +| Frontend | Status | +| ---------------- | ----------------------------------------------------------------------------- | +| Parsing (legacy) | current default | +| Tracing | available in NKI **0.6.0**, default in **0.7.0**, parser removed in **0.8.0** | Tracing removes support for `for i in nl.dynamic_range(...)` and bare `while reg:`. These forms build an on-device loop from a runtime register; under @@ -26,12 +26,12 @@ safe and lets you fall back to the parser during the transition ## Do I need to migrate? -| Construct | Migrate? | -|-----------|----------| -| `for i in nl.dynamic_range(...)` | **Yes** | -| bare `while reg:` (register condition) | **Yes** | +| Construct | Migrate? | +| ------------------------------------------------------------- | ----------------------------------- | +| `for i in nl.dynamic_range(...)` | **Yes** | +| bare `while reg:` (register condition) | **Yes** | | `nl.affine_range` / `nl.sequential_range` / `nl.static_range` | No — compile-time loops, unaffected | -| `range(...)` over a Python int | No | +| `range(...)` over a Python int | No | Only loops whose bound or condition is a **runtime hardware register** are affected. @@ -44,12 +44,14 @@ loop body becomes a callable that receives the iteration value as a (standard Python LEGB scoping). **Before (parser-only, removed under tracing):** + ```python for i in nl.dynamic_range(reg): nisa.dma_copy(dst=temp, src=data.ap(scalar_offset=i, indirect_dim=1)) ``` **After (parser + tracer):** + ```python def body(i): nisa.dma_copy(dst=temp, src=data.ap(scalar_offset=i, indirect_dim=1)) @@ -67,6 +69,7 @@ condition register. It is a true `while` (skips the body entirely if `init` is zero), not a do-while. **Before (parser-only, removed under tracing):** + ```python while reg: nisa.tensor_tensor(dst=acc, data1=acc, data2=val, op=nl.add) @@ -75,6 +78,7 @@ while reg: ``` **After (parser + tracer):** + ```python def body(r): nisa.tensor_tensor(dst=acc, data1=acc, data2=val, op=nl.add) @@ -94,7 +98,9 @@ nkilib refactor): for VAR in nl.dynamic_range(LB, UB): BODY ``` + becomes + ``` def _fori_body_N(VAR): BODY diff --git a/skills/neuron-nki-docs/references/reference/migration/nki-migration-guide.md b/skills/neuron-nki-docs/references/reference/migration/nki-migration-guide.md index 5e8b3ae..6ee52dd 100644 --- a/skills/neuron-nki-docs/references/reference/migration/nki-migration-guide.md +++ b/skills/neuron-nki-docs/references/reference/migration/nki-migration-guide.md @@ -22,7 +22,6 @@ decorator. Just as before, you mark your NKI kernels with the `nki.jit` decorator. However, unlike before, the functions under this decorator will be passed to the NKI Compiler and not be evaluated by the Python interpreter. - ```python def a_function(x,y,z): # this is Python code @@ -32,7 +31,6 @@ def kernel(x,y,z): # this is NKI code ``` - If you use Python features within a NKI kernel that are not supported, the NKI Compiler will give an error. The goal is that programming in NKI is intuitive and convenient and all of the features you need are available and behave as @@ -49,35 +47,35 @@ These are the key items to migrate existing kernel to the Beta 2 NKI Compiler. ### What new features are available in NKI Beta 2? -* A new namespace for NKI Beta 2, `nki.*` +- A new namespace for NKI Beta 2, `nki.*` -* `device_print` is available to inspect tensor values +- `device_print` is available to inspect tensor values -* The behavior of loops and branching is consistent with regular Python +- The behavior of loops and branching is consistent with regular Python -* Lists and dictionaries are available and their behavior in loops is consistent with regular Python +- Lists and dictionaries are available and their behavior in loops is consistent with regular Python -* Direct allocation APIs have been reworked +- Direct allocation APIs have been reworked ### What features in `neuronxcc.nki.*` are not available in `nki.*`? -* `arange` has been removed, use slicing or [NKI Access Patterns](../../programming/nki-aps.md#nki-aps) +- `arange` has been removed, use slicing or [NKI Access Patterns](../../programming/nki-aps.md#nki-aps) -* The `mask` parameter is no longer supported +- The `mask` parameter is no longer supported -* Block dimensions of tensors have been removed +- Block dimensions of tensors have been removed -* Explicit `dst` parameter is now required for `nki.isa` instructions and is always the first argument +- Explicit `dst` parameter is now required for `nki.isa` instructions and is always the first argument -* `nl.load` and `nl.store` have been removed, use `nisa.dma_copy` +- `nl.load` and `nl.store` have been removed, use `nisa.dma_copy` -* Nested slicing is not available +- Nested slicing is not available -* Dynamic Access syntax has changed +- Dynamic Access syntax has changed -* Decorators on sub-kernels need to be removed +- Decorators on sub-kernels need to be removed -* Dictionaries support only string keys +- Dictionaries support only string keys ## New Features in NKI Beta 2 @@ -89,7 +87,6 @@ supports both versions of the language via namespaces. The Beta 1 APIs can be used via the `neuronxcc.nki.*` namespace, while Beta 2 has moved to the `nki.*` namespace. - ```python # Legacy Beta 1 APIs import neuronxcc.nki as nki @@ -100,7 +97,6 @@ import nki import nki.isa as nisa ``` - We have made improvements to the APIs, like consistent naming, order of arguments, and matching more closely the hardware ISA so that what developers write in NKI and what they see in the profiler are the same. There is one @@ -112,20 +108,16 @@ destination parameter. In Beta 2, all of the ISA functions now require a `dst` parameter instead of returning a result. So, instead of writing: - ```python result[...] = nisa.reciprocal(src) ``` - Developers must write: - ```python nisa.reciprocal(dst=result[...], src) ``` - This change makes the behavior of the APIs more consistent and matches cases where APIs may perform accumulation or return multiple results. It also helps avoid scenarios where developers might inadvertently write to the wrong buffer @@ -137,7 +129,7 @@ or inadvertently introduce additional copy operations. > 0.6.0, default in 0.7.0, parser removed in 0.8.0), the `for i in dynamic_range(...)` and bare > `while reg:` forms below become **parser-only and are removed under tracing** — a runtime register > has no value at trace time. Migrate to the structured constructs `nl.fori_loop(lower, upper, -> body_fun, step=1)` (counted loop with a runtime bound) and `nl.while_loop(init, body_fun)` +body_fun, step=1)` (counted loop with a runtime bound) and `nl.while_loop(init, body_fun)` > (data-dependent loop); both compile on the parser and the tracer. See the > [NKI 0.6.0 Dynamic Loop Migration Guide](nki-060-dynamic-loop-migration-guide.md) for > before/after patterns, rules, and the mechanical transform recipe. @@ -155,7 +147,6 @@ To support dynamic control flow, NKI has a new set of `nki.isa` APIs for reading and writing to hardware registers. See [NKI API Reference Manual](../../programming/api/index.md) for more information. - ```python # Define a register def register_alloc(x: Optional[int]) -> register: ... @@ -170,12 +161,10 @@ def register_load(dst: register, src: tensor): ... def register_store(dst: tensor, src: register): ... ``` - The most basic dynamic loop is a `for` loop that uses a register value for the iteration value and another register for the upper bound. Developers can write this kind of loop using `dynamic_range`: - ```python # dynamic loop with dynamically computed upper bounds # upper_bound is a hardware register @@ -186,11 +175,9 @@ for i in dynamic_range(5, upper_bound, 2): ... ``` - Developers can also write dynamic while loops. When using a dynamic while loop, the developer should update the register within the body of the loop. - ```python # initialize a conditional tensor which will be updated in the loop cond = nl.ndarray((1, 1), buffer=nl.sbuf, dtype=np.int32) @@ -208,14 +195,12 @@ while reg: # loop will terminate when the value reaches 0 nisa.register_load(reg, cond) ``` - ### Update indexing syntax for `mgrid` and `arange` If using `nl.mgrid/arange` to access continuous elements in an existing NKI kernel, this should be replaced with integer slicing. Take a look at the following example. - ```python # Example 1 t = nl.ndarray(shape=(128, 16, 64), ...) @@ -234,7 +219,6 @@ t[i_p, if0*64+i_f1] t[0:128, 0:8*64] ``` - If your use case cannot be represented with the slicing syntax above, see [NKI Access Patterns](../../programming/nki-aps.md#nki-aps). @@ -248,7 +232,6 @@ by the Python evaluator, which could lead to some surprising results. For exampl in the code below, the normal Python variable `var` ends up with a value of 1 rather than the expected value of 8. This has been solved in the new NKI Compiler. - ```python val = 0 for i in range(8): @@ -256,13 +239,11 @@ for i in range(8): print(val) # will print 1 in Beta 1, prints 8 in Beta 2 ``` - For similar reasons, sometimes Python control flow constructs, such as `if` statements, could not be handled properly when nested within a `for` loop. For example, in Beta 1 the code below produces an undefined result. In Beta 2, this code produces the expected result. - ```python val = 0 for i in range(8): @@ -273,7 +254,6 @@ for i in range(8): print(val) # undefined behaviour in Beta 1, prints 2 in Beta 2 ``` - Many other examples of troublesome control flow have been fixed, which should make using NKI easier and more intuitive. @@ -287,7 +267,6 @@ used to avoid out-of-bounds access. For example, suppose a developer is tiling a tensor of size 129 x 513, and you want to use tiles of size 128 x 512. A typical way to write a tiling loop in Beta 1 is shown below. - ```python t = nl.ndarray(shape=(129, 513), ...) result = nl.ndarray(shape=(129, 513), ...) @@ -298,7 +277,6 @@ for i in range(2): mask=(i_p+128*i<129) & (i_f+512*i<513)) ``` - Note, when `i` (or `j`) is equal to 1, then the index expression `result[i_p+128*i, i_f+512*i]` would overflow the tensor dimension. The mask expression `mask=(i_p+128*i<129) & (i_f+512*i<513)` modifies the indexing so @@ -310,7 +288,6 @@ In NKI Beta 2, developers can use standard constructs from Python such as `min` and `slice` to build indexing expressions that are in bounds for the tensor. For example, the above code can now be written as: - ```python for i in range(2): p_start = i * 128 @@ -325,17 +302,14 @@ for i in range(2): nisa.tensor_copy(result[p, f], t[p, f]) ``` - The developer may also choose to inline the slices, if that is more natural. The below syntax is common in NKI Beta 1. - ```python nisa.tensor_copy(result[p_start:p_end, f_start:f_end], t[p_start:p_end, f_start:f_end]) ``` - ### Improved Allocation API The manual allocation API has been simplified. In Beta 2 the there is a new @@ -345,14 +319,12 @@ corresponds to a physical partition lane on the hardware, the free dimension off is the element offset within each partition. The free dimension offset is translated into physical SBUF address in the compiler. - ```python # creates your buffer on parition 0, offset by 128 elements of your data type a_result = nl.ndarray(dtype=a.dtype, shape=a.shape, name="result", address=(0, 128), buffer=nl.sbuf) ``` - The address space for PSUM is now also 2D to be consistent with the hardware. Recall that PSUM on NeuronCore v2/v3/v4 is organized into 128 partitions, each consisting of 16KB of memory. Each partition is further divided into 8 PSUM banks, @@ -362,7 +334,6 @@ error otherwise. For example, the following code will allocate a PSUM tensor on bank 3: - ```python bank_id = 3 PSUM_BANK_SIZE = 2048 @@ -370,7 +341,6 @@ psum_t = nl.ndarray(dtype=nl.bfloat16, shape=(128, 1024), address=(0, bank_id*PSUM_BANK_SIZE)) ``` - ### Translate from the Beta 1 Direct Allocation API To translate the direct allocated kernel in Beta 1, all data structures must @@ -382,7 +352,6 @@ multi-dimensional tensors for the rest of your dimensions. See After this, translate the address of each block. For example, given the following tensor in the Beta 1 that uses the modular allocation. - ```python # beta 1 - uses block dimension and mod allocator k_loaded = nl.ndarray((num_512_tiles_cur_section, nl.par_dim(p_k), n_k), @@ -390,11 +359,9 @@ k_loaded = nl.ndarray((num_512_tiles_cur_section, nl.par_dim(p_k), n_k), buffer=sb_mod(base_addr=sca, num_free_tiles=(num_512_tiles_cur_section, ) ``` - Now with Beta 2, developers can translate the block dimension into a list and compute the address for each block. - ```python # beta 2 - use lists of tensors and get lists of virtual byte addresses k_loaded_tensors = [] @@ -403,7 +370,6 @@ for i in range(num_512_tiles_cur_section): buffer=nl.sbuf, address=(0, sca + (i%num_512_tiles_cur_section)*n_k*2 ) ) ``` - ### Remove nki.jit decorator on sub-kernels For kernels that call other kernels, or call any other functions that are @@ -418,7 +384,6 @@ needing to inherent from `nl.NKIObject` thrown from the callsite of the sub-kern If a kernel is being called by another kernel and it is also called standalone, the decorator can be applied on-the-fly at the call site to avoid this problem. - ```python # Do not apply the decorator on the kernel definition def my_kernel(...): @@ -430,7 +395,6 @@ kernel_decorated = nki.jit(my_kernel) result = kernel_decorated(a) ``` - ### Translation of Block Dimensions If the kernel uses block dimension, defined as a tensor with a partition @@ -445,7 +409,6 @@ Block dimension of tensors in Beta 1 was syntactic sugar for a list of tensors managed by the compiler. In NKI Beta 2, users can directly code this patten using standard lists, without extra compiler support. - ```python # Before migration t = nl.ndarray((8, nl.par_dim(128), 256), dtype=nl.float32, buffer=nl.sbuf) @@ -461,7 +424,6 @@ for i in range(8): t_list[i] ``` - With this approach, the programs generated before and after migration are identical and should yield the same performance. @@ -470,7 +432,6 @@ identical and should yield the same performance. If blocks need to be alive at the same time, move the block dimension into free dimension - ```python a = nl.ndarray((8, par_dim(128), 512), buffer=nl.sbuf, dtype=bfloat16) @@ -478,11 +439,9 @@ a = nl.ndarray((8, par_dim(128), 512), buffer=nl.sbuf, dtype=bfloat16) a = nl.ndarray((128, 8, 512), buffer=nl.sbuf, dtype=bfloat16) ``` - As an example, if all 8 blocks of add_buf need to be live at the same time, then the block dimension needs to be folded into the free dimension. - ```python @nki.jit def sb_blocks(inp): @@ -506,11 +465,9 @@ def sb_blocks_migrated(inp): return res ``` - If blocks do not need to be alive at the same time, remove the block dimension and relocate tensor declaration. - ```python a = nl.ndarray((8, par_dim(128), 256)) for i in nl.affine_range(8): @@ -522,12 +479,10 @@ for i in nl.affine_range(8): ``` - As an example, if all 8 blocks of add_buf do not need to be live at the same time, then remove the block dimension and relocate the tensor declaration inside the loop. - ```python @nki.jit def sb_blocks(inp): @@ -549,12 +504,10 @@ def sb_blocks_migrated(inp): return res ``` - It is important to note that the dependency relationship between loop iterations is different in `sb_blocks_migrated` and the following `sb_blocks_migrated_incorrect` shown below. - ```python @nki.jit def sb_blocks_migrated_incorrect(inp): @@ -566,7 +519,6 @@ def sb_blocks_migrated_incorrect(inp): return res ``` - In `sb_blocks_migrated`, the compiler could unroll the loop and materialize multiple copies of the tensor `add_buf`. However, in the `sb_blocks_migrated_incorrect`, the execution will be serialized because the loop carries a dependency on `add_buf`. @@ -580,7 +532,6 @@ The syntax for representing dynamic access patterns has changed. In NKI Beta 1, an access with a dynamic scalar offset could be represented as shown below where `batch_idx` is a dynamic value in the SBUF: - ```python batch_idx = nl.multiply(nl.bitwise_and(nl.load(dynamic_idx), y=3), 128) result = nl.ndarray((128, 256), A.dtype, buffer=nl.shared_hbm) @@ -589,13 +540,11 @@ i_p, i_f = nl.mgrid[0:128, 0:256] nisa.dma_copy(src=A[batch_idx, i_p, i_f], dst=result[...]) ``` - #### Scalar Dynamic Access In Beta 2, we need to use a physical access pattern, specified with the `.ap` method, to represent this. - ```python def indirect_scalar_dynamic_dma(A): # Assume input A is of shape (4*128, 512). We want to copy from A[3*128:, 0:256] @@ -615,34 +564,28 @@ def indirect_scalar_dynamic_dma(A): return result ``` - The `scalar_offset` is an SBUF value that specifies the index on the `indirect_dim` of the tensor. For example, the code block above accesses `batch_idx` on the 0-th dimension of the tensor `A`. This example will access the memory from `A` starting at the element offset below. - ```python # prod(A.shape[indirect_dim+1:]) is the accumulated shape # to the right of indirect_dim offset + scalar_offset * prod(A.shape[indirect_dim+1:]) ``` - In the example above, the access would start from: - ```python 0 + batch_idx * 512 ``` - In conventional NumPy syntax, the above means that we will are accessing `A[batch_idx:batch_idx+128, 0:256]`. Writing this in the canonical loop form, the result of the access is the following: - ```python result = nl.ndarray(shape=(128, 256), dtype=A.dtype, buffer=nl.sbuf) for x in range(128): @@ -650,13 +593,11 @@ for x in range(128): result[x, y] = A.flatten()[0 + batch_idx*512 + x*512 + y*1] ``` - #### Vector Dynamic Access Vector dynamic access is similar to that of scalar, except that we need to specify the field `vector_offset`. Currently only `indirect_dim=0` is supported. - ```python def indirect_vector_dynamic_dma(A): # shape of A is (128, 512) @@ -675,10 +616,8 @@ def indirect_vector_dynamic_dma(A): return result_hbm ``` - For this particular case, the semantics of the access are: - ```python indirect_dimension = 0 @@ -693,10 +632,8 @@ for w in range(64): ] ``` - In general, the semantics are as follows. (Note: `indirect_dimension=0` is the only supported configuration at the moment). - ```python # For access pattern [s3, W],[s2, Z],[s1, Y],[s0,X], with vector offset indirect_tensor @@ -717,9 +654,8 @@ for w in range(W): ] ``` - ### Further reading -* [About the NKI Compiler](../../programming/nki-compiler.md) +- [About the NKI Compiler](../../programming/nki-compiler.md) -* [NKI API Reference Manual](../../programming/api/index.md) \ No newline at end of file +- [NKI API Reference Manual](../../programming/api/index.md) diff --git a/skills/neuron-nki-docs/references/reference/migration/nki_block_dimension_migration_guide.md b/skills/neuron-nki-docs/references/reference/migration/nki_block_dimension_migration_guide.md index c94ca1f..520ca88 100644 --- a/skills/neuron-nki-docs/references/reference/migration/nki_block_dimension_migration_guide.md +++ b/skills/neuron-nki-docs/references/reference/migration/nki_block_dimension_migration_guide.md @@ -3,15 +3,15 @@ NKI Block Dimension Migration Guide The SBUF/PSUM tensors in NKI used to allow block dimensions in front of the partition dimension. The block dimension support has been removed due the following reasons. -* Removing block dimensions does not hurt the expressivity of NKI. +- Removing block dimensions does not hurt the expressivity of NKI. -* Block dimension is a pure software concept and does not have direct hardware mapping. +- Block dimension is a pure software concept and does not have direct hardware mapping. -* The block dimension is unintuitive and causes confusion. +- The block dimension is unintuitive and causes confusion. -* Using block dimension has no inherit performance benefit, particularly using block dimension has no relationship with memory throughput whatsoever. +- Using block dimension has no inherit performance benefit, particularly using block dimension has no relationship with memory throughput whatsoever. -* Multi-buffering is implicit with block dimension. Removing block dimension will make multi-buffering more natural. +- Multi-buffering is implicit with block dimension. Removing block dimension will make multi-buffering more natural. This document will first explain the semantics of block dimensions in detail, then it will provide information on how to migrate existing code that uses block dimensions while maintain the functional correctness and performance. @@ -19,7 +19,6 @@ This document will first explain the semantics of block dimensions in detail, th Consider the following NKI tensor. - ```python a = nl.ndarray((4, 8, nl.par_dim(128), 2, 512), buffer=nl.sbuf) @@ -28,12 +27,10 @@ a = nl.ndarray((4, 8, nl.par_dim(128), 2, 512), buffer=nl.sbuf) # - (2, 512): (F) free dimension ``` - A NKI tensor has three types of dimensions: (B, P, F) . The partition dimension maps to the partition dimension of the physical memory, and the free dimensions describe how data is organized in each SBUF/PSUM partition. The block dimensions described how many physical (P, F) tiles the tensor has. The block dimension of tensors is a **logical** dimension and is a pure software concept. The compiler analyzes the memory dependency and allocates physical address to each tiles. **This means that the physical tiles may not be alive in the memory simultaneously**, and in most of the cases they don not. Consider the following code snippet that access the tensor a. - ```python @nki.jit def exp_func(inp): @@ -47,10 +44,8 @@ def exp_func(inp): nl.store(output[i, j], value=result) ``` - At the very minimum, only 1 physical tile of a needs to be alive. Then the execution is completely serialized. Essentially, all physical tiles would have the exact same memory address. - ```python Physical Address Map @@ -59,10 +54,8 @@ output[0, 1] --> Partition 0 - 128, Free 0 - 2048B ... ``` - Instead, compiler could choose to allocate 2 physical tiles to a, then the dma copy from HBM to SBUF can overlap with the exponential operation. In other word, **the block dimension allows compiler to perform space-time tradeoff at liberty.** - ```python Physical Address Map @@ -73,14 +66,12 @@ output[0, 3] --> Partition 0 - 128, Free 2048 - 4096B ... ``` - When performing the migration, it is important to understand the dependency relationship between blocks and choose the correct migration method accordingly. ## Migration for SBUF tensors ### If blocks need to be alive at the same time, move the block dimension into free dimension - ```python a = nl.ndarray((8, par_dim(128), 512), buffer=nl.sbuf, dtype=bfloat16) @@ -88,10 +79,8 @@ a = nl.ndarray((8, par_dim(128), 512), buffer=nl.sbuf, dtype=bfloat16) a = nl.ndarray((128, 8, 512), buffer=nl.sbuf, dtype=bfloat16) ``` - As an example, all 8 blocks of `add_buf` needs to be alive at the same time when the first for loop finishes. Therefore, the block dimension need to be fold into the free dimension. - ```python @nki.jit def sb_blocks(inp): @@ -115,10 +104,8 @@ def sb_blocks_migrated(inp): return res ``` - ### If blocks does not need to be alive at the same time, remove the block dimension and hoist it down - ```python a = nl.ndarray((8, par_dim(128), 256)) for i in nl.affine_range(8): @@ -130,10 +117,8 @@ for i in nl.affine_range(8): ``` - As an example, all 8 blocks of `add_buf` does not need to be alive at the same time. We can remove the block dimension and hoist down the tensor inside the loop. - ```python @nki.jit def sb_blocks(inp): @@ -155,17 +140,14 @@ def sb_blocks_migrated(inp): return res ``` - > **Note** > > Warning -> -> +> > To preserve performance, it is important to hoist down the tensor inside the loop. It is important to note that the dependency relationship betweens loop iterations is different in `sb_blocks_migrated` and the following `sb_blocks_migrated_incorrect`. - ```python @nki.jit def sb_blocks_migrated_incorrect(inp): @@ -177,7 +159,6 @@ def sb_blocks_migrated_incorrect(inp): return res ``` - In `sb_blocks_migrated`, compiler could unroll the loop and materialize multiple copies of the tensor `add_buf`. However, in the `sb_blocks_migrated_incorrect`, the execution will be serialized because the loop carries dependency on `add_buf`. ## Migration for PSUM tensors @@ -185,15 +166,13 @@ In `sb_blocks_migrated`, compiler could unroll the loop and materialize multiple > **Note** > > Note -> -> +> > To be filled, the backend support for removing blocks in PSUM tensor is still in progress. ## Migration of direct allocation & multi-buffering When we have block dimensions, we allocate interleaved address for blocks to achieve multi-buffering. - ```python def interleave_alloc_func(idx, pdim_size, fdim_size): """ @@ -220,10 +199,8 @@ def copy_func(inp): nl.store(output[i], value=a[i]) ``` - After removing the block dimension, we could write the following to implement the same multi-buffering, which is actually more natural and closer to that on CPU. - ```python def interleave_alloc_func(idx, pdim_size, fdim_size): """ @@ -249,4 +226,4 @@ def exp_func(inp): for i in range(4): a[0:128, i % 2, 0:512] = nl.load(inp[i]) nl.store(output[i], value=a[0:128, i % 2, 0:512]) -``` \ No newline at end of file +``` diff --git a/skills/neuron-nki-docs/references/reference/nki_faq.md b/skills/neuron-nki-docs/references/reference/nki_faq.md index aeff1d3..fdd479c 100644 --- a/skills/neuron-nki-docs/references/reference/nki_faq.md +++ b/skills/neuron-nki-docs/references/reference/nki_faq.md @@ -85,4 +85,4 @@ Architecture) APIs in upcoming Neuron releases. The [NKI APIs](../programming/api/index.md) follow the Neuron Software Maintenance policy for Neuron APIs. For more information, see the -[SDK Maintenance Policy](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/about-neuron/sdk-policy.html). \ No newline at end of file +[SDK Maintenance Policy](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/about-neuron/sdk-policy.html). diff --git a/skills/neuron-nki-docs/references/reference/nki_rn.md b/skills/neuron-nki-docs/references/reference/nki_rn.md index e7dfd48..38bbbf5 100644 --- a/skills/neuron-nki-docs/references/reference/nki_rn.md +++ b/skills/neuron-nki-docs/references/reference/nki_rn.md @@ -8,333 +8,333 @@ Date: 2026 NKI 0.3.0 moves NKI to General Availability with a new open-source NKI Standard Library (nki-stdlib), a built-in CPU Simulator, `nki.language` APIs, and several API improvements for correctness and consistency. -* new features: +- new features: NKI Standard Library (nki-stdlib) — open-source, developer-visible code for all NKI APIs and native language objects -* NKI CPU Simulator — `nki.simulate(kernel)` executes NKI kernels on CPU without NeuronDevice hardware (experimental) +- NKI CPU Simulator — `nki.simulate(kernel)` executes NKI kernels on CPU without NeuronDevice hardware (experimental) -* `nki.typing` module — type-annotate kernel tensor parameters with `nt.tensor[shape]` +- `nki.typing` module — type-annotate kernel tensor parameters with `nt.tensor[shape]` -* `nki.language` convenience APIs (experimental) — `nl.load`, `nl.store`, `nl.copy`, `nl.matmul`, `nl.transpose`, `nl.softmax` +- `nki.language` convenience APIs (experimental) — `nl.load`, `nl.store`, `nl.copy`, `nl.matmul`, `nl.transpose`, `nl.softmax` -* new `nki.isa` APIs: +- new `nki.isa` APIs: `nki.isa.exponential` — dedicated exponential instruction (Trn3/NeuronCore-v4 only) -* new `nki.collectives` APIs: +- new `nki.collectives` APIs: `nki.collectives.all_to_all_v` — variable-length all-to-all collective -* changes to existing APIs: +- changes to existing APIs: `nki.isa.nc_matmul` and `nki.isa.nc_matmul_mx` — new `accumulate` parameter for controlling overwrite vs accumulation on PSUM -* `nki.language.ndarray` — new `address` parameter for explicit memory placement +- `nki.language.ndarray` — new `address` parameter for explicit memory placement -* `nki.isa.dma_copy` — no longer supports reading directly from PSUM; `dst_rmw_op` and `unique_indices` parameters removed (use `nisa.dma_compute` instead); enforces type matching with `dge_mode=hwdge` +- `nki.isa.dma_copy` — no longer supports reading directly from PSUM; `dst_rmw_op` and `unique_indices` parameters removed (use `nisa.dma_compute` instead); enforces type matching with `dge_mode=hwdge` -* `nki.isa.dma_compute` — `scales` and `reduce_op` parameter positions swapped; `unique_indices` parameter added +- `nki.isa.dma_compute` — `scales` and `reduce_op` parameter positions swapped; `unique_indices` parameter added -* `nki.isa.memset` — `value` must match destination dtype; x4 packed types enforce `value=0` +- `nki.isa.memset` — `value` must match destination dtype; x4 packed types enforce `value=0` -* `nki.isa.tensor_reduce` — fixed incorrect axis handling for 3D/4D tensors +- `nki.isa.tensor_reduce` — fixed incorrect axis handling for 3D/4D tensors -* `nki.isa.sendrecv` — `use_gpsimd_dma` replaced by `dma_engine` enum +- `nki.isa.sendrecv` — `use_gpsimd_dma` replaced by `dma_engine` enum -* `nki.isa.affine_select` — `offset` parameter moved to keyword argument +- `nki.isa.affine_select` — `offset` parameter moved to keyword argument -* `nki.isa.register_move` — `imm` parameter renamed to `src`, now accepts `VirtualRegister` +- `nki.isa.register_move` — `imm` parameter renamed to `src`, now accepts `VirtualRegister` -* `nki.jit` — `platform_target` parameter removed (use `NEURON_PLATFORM_TARGET_OVERRIDE` env var); `mode` parameter deprecated and ignored +- `nki.jit` — `platform_target` parameter removed (use `NEURON_PLATFORM_TARGET_OVERRIDE` env var); `mode` parameter deprecated and ignored -* Output tensors must use `buffer=nl.shared_hbm` +- Output tensors must use `buffer=nl.shared_hbm` -* Integer enum constants no longer supported (use named enum members) +- Integer enum constants no longer supported (use named enum members) -* String buffer names no longer supported (use buffer objects like `nl.sbuf`, `nl.psum`) +- String buffer names no longer supported (use buffer objects like `nl.sbuf`, `nl.psum`) -* `nki.isa.tensor_copy_dynamic_src` / `nki.isa.tensor_copy_dynamic_dst` deprecated (use `nisa.tensor_copy()` with `.ap()` and `scalar_offset`) +- `nki.isa.tensor_copy_dynamic_src` / `nki.isa.tensor_copy_dynamic_dst` deprecated (use `nisa.tensor_copy()` with `.ap()` and `scalar_offset`) -* default value changes: +- default value changes: `nki.isa.iota` — `offset` now optional with default `0` -* `nki.isa.core_barrier` — `engine` default changed from `unknown` to `gpsimd` +- `nki.isa.core_barrier` — `engine` default changed from `unknown` to `gpsimd` -* `nki.language.num_programs` — `axes` default changed from `None` to `0` +- `nki.language.num_programs` — `axes` default changed from `None` to `0` -* `nki.language.program_id` — `axis` now has default value of `0` +- `nki.language.program_id` — `axis` now has default value of `0` -* `nki.language.ndarray` — `buffer` default changed from `None` to `nl.sbuf` +- `nki.language.ndarray` — `buffer` default changed from `None` to `nl.sbuf` -* `nki.language.zeros` — `buffer` default changed from `None` to `nl.sbuf` +- `nki.language.zeros` — `buffer` default changed from `None` to `nl.sbuf` -* `nki.language.sequential_range` — `stop` and `step` now have default values (`None` and `1`) +- `nki.language.sequential_range` — `stop` and `step` now have default values (`None` and `1`) -* language restrictions: +- language restrictions: Keyword-only argument separator (`*`) not supported in kernel function signatures -* `is` / `is not` operators not supported; use `==` / `!=` instead +- `is` / `is not` operators not supported; use `==` / `!=` instead -* `list` not supported as kernel argument type; use tuples instead +- `list` not supported as kernel argument type; use tuples instead -* Collectives — `num_channels` removed from `collective_permute_implicit_current_processing_rank_id` +- Collectives — `num_channels` removed from `collective_permute_implicit_current_processing_rank_id` ## Neuron Kernel Interface (NKI) (Beta) [2.27] Date: 12/25/2025 -* new `nki.language` APIs: +- new `nki.language` APIs: `nki.language.device_print` -* new `nki.isa` APIs: +- new `nki.isa` APIs: `nki.isa.dma_compute` -* `nki.isa.nki.isa.quantize_mx` +- `nki.isa.nki.isa.quantize_mx` -* `nki.isa.nki.isa.nc_matmul` +- `nki.isa.nki.isa.nc_matmul` -* `nki.isa.nki.isa.nc_n_gather` [used to be `nl.gather_flattened` with free partition limited to 512] +- `nki.isa.nki.isa.nc_n_gather` [used to be `nl.gather_flattened` with free partition limited to 512] -* `nki.isa.rand2` +- `nki.isa.rand2` -* `nki.isa.rand_set_state` +- `nki.isa.rand_set_state` -* `nki.isa.rand_get_state` +- `nki.isa.rand_get_state` -* `nki.isa.set_rng_seed` +- `nki.isa.set_rng_seed` -* `nki.isa.rng` +- `nki.isa.rng` -* new `dtypes`: +- new `dtypes`: `nki.language.float8_e5m2_x4` -* `nki.language.float4_e2m1fn_x4` +- `nki.language.float4_e2m1fn_x4` -* `nki.language.float8_e4m3fn_x4` +- `nki.language.float8_e4m3fn_x4` -* changes to existing APIs: +- changes to existing APIs: several `nki.language` APIs have been removed in NKI Beta 2 -* all nki.isa APIs have `dst` as an input param +- all nki.isa APIs have `dst` as an input param -* all nki.isa APIs removed `dtype` and `mask` support +- all nki.isa APIs removed `dtype` and `mask` support -* `nki.isa.memset` — removed `shape` positional arg , since we have `dst` +- `nki.isa.memset` — removed `shape` positional arg , since we have `dst` -* `nki.isa.affine_select` — instead of `pred`, we now take `pattern` and `cmp_op` params +- `nki.isa.affine_select` — instead of `pred`, we now take `pattern` and `cmp_op` params -* `nki.isa.iota` — `expr` replaced with `pattern` and `offset` +- `nki.isa.iota` — `expr` replaced with `pattern` and `offset` -* `nki.isa.nc_stream_shuffle` - `src` and `dst` order changed +- `nki.isa.nc_stream_shuffle` - `src` and `dst` order changed -* docs improvements: +- docs improvements: restructured NKI Documentation to align with workflows -* added [Trainium3 Architecture Guide for NKI](../architecture/trainium3_arch.md) +- added [Trainium3 Architecture Guide for NKI](../architecture/trainium3_arch.md) -* added [About Neuron Kernel Interface (NKI)](../programming/api/index.md) +- added [About Neuron Kernel Interface (NKI)](../programming/api/index.md) -* added [NKI Environment Setup Guide](../programming/setup-env.md) +- added [NKI Environment Setup Guide](../programming/setup-env.md) -* added [Get Started with NKI](../programming/quickstart-implement-run-kernel.md) +- added [Get Started with NKI](../programming/quickstart-implement-run-kernel.md) -* added [NKI Language Guide](../programming/nki-language-guide.md) +- added [NKI Language Guide](../programming/nki-language-guide.md) -* added [About the NKI Compiler](../programming/nki-compiler.md) +- added [About the NKI Compiler](../programming/nki-compiler.md) -* added [About NKI Beta Versions](../optimization/nki-beta-versions.md) +- added [About NKI Beta Versions](../optimization/nki-beta-versions.md) -* added [MXFP Matrix Multiplication with NKI](../optimization/mxfp-matmul.md) +- added [MXFP Matrix Multiplication with NKI](../optimization/mxfp-matmul.md) -* updated [Matrix Multiplication Tutorial](../programming/tutorials/matrix_multiplication.md) +- updated [Matrix Multiplication Tutorial](../programming/tutorials/matrix_multiplication.md) -* updated [Profile a NKI Kernel](../optimization/use-neuron-profile.md) +- updated [Profile a NKI Kernel](../optimization/use-neuron-profile.md) -* updated [NKI APIs](../programming/api/index.md) +- updated [NKI APIs](../programming/api/index.md) -* updated [NKI Library docs](../programming/api/index.md) +- updated [NKI Library docs](../programming/api/index.md) -* removed NKI Error Guide +- removed NKI Error Guide -* known issues: +- known issues: `nki.isa.nki.isa.nc_matmul` - `is_moving_onezero` was incorrectly named `is_moving_zero` in this release -* NKI ISA semantic checks are not available with Beta 2, workaround is to reference the API docs +- NKI ISA semantic checks are not available with Beta 2, workaround is to reference the API docs -* NKI Collectives are not available with Beta 2 +- NKI Collectives are not available with Beta 2 -* `nki.benchmark` and `nki.profile` are not available with Beta 2 +- `nki.benchmark` and `nki.profile` are not available with Beta 2 ## Neuron Kernel Interface (NKI) (Beta) [2.26] Date: 09/18/2025 -* new `nki.language` APIs: +- new `nki.language` APIs: `nki.language.gelu_apprx_sigmoid` - Gaussian Error Linear Unit activation function with sigmoid approximation. -* `nki.language.tile_size.total_available_sbuf_size` to get total available SBUF size +- `nki.language.tile_size.total_available_sbuf_size` to get total available SBUF size -* new `nki.isa` APIs: +- new `nki.isa` APIs: `nki.isa.select_reduce` - selectively copy elements with max reduction -* `nki.isa.sequence_bounds` - compute sequence bounds of segment IDs +- `nki.isa.sequence_bounds` - compute sequence bounds of segment IDs -* `nki.isa.dma_transpose` +- `nki.isa.dma_transpose` `axes` param to define 4D transpose for some supported cases -* `dge_mode` to specify Descriptor Generation Engine (DGE). +- `dge_mode` to specify Descriptor Generation Engine (DGE). -* `nl.gelu_apprx_sigmoid` op support on `nki.isa.activation` +- `nl.gelu_apprx_sigmoid` op support on `nki.isa.activation` -* fixes / improvements: +- fixes / improvements: `nki.language.store` supports PSUM buffer with extra additional copy inserted. -* docs/tutorial improvements: +- docs/tutorial improvements: `nki.isa.dma_transpose` API doc and example -* `nki.simulate_kernel` example improvement +- `nki.simulate_kernel` example improvement -* use `nl.fp32.min` in tutorial code instead of a magic number +- use `nl.fp32.min` in tutorial code instead of a magic number -* better error reporting: +- better error reporting: indirect indexing on transpose -* mask expressions +- mask expressions ## Neuron Kernel Interface (NKI) (Beta) [2.24] Date: 06/24/2025 -* `sqrt` valid data range extended for accuracy improvement with wider numerical values support. +- `sqrt` valid data range extended for accuracy improvement with wider numerical values support. -* `nki.language.gather_flattened` new API +- `nki.language.gather_flattened` new API -* `nki.isa.nc_match_replace8` additional param `dst_idx` +- `nki.isa.nc_match_replace8` additional param `dst_idx` -* improved docs/examples on `nki.isa.nc_match_replace8`, `nki.isa.nc_stream_shuffle` +- improved docs/examples on `nki.isa.nc_match_replace8`, `nki.isa.nc_stream_shuffle` -* improved error messages +- improved error messages ## Neuron Kernel Interface (NKI) (Beta) [2.23] Date: 05/20/2025 -* `nki.isa.range_select` (for trn2) new instruction +- `nki.isa.range_select` (for trn2) new instruction -* `abs`, `power` ops supported on to nki.isa tensor instruction +- `abs`, `power` ops supported on to nki.isa tensor instruction -* `abs` op supported on `nki.isa.activation` instruction +- `abs` op supported on `nki.isa.activation` instruction -* GpSIMD engine support added to `add`, `multiply` in 32bit integer to nki.isa tensor operations +- GpSIMD engine support added to `add`, `multiply` in 32bit integer to nki.isa tensor operations -* `nki.isa.tensor_copy_predicated` support for reversing predicate. +- `nki.isa.tensor_copy_predicated` support for reversing predicate. -* `nki.isa.tensor_copy_dynamic_src`, `tensor_copy_dynamic_dst` engine selection. +- `nki.isa.tensor_copy_dynamic_src`, `tensor_copy_dynamic_dst` engine selection. -* `nki.isa.dma_copy` additional support with `dge_mode`, `oob_mode`, and in-place add `rmw_op`. +- `nki.isa.dma_copy` additional support with `dge_mode`, `oob_mode`, and in-place add `rmw_op`. -* `+=, -=, /=, *=` operators now work consistently across loop types, PSUM, and SBUF, +- `+=, -=, /=, *=` operators now work consistently across loop types, PSUM, and SBUF, -* fixed simulation for instructions: `nki.language.rand`, `random_seed`, `nki.isa.dropout` +- fixed simulation for instructions: `nki.language.rand`, `random_seed`, `nki.isa.dropout` -* fixed simulation masking behavior +- fixed simulation masking behavior -* Added warning when the block dimension is used for SBUF and PSUM tensors, see: [NKI Block Dimension Migration Guide](migration/nki_block_dimension_migration_guide.md#nki-block-dimension-migration-guide) +- Added warning when the block dimension is used for SBUF and PSUM tensors, see: [NKI Block Dimension Migration Guide](migration/nki_block_dimension_migration_guide.md#nki-block-dimension-migration-guide) ## Neuron Kernel Interface (NKI) (Beta) [2.22] Date: 04/03/2025 -* New modules and APIs: +- New modules and APIs: `nki.profile` -* `nki.isa` new APIs: +- `nki.isa` new APIs: `tensor_copy_dynamic_dst` -* `tensor_copy_predicated` +- `tensor_copy_predicated` -* `max8`, `nc_find_index8`, `nc_match_replace8` +- `max8`, `nc_find_index8`, `nc_match_replace8` -* `nc_stream_shuffle` +- `nc_stream_shuffle` -* `nki.language` new APIs: `mod`, `fmod`, `reciprocal`, `broadcast_to`, `empty_like` +- `nki.language` new APIs: `mod`, `fmod`, `reciprocal`, `broadcast_to`, `empty_like` -* Improvements: +- Improvements: `nki.isa.nc_matmul` now supports PE tiling feature -* `nki.isa.activation` updated to support reduce operation and `reduce` commands +- `nki.isa.activation` updated to support reduce operation and `reduce` commands -* `nki.isa.engine` enum +- `nki.isa.engine` enum -* `engine` parameter added to more `nki.isa` APIs that support engine selection (ie, `tensor_scalar`, `tensor_tensor`, `memset`) +- `engine` parameter added to more `nki.isa` APIs that support engine selection (ie, `tensor_scalar`, `tensor_tensor`, `memset`) -* Documentation for `nki.kernels` have been moved to the GitHub: [https://aws-neuron.github.io/nki-samples](https://aws-neuron.github.io/nki-samples). -The source code can be viewed at [aws-neuron/nki-samples](https://github.com/aws-neuron/nki-samples). +- Documentation for `nki.kernels` have been moved to the GitHub: [https://aws-neuron.github.io/nki-samples](https://aws-neuron.github.io/nki-samples). + The source code can be viewed at [aws-neuron/nki-samples](https://github.com/aws-neuron/nki-samples). These kernels are still shipped as part of Neuron package in `neuronxcc.nki.kernels` module -* Documentation updates: +- Documentation updates: Kernels public repository [https://aws-neuron.github.io/nki-samples](https://aws-neuron.github.io/nki-samples) -* Updated [profiling guide](../optimization/use-neuron-profile.md) to use `nki.profile` instead of `nki.benchmark` +- Updated [profiling guide](../optimization/use-neuron-profile.md) to use `nki.profile` instead of `nki.benchmark` -* NKI ISA Activation functions table now have [valid input data ranges](../programming/api/nki.api.shared.md#tbl-act-func) listed +- NKI ISA Activation functions table now have [valid input data ranges](../programming/api/nki.api.shared.md#tbl-act-func) listed -* NKI ISA Supported Math operators now have [supported engine](../programming/api/nki.api.shared.md#tbl-aluop) listed +- NKI ISA Supported Math operators now have [supported engine](../programming/api/nki.api.shared.md#tbl-aluop) listed -* Clarify `+=` syntax support/limitation +- Clarify `+=` syntax support/limitation ## Neuron Kernel Interface (NKI) (Beta) [2.21] Date: 12/16/2024 -* New modules and APIs: +- New modules and APIs: `nki.compiler` module with Allocation Control and Kernel decorators, see guide for more info. -* `nki.isa`: new APIs (`activation_reduce`, `tensor_partition_reduce`, -`scalar_tensor_tensor`, `tensor_scalar_reduce`, `tensor_copy`, -`tensor_copy_dynamic_src`, `dma_copy`), new activation functions(`identity`, -`silu`, `silu_dx`), and target query APIs (`nc_version`, `get_nc_version`). +- `nki.isa`: new APIs (`activation_reduce`, `tensor_partition_reduce`, + `scalar_tensor_tensor`, `tensor_scalar_reduce`, `tensor_copy`, + `tensor_copy_dynamic_src`, `dma_copy`), new activation functions(`identity`, + `silu`, `silu_dx`), and target query APIs (`nc_version`, `get_nc_version`). -* `nki.language`: new APIs (`shared_identity_matrix`, `tan`, -`silu`, `silu_dx`, `left_shift`, `right_shift`, `ds`, `spmd_dim`, `nc`). +- `nki.language`: new APIs (`shared_identity_matrix`, `tan`, + `silu`, `silu_dx`, `left_shift`, `right_shift`, `ds`, `spmd_dim`, `nc`). -* New `datatype `: `float8_e5m2` +- New `datatype `: `float8_e5m2` -* New `kernels` (`allocated_fused_self_attn_for_SD_small_head_size`, -`allocated_fused_rms_norm_qkv`) added, kernels moved to public repository. +- New `kernels` (`allocated_fused_self_attn_for_SD_small_head_size`, + `allocated_fused_rms_norm_qkv`) added, kernels moved to public repository. -* Improvements: +- Improvements: Semantic analysis checks for nki.isa APIs to validate supported ops, dtypes, and tile shapes. -* Standardized naming conventions with keyword arguments for common optional parameters. +- Standardized naming conventions with keyword arguments for common optional parameters. -* Transition from function calls to kernel decorators (`jit`, -`benchmark`, `baremetal`, `simulate_kernel`). +- Transition from function calls to kernel decorators (`jit`, + `benchmark`, `baremetal`, `simulate_kernel`). -* Documentation updates: +- Documentation updates: Tutorial for [SPMD usage with multiple Neuron Cores on Trn2](../programming/tutorials/spmd_multiple_nc_tensor_addition.md) @@ -342,25 +342,25 @@ Tutorial for [SPMD usage with multiple Neuron Cores on Trn2](../programming/tuto Date: 12/03/2024 -* NKI support for Trainium2, including full integration with Neuron Compiler. -Users can directly shard NKI kernels across multiple Neuron Cores from an SPMD launch grid. -See [tutorial](../programming/tutorials/spmd_multiple_nc_tensor_addition.md) for more info. -See [Trainium2 Architecture Guide](../architecture/trainium2_arch.md) for an initial version of the architecture specification -(more details to come in future releases). +- NKI support for Trainium2, including full integration with Neuron Compiler. + Users can directly shard NKI kernels across multiple Neuron Cores from an SPMD launch grid. + See [tutorial](../programming/tutorials/spmd_multiple_nc_tensor_addition.md) for more info. + See [Trainium2 Architecture Guide](../architecture/trainium2_arch.md) for an initial version of the architecture specification + (more details to come in future releases). -* New calling convention in NKI kernels, where kernel output tensors are explicitly returned from the kernel instead -of pass-by-reference. See any [NKI tutorial](../programming/api/index.md) for code examples. +- New calling convention in NKI kernels, where kernel output tensors are explicitly returned from the kernel instead + of pass-by-reference. See any [NKI tutorial](../programming/api/index.md) for code examples. ## Neuron Kernel Interface (NKI) (Beta) [2.20] Date: 09/16/2024 -* This release includes the beta launch of the Neuron Kernel Interface (NKI) (Beta). -NKI is a programming interface enabling developers to build optimized compute kernels -on top of Trainium and Inferentia. NKI empowers developers to enhance deep learning models -with new capabilities, performance optimizations, and scientific innovation. -It natively integrates with PyTorch and JAX, providing a Python-based programming environment -with Triton-like syntax and tile-level semantics offering a familiar programming experience -for developers. Additionally, to enable bare-metal access precisely programming the instructions -used by the chip, this release includes a set of NKI APIs (`nki.isa`) that directly emit -Neuron Instruction Set Architecture (ISA) instructions in NKI kernels. \ No newline at end of file +- This release includes the beta launch of the Neuron Kernel Interface (NKI) (Beta). + NKI is a programming interface enabling developers to build optimized compute kernels + on top of Trainium and Inferentia. NKI empowers developers to enhance deep learning models + with new capabilities, performance optimizations, and scientific innovation. + It natively integrates with PyTorch and JAX, providing a Python-based programming environment + with Triton-like syntax and tile-level semantics offering a familiar programming experience + for developers. Additionally, to enable bare-metal access precisely programming the instructions + used by the chip, this release includes a set of NKI APIs (`nki.isa`) that directly emit + Neuron Instruction Set Architecture (ISA) instructions in NKI kernels. diff --git a/skills/neuron-nki-profile-querying/SKILL.md b/skills/neuron-nki-profile-querying/SKILL.md index f8b0107..33a9374 100644 --- a/skills/neuron-nki-profile-querying/SKILL.md +++ b/skills/neuron-nki-profile-querying/SKILL.md @@ -28,7 +28,7 @@ on localhost. No deployment, no remote service — just the CLI and curl. For more advanced analysis, use python on parquet to compute performance bounds and investigate precise inefficiencies within arbitrary execution -intervals. +intervals. **What you need:** A compiled NEFF file and a captured NTFF trace file. These come from `/neuron-nki-profiling` or from running a kernel with the right @@ -62,9 +62,11 @@ That's it. Ingest, serve, query. - NEFF file (compiled kernel binary) + NTFF file (execution trace) Check availability: + ```bash which neuron-explorer && neuron-explorer --version ``` + If not found, check `/opt/aws/neuron/bin/neuron-explorer`. --- @@ -81,6 +83,7 @@ If not found, check `/opt/aws/neuron/bin/neuron-explorer`. > DmaPacketAggregated) may be empty and source-level attribution will be missing. Check whether the profile has the data you need: + ```bash # After ingesting (Step 2), check for DMA packet data curl -s -X POST http://localhost:3002/api/v1/db/${PROFILE_NAME}/_search \ @@ -104,10 +107,11 @@ os.environ["NEURON_RT_VISIBLE_CORES"] = ... # Restrict available cores when runn os.environ["NEURON_RT_INSPECT_ENABLE"] = "1" os.environ["NEURON_RT_INSPECT_DEVICE_PROFILE"] = "1" os.environ["NEURON_RT_INSPECT_SYSTEM_PROFILE"] = "0" -os.environ["NEURON_RT_INSPECT_OUTPUT_DIR"] = ... # This is for the NEFF generation if needed. NTFF will go to the -s capture path in the next command. +os.environ["NEURON_RT_INSPECT_OUTPUT_DIR"] = ... # This is for the NEFF generation if needed. NTFF will go to the -s capture path in the next command. ``` Then re-capture with DGE notifications enabled: + ```bash NEFF_PATH=$(find ./output -name "*.neff" | head -1) NEURON_RT_ENABLE_DGE_NOTIFICATIONS=1 neuron-explorer capture \ @@ -119,25 +123,26 @@ NEURON_RT_ENABLE_DGE_NOTIFICATIONS=1 neuron-explorer capture \ With `--profile-nth-exec=2`, the output file is `profile_exec_2.ntff` (not `profile.ntff`), written to the directory specified by the `-s` flag. -| Env Var | What it enables | -|---------|----------------| -| `XLA_IR_DEBUG` / `XLA_HLO_DEBUG` | HLO-level debug info in NEFF | -| `NEURON_FRAMEWORK_DEBUG` | Framework-level source attribution | +| Env Var | What it enables | +| ------------------------------------ | -------------------------------------------------- | +| `XLA_IR_DEBUG` / `XLA_HLO_DEBUG` | HLO-level debug info in NEFF | +| `NEURON_FRAMEWORK_DEBUG` | Framework-level source attribution | | `NEURON_RT_ENABLE_DGE_NOTIFICATIONS` | DMA packet tables (DmaPacket, DmaPacketAggregated) | -| `NEURON_RT_INSPECT_DEVICE_PROFILE` | Device-level profiling in NEFF output | +| `NEURON_RT_INSPECT_DEVICE_PROFILE` | Device-level profiling in NEFF output | If the existing profile has the data you need, skip this step entirely. Another thing to look out for is running torch functions on device like randomnly generating -inputs. This will be fused into the kernel execution and obfuscate it's profile. Move those -commands off device if you want to isolate kernel execution. +inputs. This will be fused into the kernel execution and obfuscate it's profile. Move those +commands off device if you want to isolate kernel execution. -### Step 1: Ingest and Start Server +### Step 1: Ingest and Start Server If you want to run SQL queries against the Neuron Explorer DuckDB engine, use the view -command with --disable-ui to start the server. +command with --disable-ui to start the server. Set variables: + ```bash NEFF_PATH= NTFF_PATH= @@ -145,12 +150,14 @@ PROFILE_NAME= NE_DATA_PATH=~/.local/share/neuron-profile ``` -Check if the neuron-explorer server is already running: +Check if the neuron-explorer server is already running: + ```bash curl -s http://localhost:3002/api/v1/health ``` -If the server is already running or if you are running python directly -on the parquet, use --ingest-only in the following command instead of --disable-ui. + +If the server is already running or if you are running python directly +on the parquet, use --ingest-only in the following command instead of --disable-ui. ```bash neuron-explorer view \ @@ -163,11 +170,13 @@ neuron-explorer view \ NE_PID=$! echo "neuron-explorer started (PID: $NE_PID), waiting for API..." ``` -The command may fail on an conflicting port from the existing server but the ingestion + +The command may fail on an conflicting port from the existing server but the ingestion may have still succeeded. If so, check for `Processing for ... is complete` before the error message or rerun with --ingest-only. Wait for API: + ```bash for i in $(seq 1 60); do if curl -s http://localhost:3002/api/v1/health 2>/dev/null | grep -q healthy; then @@ -192,17 +201,19 @@ curl -s -X POST http://localhost:3002/api/v1/db/${PROFILE_NAME}/_search \ ``` List all available tables: + ```bash curl -s -X POST http://localhost:3002/api/v1/db/${PROFILE_NAME}/_search \ -H 'Content-Type: application/json' \ -d '{"type": "listDbFiles"}' | python3 -m json.tool ``` -### Step 3a: Execute SQL Queries +### Step 3a: Execute SQL Queries Use `databaseExplorerQuery` for arbitrary SQL (SELECT only). **Summary metrics — which engine is the bottleneck?** + ```bash curl -s -X POST http://localhost:3002/api/v1/db/${PROFILE_NAME}/_search \ -H 'Content-Type: application/json' \ @@ -210,6 +221,7 @@ curl -s -X POST http://localhost:3002/api/v1/db/${PROFILE_NAME}/_search \ ``` **Instruction breakdown — what is each engine doing and waiting on?** + ```bash curl -s -X POST http://localhost:3002/api/v1/db/${PROFILE_NAME}/_search \ -H 'Content-Type: application/json' \ @@ -217,6 +229,7 @@ curl -s -X POST http://localhost:3002/api/v1/db/${PROFILE_NAME}/_search \ ``` **NKI source line hotspots — which lines of the kernel are slowest?** + ```bash curl -s -X POST http://localhost:3002/api/v1/db/${PROFILE_NAME}/_search \ -H 'Content-Type: application/json' \ @@ -309,8 +322,7 @@ it tells you what happened, not always why. or `bir_debug_info_source_location` is mostly NULL, the query results are incomplete — re-profile before interpreting. - -### Step 5: Cleanup +### Step 5: Cleanup ```bash kill $NE_PID 2>/dev/null @@ -326,7 +338,7 @@ selection — lives in [performance-bounds.md](references/performance-bounds.md) Follow the **"The bounds"** section of performance-bounds.md to compute all three families (memory, compute, pipeline). These require Python on parquet -(Step 3c). +(Step 3c). ### 2. Identify the dominant gaps @@ -346,34 +358,35 @@ typically has multiple active inefficiencies. Present a single summary: -- **Bounds table**: all bounds with values and the gap between each pair. -Also report each engine's total time pointing out the largest one(s) as -the bottleneck(s). If neither DMA nor Tensor Engine is the bottleneck, -explain which engine is the bottleneck and that supporting it is still WIP. +- **Bounds table**: all bounds with values and the gap between each pair. + Also report each engine's total time pointing out the largest one(s) as + the bottleneck(s). If neither DMA nor Tensor Engine is the bottleneck, + explain which engine is the bottleneck and that supporting it is still WIP. - **Per-investigation findings**: gap size, source lines responsible, and their contributions. Include investigations that found nothing so the analysis is visibly complete. Order the presented inefficiencies and investigation findings according - to it's relevance to the bottlenecks and the measured gaps. +to it's relevance to the bottlenecks and the measured gaps. ### 5. Follow up (After an optimization step/attempt) -After an optimization step or attempt, investigate the new profile to +After an optimization step or attempt, investigate the new profile to identify exactly what improved or regressed. Follow the full process and -present a side by side report of all of the bounds and engine times as well -as the new investigation findings. Highlight changes but do not over-interpret, +present a side by side report of all of the bounds and engine times as well +as the new investigation findings. Highlight changes but do not over-interpret, only relay what the evidence shows. Static code analysis is faulty, you will be -tempted to over-intepret the causes and effects, DON'T (unless EXPLICITELY) asked -to. +tempted to over-intepret the causes and effects, DON'T (unless EXPLICITELY) asked +to. ### Worked Examples For end-to-end examples of profile-guided optimization, see: -| Investigation | What it covers | -|--------------|----------------| +| Investigation | What it covers | +| ---------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | | [Optimizing-Matmul](references/example-bounds-analysis.md) | End-to-end bounds analysis of a 4096x4096 bf16 matmul across three versions: V0 (naive tiling, DMA-bound), V1 (free-dimension blocking, reduces reloads, flips bottleneck to TE), V2 (row loads, near-peak TE utilization). Shows bounds tables, gap analysis, and investigation results at each step. | + --- ## Multi-Kernel Querying @@ -429,10 +442,10 @@ lsof -i :3002 | head -5 ## Related Skills -| Skill | Purpose | -|-------|---------| -| `/neuron-nki-profiling` | Capture NEFF/NTFF on hardware | +| Skill | Purpose | +| --------------------------------- | ---------------------------------------------------------------------------------------- | +| `/neuron-nki-profiling` | Capture NEFF/NTFF on hardware | | `/neuron-explorer-profile-schema` | Reference for the parquet schema (table modalities, field origins, version-matched YAML) | -| `/neuron-nki-writing` | Write NKI kernels | -| `/neuron-nki-debugging` | Debug compilation errors | -| `/neuron-nki-docs` | Look up API documentation | \ No newline at end of file +| `/neuron-nki-writing` | Write NKI kernels | +| `/neuron-nki-debugging` | Debug compilation errors | +| `/neuron-nki-docs` | Look up API documentation | diff --git a/skills/neuron-nki-profile-querying/references/example-bounds-analysis.md b/skills/neuron-nki-profile-querying/references/example-bounds-analysis.md index 896abd7..de47785 100644 --- a/skills/neuron-nki-profile-querying/references/example-bounds-analysis.md +++ b/skills/neuron-nki-profile-querying/references/example-bounds-analysis.md @@ -1,6 +1,6 @@ # Profile Analysis: blocking-tiled matmul -## Kernel (V0) +## Kernel (V0) A 4096x4096x4096 bf16 matmul: `C[M,N] = A^T[K,M] @ B[K,N]`. Input `lhsT` is pre-transposed to [K, M]. Tiles into 128x128 stationary and 128x512 moving tiles. @@ -43,68 +43,69 @@ Hardware: trn2, dma_ddr_bandwidth = 435 GB/s, TE peak = 78.6 TFLOPS. ### Bounds -| Bound | Value (us) | -|-------|-----------| -| total_time | 8,602 | -| memory_bound | 7,906 | -| memory_bound_ideal | 2,776 | -| memory_bound_no_reloads | 231 | -| compute_bound | 4,806 | -| compute_bound_ideal | 1,748 | -| compute_bound_ideal_useful | 1,748 | -| perfect_pipeline (DMA) | 7,906 | +| Bound | Value (us) | +| -------------------------- | ---------- | +| total_time | 8,602 | +| memory_bound | 7,906 | +| memory_bound_ideal | 2,776 | +| memory_bound_no_reloads | 231 | +| compute_bound | 4,806 | +| compute_bound_ideal | 1,748 | +| compute_bound_ideal_useful | 1,748 | +| perfect_pipeline (DMA) | 7,906 | ### Engine active times -| Engine | Active time | -|--------|------------| -| DMA | 7,906 us (91.9%) | +| Engine | Active time | +| ------ | ---------------- | +| DMA | 7,906 us (91.9%) | | Tensor | 4,806 us (55.9%) | -| Vector | 176 us (2.0%) | -| Scalar | 0 us | -| GpSimd | 0 us | +| Vector | 176 us (2.0%) | +| Scalar | 0 us | +| GpSimd | 0 us | Bottleneck: DMA ### Memory family gaps -| Gap | Value (us) | % of total | -|-----|-----------|------------| -| DMA idle | 696 | 8.1% | -| DMA inefficiency | 5,130 | 59.6% | -| Excess traffic | 2,545 | 29.6% (91.7% of traffic) | +| Gap | Value (us) | % of total | +| ---------------- | ---------- | ------------------------ | +| DMA idle | 696 | 8.1% | +| DMA inefficiency | 5,130 | 59.6% | +| Excess traffic | 2,545 | 29.6% (91.7% of traffic) | ### Compute family gaps -| Gap | Value (us) | % of total | -|-----|-----------|------------| -| TE idle | 3,796 | 44.1% | -| TE underutil | 3,058 | 35.6% | -| Transpose | 0 | 0.0% (0.0% of Flops)| +| Gap | Value (us) | % of total | +| ------------ | ---------- | -------------------- | +| TE idle | 3,796 | 44.1% | +| TE underutil | 3,058 | 35.6% | +| Transpose | 0 | 0.0% (0.0% of Flops) | ### Summary -Although TE instructions seem to be inefficient, DMA is the clear bottleneck -and transfer inefficiency + redundant transfers contribute (59.6% + 29.6% = 89%) of total -kernel execution time. Eliminating this would lead to the next unaffected engine (TE) -to become the bottleneck at 4,806ns (1.79x speedup). +Although TE instructions seem to be inefficient, DMA is the clear bottleneck +and transfer inefficiency + redundant transfers contribute (59.6% + 29.6% = 89%) of total +kernel execution time. Eliminating this would lead to the next unaffected engine (TE) +to become the bottleneck at 4,806ns (1.79x speedup). -Investigations (Appendix V0): -- Redundant dma transfers: Using the investigations/redundant_dma_transfers.md steps -we find that all excess dma transfers are input reloads. Almost entirely -dma transfers are input reloads. Almost entirely from reloading the rhs tensor. -- DMA efficiency: Using the investigations/dma_efficiency.md steps, we find that loads and -stores are also well below ideal transfer size. +Investigations (Appendix V0): + +- Redundant dma transfers: Using the investigations/redundant_dma_transfers.md steps + we find that all excess dma transfers are input reloads. Almost entirely + dma transfers are input reloads. Almost entirely from reloading the rhs tensor. +- DMA efficiency: Using the investigations/dma_efficiency.md steps, we find that loads and + stores are also well below ideal transfer size. DMA is the bottleneck and it's execution is dominated by redundant input reloads, this is the gap we will prioritize for V1 even if we keep in mind that transfer sizes should be increased in the future -as well. +as well. ## V1: Blocking free dimension -V0 has three nested loops where both operands are (attempted to be) loaded on every +V0 has three nested loops where both operands are (attempted to be) loaded on every inner iteration. To reduce the excessive input reloads, in V1, we block the loads over -the M and N dimension to localize computation. +the M and N dimension to localize computation. Key structural change from V0: @@ -139,66 +140,67 @@ lhsT is loaded once per M-block and reused across all 8 N-tiles. ### V1 Bounds (BM=8, BN=2) -| Bound | V0 (us) | V1 (us) | V0 -> V1 | -|-------|---------|---------|--------| -| total_time | 8,602 | 2,451 | -6,151 us (3.5x) | -| memory_bound | 7,906 | 1,325 | -6,581 us (6.0x) | -| memory_bound_ideal | 2,776 | 463 | -2,313 us (6.0x) | -| memory_bound_no_reloads | 231 | 231 | — | -| compute_bound | 4,806 | 2,163 | -2,643 us (2.2x) | -| compute_bound_ideal | 1,748 | 1,748 | — | -| compute_bound_ideal_useful | 1,748 | 1,748 | — | -| perfect_pipeline | 7,906 (DMA) | 2,163 (Tensor) | -5,743 us (3.7x) | +| Bound | V0 (us) | V1 (us) | V0 -> V1 | +| -------------------------- | ----------- | -------------- | ---------------- | +| total_time | 8,602 | 2,451 | -6,151 us (3.5x) | +| memory_bound | 7,906 | 1,325 | -6,581 us (6.0x) | +| memory_bound_ideal | 2,776 | 463 | -2,313 us (6.0x) | +| memory_bound_no_reloads | 231 | 231 | — | +| compute_bound | 4,806 | 2,163 | -2,643 us (2.2x) | +| compute_bound_ideal | 1,748 | 1,748 | — | +| compute_bound_ideal_useful | 1,748 | 1,748 | — | +| perfect_pipeline | 7,906 (DMA) | 2,163 (Tensor) | -5,743 us (3.7x) | ### Engine active times -| Engine | V0 | V1 | V0 -> V1 | -|--------|----|----|----------| -| DMA | 7,906 us (91.9%) | 1,325 us (54.1%) | -6,581 us (6.0x) | +| Engine | V0 | V1 | V0 -> V1 | +| ------ | ---------------- | ---------------- | ---------------- | +| DMA | 7,906 us (91.9%) | 1,325 us (54.1%) | -6,581 us (6.0x) | | Tensor | 4,806 us (55.9%) | 2,163 us (88.2%) | -2,643 us (2.2x) | -| Vector | 176 us (2.0%) | 174 us (7.1%) | — | -| Scalar | 0 us | 0 us | — | -| GpSimd | 0 us | 0 us | — | +| Vector | 176 us (2.0%) | 174 us (7.1%) | — | +| Scalar | 0 us | 0 us | — | +| GpSimd | 0 us | 0 us | — | Bottleneck: DMA → Tensor ### Memory family gaps -| Gap | V0 | V1 | V0 -> V1 | -|-----|----|----|----------| -| DMA idle | 696 us (8.1%) | 1,126 us (46.0%) | +430 us | -| DMA inefficiency | 5,130 us (59.6%) | 862 us (35.2%) | -4,268 us (6.0x) | -| Excess traffic | 2,545 us (91.7% of traffic) | 231 us (50.0% of traffic) | -2,314 us (11.0x) | +| Gap | V0 | V1 | V0 -> V1 | +| ---------------- | --------------------------- | ------------------------- | ----------------- | +| DMA idle | 696 us (8.1%) | 1,126 us (46.0%) | +430 us | +| DMA inefficiency | 5,130 us (59.6%) | 862 us (35.2%) | -4,268 us (6.0x) | +| Excess traffic | 2,545 us (91.7% of traffic) | 231 us (50.0% of traffic) | -2,314 us (11.0x) | ### Compute family gaps -| Gap | V0 | V1 | V0 -> V1 | -|-----|----|----|----------| -| TE idle | 3,796 us (44.1%) | 288 us (11.8%) | -3,508 us (13.2x) | -| TE underutil | 3,058 us (35.6%) | 416 us (17.0%) | -2,642 us (7.4x) | -| Transpose | 0 (0.0% of flops) | 0 (0.0% of flops) | — | +| Gap | V0 | V1 | V0 -> V1 | +| ------------ | ----------------- | ----------------- | ----------------- | +| TE idle | 3,796 us (44.1%) | 288 us (11.8%) | -3,508 us (13.2x) | +| TE underutil | 3,058 us (35.6%) | 416 us (17.0%) | -2,642 us (7.4x) | +| Transpose | 0 (0.0% of flops) | 0 (0.0% of flops) | — | ### Summary V0 -> V1: Reducing input reloads reduced both the redundant transfers and dma efficiency gap as expected. As DMA stopped being the bottleneck, DMA idle gaps also increased -as expected while TE idle gaps reduced as it became the new bottleneck. More interestingly, TE engine active time went down which seems to drive a conveniently reduced gap (less throttling). +as expected while TE idle gaps reduced as it became the new bottleneck. More interestingly, TE engine active time went down which seems to drive a conveniently reduced gap (less throttling). V1: Tensor Engine is the bottleneck at 2,163 us. TE idle (11.8%) and TE underutilization (17.0%) together account for 28.8% of total time in the -bottleneck family. These gaps should be the new focus. +bottleneck family. These gaps should be the new focus. Investigations (Appendix V1): -- TE idle gaps: Examining TE idle gaps tells us that the majority -of TE idle spans are caused by waiting on dma transfers. + +- TE idle gaps: Examining TE idle gaps tells us that the majority + of TE idle spans are caused by waiting on dma transfers. - TE inefficiency: Running the investigations/te_inefficiency.md steps tells us that the tile -sizes are the right size and inefficiency seems to come from throttling due to idle gaps. -- DMA efficiency: Running the investigations/dma_efficiency.md steps tells us that all transfers -are well below the ideal threshold causing low dma BW utilization. + sizes are the right size and inefficiency seems to come from throttling due to idle gaps. +- DMA efficiency: Running the investigations/dma_efficiency.md steps tells us that all transfers + are well below the ideal threshold causing low dma BW utilization. -Since TE engine is the bottleneck, and TE inefficiency seems to all stem from idle gaps specifically +Since TE engine is the bottleneck, and TE inefficiency seems to all stem from idle gaps specifically waiting on dma, this is the performance gap we will tackle. We will do this by improving dma efficiency -through larger transfers. We could also look at techniques for pipelining. +through larger transfers. We could also look at techniques for pipelining. ## V2: Row loads @@ -236,44 +238,44 @@ for m in nl.affine_range(M // BLOCK_M): ### V2 Bounds (BM=8, BN=2, row loads) -| Bound | V0 (us) | V1 (us) | V2 (us) | V1 -> V2 | -|-------|---------|---------|---------|----------| -| total_time | 8,602 | 2,451 | 1,801 | -650 us (1.4x) | -| memory_bound | 7,906 | 1,325 | 819 | -506 us (1.6x) | -| memory_bound_ideal | 2,776 | 463 | 463 | — | -| memory_bound_no_reloads | 231 | 231 | 231 | — | -| compute_bound | 4,806 | 2,163 | 1,781 | -382 us (1.2x) | -| compute_bound_ideal | 1,748 | 1,748 | 1,748 | — | -| compute_bound_ideal_useful | 1,748 | 1,748 | 1,748 | — | -| perfect_pipeline | 7,906 (DMA) | 2,163 (Tensor) | 1,781 (Tensor) | -382 us (1.2x) | +| Bound | V0 (us) | V1 (us) | V2 (us) | V1 -> V2 | +| -------------------------- | ----------- | -------------- | -------------- | -------------- | +| total_time | 8,602 | 2,451 | 1,801 | -650 us (1.4x) | +| memory_bound | 7,906 | 1,325 | 819 | -506 us (1.6x) | +| memory_bound_ideal | 2,776 | 463 | 463 | — | +| memory_bound_no_reloads | 231 | 231 | 231 | — | +| compute_bound | 4,806 | 2,163 | 1,781 | -382 us (1.2x) | +| compute_bound_ideal | 1,748 | 1,748 | 1,748 | — | +| compute_bound_ideal_useful | 1,748 | 1,748 | 1,748 | — | +| perfect_pipeline | 7,906 (DMA) | 2,163 (Tensor) | 1,781 (Tensor) | -382 us (1.2x) | ### Engine active times -| Engine | V0 | V1 | V2 | V1 -> V2 | -|--------|----|----|----|----| -| DMA | 7,906 us (91.9%) | 1,325 us (54.1%) | 819 us (45.5%) | -506 us (1.6x) | +| Engine | V0 | V1 | V2 | V1 -> V2 | +| ------ | ---------------- | ---------------- | ---------------- | -------------- | +| DMA | 7,906 us (91.9%) | 1,325 us (54.1%) | 819 us (45.5%) | -506 us (1.6x) | | Tensor | 4,806 us (55.9%) | 2,163 us (88.2%) | 1,781 us (98.9%) | -382 us (1.2x) | -| Vector | 176 us (2.0%) | 174 us (7.1%) | 172 us (9.6%) | — | -| Scalar | 0 us | 0 us | 0 us | — | -| GpSimd | 0 us | 0 us | 0 us | — | +| Vector | 176 us (2.0%) | 174 us (7.1%) | 172 us (9.6%) | — | +| Scalar | 0 us | 0 us | 0 us | — | +| GpSimd | 0 us | 0 us | 0 us | — | Bottleneck: Tensor ### Memory family gaps -| Gap | V0 | V1 | V2 | V1 -> V2 | -|-----|----|----|----|----| -| DMA idle | 696 us (8.1%) | 1,126 us (46.0%) | 982 us (54.5%) | -144 us | -| DMA inefficiency | 5,130 us (59.6%) | 862 us (35.2%) | 356 us (19.8%) | -506 us (2.4x) | -| Excess traffic | 2,545 us (91.7% of traffic) | 231 us (50.0% of traffic) | 231 us (50.0% of traffic) | — | +| Gap | V0 | V1 | V2 | V1 -> V2 | +| ---------------- | --------------------------- | ------------------------- | ------------------------- | -------------- | +| DMA idle | 696 us (8.1%) | 1,126 us (46.0%) | 982 us (54.5%) | -144 us | +| DMA inefficiency | 5,130 us (59.6%) | 862 us (35.2%) | 356 us (19.8%) | -506 us (2.4x) | +| Excess traffic | 2,545 us (91.7% of traffic) | 231 us (50.0% of traffic) | 231 us (50.0% of traffic) | — | ### Compute family gaps -| Gap | V0 | V1 | V2 | V1 -> V2 | -|-----|----|----|----|----| -| TE idle | 3,796 us (44.1%) | 288 us (11.8%) | 20 us (1.1%) | -268 us (14.4x) | -| TE underutil | 3,058 us (35.6%) | 416 us (17.0%) | 33 us (1.8%) | -383 us (12.6x) | -| Transpose | 0 (0.0% of flops) | 0 (0.0% of flops) | 0 (0.0% of flops) | — | +| Gap | V0 | V1 | V2 | V1 -> V2 | +| ------------ | ----------------- | ----------------- | ----------------- | --------------- | +| TE idle | 3,796 us (44.1%) | 288 us (11.8%) | 20 us (1.1%) | -268 us (14.4x) | +| TE underutil | 3,058 us (35.6%) | 416 us (17.0%) | 33 us (1.8%) | -383 us (12.6x) | +| Transpose | 0 (0.0% of flops) | 0 (0.0% of flops) | 0 (0.0% of flops) | — | ### Summary @@ -289,20 +291,20 @@ V2: Tensor Engine is the bottleneck at 1,781 us , operating at near-peak. Remain ### Excess traffic decomposition -| Metric | Value | -|--------|-------| +| Metric | Value | +| ------------------ | ------------------- | | dma_transfer_bytes | 1,207,697,408 bytes | -| necessary_bytes | 100,663,296 bytes | -| excess_bytes | 1,107,034,112 bytes | -| excess_ratio | 12.0x | +| necessary_bytes | 100,663,296 bytes | +| excess_bytes | 1,107,034,112 bytes | +| excess_ratio | 12.0x | ### Per-tensor breakdown -| Tensor | Type | Size (bytes) | Transferred (bytes) | Repeat | Identity | -|--------|------|-------------|---------------------|--------|----------| -| input0 | IN | 33,554,432 | 1,073,741,824 | 32.0x | rhs (128x512 tiles) | -| input1 | IN | 33,554,432 | 100,532,224 | 3.0x | lhsT (128x128 tiles) | -| output0 | OUT | 33,554,432 | 33,423,360 | 1.0x | result | +| Tensor | Type | Size (bytes) | Transferred (bytes) | Repeat | Identity | +| ------- | ---- | ------------ | ------------------- | ------ | -------------------- | +| input0 | IN | 33,554,432 | 1,073,741,824 | 32.0x | rhs (128x512 tiles) | +| input1 | IN | 33,554,432 | 100,532,224 | 3.0x | lhsT (128x128 tiles) | +| output0 | OUT | 33,554,432 | 33,423,360 | 1.0x | result | Identity determined from transfer sizes: input0 transfers are 131,072 bytes (128x512x2 = TILE_K x TILE_N), input1 transfers are 32,768 bytes @@ -310,10 +312,10 @@ Identity determined from transfer sizes: input0 transfers are 131,072 bytes ### Decomposition -| Source | Bytes | % of excess | -|--------|-------|-------------| -| reload_excess | 1,107,165,184 | 100% | -| spill_bytes | 0 | 0% | +| Source | Bytes | % of excess | +| ------------- | ------------- | ----------- | +| reload_excess | 1,107,165,184 | 100% | +| spill_bytes | 0 | 0% | All excess traffic is input reloads. No spills. rhs (input0) is loaded 32x, lhsT (input1) is loaded 3x. @@ -325,31 +327,32 @@ V1 is Tensor-bottlenecked. Relevant groups: 2a (TE idle), 2b (TE underutil). ### Investigation 2b: Compute Tile Sizes (Step 1 & 2) #### -| Metric | Value | -|--------|-------| -| TE peak | 78.6 TFLOPS | + +| Metric | Value | +| -------- | ----------------- | +| TE peak | 78.6 TFLOPS | | Achieved | 63.5 TFLOPS (81%) | -| Source line | Tiles | % of hw_flops | K/128 | M/128 | N/512 | -|-------------|-------|---------------|-------|-------|-------| -| run_blocking_v1.py:70 | 8,192 | 100% | 128/128 | 128/128 | 512/512 | +| Source line | Tiles | % of hw_flops | K/128 | M/128 | N/512 | +| --------------------- | ----- | ------------- | ------- | ------- | ------- | +| run_blocking_v1.py:70 | 8,192 | 100% | 128/128 | 128/128 | 512/512 | All tile dimensions are at maximum. The 19% gap between achieved and peak is not from undersized tiles. ### Investigation 2a: DMA-Compute Pipelining (Step 1) -| Metric | Value | -|--------|-------| -| REGULAR MATMULs | 8,192 | -| total excess initiation interval | 683 us | -| DMA-caused excess | 95 us (14% of excess) | +| Metric | Value | +| -------------------------------- | --------------------- | +| REGULAR MATMULs | 8,192 | +| total excess initiation interval | 683 us | +| DMA-caused excess | 95 us (14% of excess) | | Last-finishing dependency | Count | -|--------------------------|-------| -| Tensor (previous MATMUL) | 7,778 | -| DPA (DMA transfer) | 258 | -| Vector | 155 | +| ------------------------- | ----- | +| Tensor (previous MATMUL) | 7,778 | +| DPA (DMA transfer) | 258 | +| Vector | 155 | 95% of MATMULs have their previous MATMUL as the last-finishing dependency — the pipeline is running normally. Only 258 MATMULs (3%) are DMA-gated. diff --git a/skills/neuron-nki-profile-querying/references/getting-started.md b/skills/neuron-nki-profile-querying/references/getting-started.md index a4d26bd..150740c 100644 --- a/skills/neuron-nki-profile-querying/references/getting-started.md +++ b/skills/neuron-nki-profile-querying/references/getting-started.md @@ -1,7 +1,7 @@ # Getting Started with Profile Analysis This skill lets you query and analyze NKI kernel execution profiles on -Neuron hardware (Trainium/Inferentia). It works with NEFF (compiled kernel), NTFF (execution trace) and corresponding +Neuron hardware (Trainium/Inferentia). It works with NEFF (compiled kernel), NTFF (execution trace) and corresponding parquet files produced by `neuron-explorer`. ## What it can do @@ -34,13 +34,14 @@ parquet files produced by `neuron-explorer`. ``` I have a profiled kernel at `./output/kernel.neff` and `./output/profile.ntff`. -What is the total execution time, and which engine is most active? +What is the total execution time, and which engine is most active? ``` This will ingest the profile, query the Summary table, and report `total_time` and per-engine active time percentages. ### 2. Full bounds analysis + ``` Profile and analyze the NKI kernel at `kernels/matmul_blocked.py`. The kernel takes lhsT[4096,4096] bf16 and rhs[4096,4096] bf16. @@ -49,6 +50,7 @@ relevant investigations. Save the report to `analysis/matmul_report.md`. ``` This will: + 1. Compile and profile the kernel with the right env vars 2. Ingest into neuron-explorer 3. Compute all bounds (memory, compute, pipeline families) @@ -65,6 +67,7 @@ This will: > are DMA-starved in that window vs the rest of the kernel? This will: + 1. Load the parquet tables 2. Set the time window to [t0, t0 + 5000 us] 3. Run the excess initiation interval analysis within that window diff --git a/skills/neuron-nki-profile-querying/references/investigations/dma_efficiency.md b/skills/neuron-nki-profile-querying/references/investigations/dma_efficiency.md index 9b0ce6a..94dc9fc 100644 --- a/skills/neuron-nki-profile-querying/references/investigations/dma_efficiency.md +++ b/skills/neuron-nki-profile-querying/references/investigations/dma_efficiency.md @@ -14,14 +14,14 @@ bandwidth, but not why. The investigation covers one potential source for this g Small transfers spend a larger fraction of their time on overhead vs actual data movement, reducing effective bandwidth. -This investigation quantifies the efficiency gap and traces it to specific source lines. -DMA transposes from HBM to SBUF contribute to dma inefficiency but aren't covered here -yet. +This investigation quantifies the efficiency gap and traces it to specific source lines. +DMA transposes from HBM to SBUF contribute to dma inefficiency but aren't covered here +yet. Note: the efficiency gap often coexists with excess HBM traffic. Excess inefficient traffic can pump up both the inefficiency gap and the excess gap. If the excess gap is proportionally -large, it potentially represents that proportion of the inefficiency gap as well. Consider -investigating that one first in such a case since reducing it will potentially reduce the +large, it potentially represents that proportion of the inefficiency gap as well. Consider +investigating that one first in such a case since reducing it will potentially reduce the number of inefficient transfers. ## Prerequisites @@ -117,16 +117,15 @@ If `efficiency_gap_us` is small relative to `memory_bound` (e.g. < 5%), DMA transfers are already reasonably efficient. If the kernel is still memory-bound, the issue might be excess traffic (reloads/spills), not per-transfer efficiency. - ## Step 2: Localize — per-source-line transfer geometry ### Check DPA coverage Step 2 joins Instruction → DmaPacketAggregated via the Flow table. DPA may only cover a subset of kernel transfers depending on the DGE -mode used (see Known Issues). Check coverage first — if DPA is missing, -skip step 2 entirely and if it undercounts significantly, note that -Step 2 results will be incomplete. +mode used (see Known Issues). Check coverage first — if DPA is missing, +skip step 2 entirely and if it undercounts significantly, note that +Step 2 results will be incomplete. ```python dma_agg_path = f"{d}/DmaPacketAggregated.parquet" @@ -215,44 +214,44 @@ transposes. Input: 128x8192 bf16. ### Step 1: Detect -| Metric | Value | -|--------|-------| -| memory\_bound | 16 us | -| memory\_bound\_ideal | 10 us | -| efficiency\_gap | 7 us | -| achieved\_bw | 259 GB/s (peak: 435 GB/s) | +| Metric | Value | +| ------------------ | ------------------------- | +| memory_bound | 16 us | +| memory_bound_ideal | 10 us | +| efficiency_gap | 7 us | +| achieved_bw | 259 GB/s (peak: 435 GB/s) | -The dma engine is clearly operating way below expected efficiency. +The dma engine is clearly operating way below expected efficiency. ### Step 2: Localize DPA coverage: 50% — swdge and unknown transfers missed variable attribution. -| Source line | Tensor | Dir | Xfers | MB | % of total | Shape | -|-------------|--------|-----|------:|---:|-----------|-------| -| dge\_mixed.py:21 | input0 | load | 4 | 0.524 | 12.3% | 128x1024B | -| dge\_mixed.py:22 | output0 | store | 4 | 0.524 | 12.3% | 128x1024B | -| dge\_mixed.py:35 | input0 | load | 4 | 0.524 | 12.3% | 128x1024B | -| dge\_mixed.py:36 | output0 | store | 4 | 0.524 | 12.3% | 128x1024B | -| dge\_mixed.py:51 | output0 | store | 4 | 0.016 | 0.4% | 16x256B | +| Source line | Tensor | Dir | Xfers | MB | % of total | Shape | +| --------------- | ------- | ----- | ----: | ----: | ---------- | --------- | +| dge_mixed.py:21 | input0 | load | 4 | 0.524 | 12.3% | 128x1024B | +| dge_mixed.py:22 | output0 | store | 4 | 0.524 | 12.3% | 128x1024B | +| dge_mixed.py:35 | input0 | load | 4 | 0.524 | 12.3% | 128x1024B | +| dge_mixed.py:36 | output0 | store | 4 | 0.524 | 12.3% | 128x1024B | +| dge_mixed.py:51 | output0 | store | 4 | 0.016 | 0.4% | 16x256B | From the covered lines analyzed, we can see that almost 50% of the kernel's transfer bytes are from transfers with small descriptor sizes (1kib instead of 4Kib). Increasing the free dimension -of these transfers may improve the DMA efficiency of the kernel. +of these transfers may improve the DMA efficiency of the kernel. ### Comparison: 128x2048 tiles Same kernel structure with 4x larger tiles (`dge_mixed_large.py`). -| Metric | 128x512 | 128x2048 | -|--------|---------|----------| -| memory\_bound | 16 us | 47 us | -| memory\_bound\_ideal | 10 us | 39 us | -| efficiency\_gap | 7 us | 9 us | -| achieved\_bw | 259 GB/s | 354 GB/s | +| Metric | 128x512 | 128x2048 | +| ------------------ | -------- | -------- | +| memory_bound | 16 us | 47 us | +| memory_bound_ideal | 10 us | 39 us | +| efficiency_gap | 7 us | 9 us | +| achieved_bw | 259 GB/s | 354 GB/s | Descriptor sizes scale from 1024B to 4096B. Achieved bandwidth improves -from 259 to 354 GB/s with larger transfers. +from 259 to 354 GB/s with larger transfers. ## Known issues diff --git a/skills/neuron-nki-profile-querying/references/investigations/redundant_dma_transfers.md b/skills/neuron-nki-profile-querying/references/investigations/redundant_dma_transfers.md index 63a330c..02a4309 100644 --- a/skills/neuron-nki-profile-querying/references/investigations/redundant_dma_transfers.md +++ b/skills/neuron-nki-profile-querying/references/investigations/redundant_dma_transfers.md @@ -1,12 +1,11 @@ -# Investigation: Redundant DMA transfers +# Investigation: Redundant DMA transfers ## Context This investigation follows from a large `memory_bound_ideal → memory_bound_ideal_no_reloads` gap in the [performance bounds](../performance-bounds.md). That gap measures the cost — at peak bandwidth — of transferring more data from HBM than the algorithm -requires. At the kernel's actual (lower) bandwidth, the real cost is proportionally larger. - +requires. At the kernel's actual (lower) bandwidth, the real cost is proportionally larger. The bounds identify that excess traffic exists but not the cause. The excess bytes could be from: @@ -18,7 +17,7 @@ be from: This investigation quantifies the excess traffic, decomposes it by source, and traces it to specific NKI source lines. DMA transposes from SBUF to SBUF can also contribute - excess redundant work but are not yet covered. +excess redundant work but are not yet covered. ## Prerequisites @@ -72,12 +71,12 @@ print(f"excess_ratio: {excess_ratio:>8.1f}x") `dma_transfer_bytes` is the total data moved across all kernel DMA packets. `necessary_bytes` is the sum of all input and output tensor sizes from `TensorInfo` — the minimum if each tensor were transferred exactly once. -An `excess_ratio` of 1.0 means no excess. +An `excess_ratio` of 1.0 means no excess. -If your kernel algorithmically doesn't use the full input and output -tensors, this will overestimate necessary bytes. Make a note of this and -replace with your actual necessary bytes if this is the case (most likely, -it's not!) +If your kernel algorithmically doesn't use the full input and output +tensors, this will overestimate necessary bytes. Make a note of this and +replace with your actual necessary bytes if this is the case (most likely, +it's not!) ### Decompose excess: reloads vs spills @@ -127,7 +126,6 @@ but not in `TensorInfo`. than once. The per-tensor excess is `total_bytes - tensor_size`. Tensors not in `TensorInfo` (spill buffers) are skipped. - ## Step 2a: Localize reloads — which source lines load which tensors? Run this step if `reload_excess` dominates excess traffic in Step 1. @@ -226,7 +224,6 @@ Each row is a source line: the NKI operation whose output the compiler spilled. `spill_cost` is the full round-trip cost (save + reload) attributed to that line. Rank by `spill_cost_mb` to find the dominant contributor. - ## Known issues - **`TensorInfo.load_to_sbuf_repeat_factor`** (and `load_to_sbuf_dma_count`, @@ -256,4 +253,3 @@ Rank by `spill_cost_mb` to find the dominant contributor. (SB→VIRTUAL) flow edges may be missing for 5-40% of spill transfers. Since each buffer is saved and reloaded for identical bytes, `spill_cost = 2 × reload_bytes`. - diff --git a/skills/neuron-nki-profile-querying/references/investigations/redundant_te_transposes.md b/skills/neuron-nki-profile-querying/references/investigations/redundant_te_transposes.md index 1d6ee68..9bc9825 100644 --- a/skills/neuron-nki-profile-querying/references/investigations/redundant_te_transposes.md +++ b/skills/neuron-nki-profile-querying/references/investigations/redundant_te_transposes.md @@ -12,7 +12,7 @@ This investigation quantifies the transpose FLOPs overhead at peak throughput and traces it to specific NKI source lines. Note that there may be more redundant TE instructions like unnecessary computation -but this is not as easily / objectively quantifiable so the gap doesn't include them. +but this is not as easily / objectively quantifiable so the gap doesn't include them. ## Prerequisites @@ -98,6 +98,7 @@ Override `t0` and `t1`, then re-run the query above. Useful when: against regular MATMULs in a specific window Example — isolate the transpose phase of the first m-block: + ```python tp_sorted = transpose.sort_values('start_ts') t0 = int(tp_sorted.iloc[0]['start_ts']) @@ -105,7 +106,6 @@ t1 = int(tp_sorted.iloc[15]['end_ts']) # first 16 transposes # Re-run the Step 1 query with this t0, t1 ``` - ## Step 2: Localize — which source lines produce the transposes? This operates on the `transpose` DataFrame from Step 1. If Step 1 used a @@ -133,7 +133,6 @@ for src, row in by_src.iterrows(): Each group corresponds to an `nisa.nc_transpose` call in the NKI source. Rank by `flops_g` to identify which transpose call contributes most to the gap. - ## Worked example ### The kernels @@ -159,14 +158,14 @@ for m in nl.affine_range(M // TILE_M): # 16 ### Step 1 -| Metric | V0 | V1 | -|--------|----|----| -| hw\_flops | 19.33 G | 17.18 G | -| transpose\_flops | 2.15 G (11.1% of hw) | 0 | -| useful\_flops | 17.18 G | 17.18 G | -| TRANSPOSE MATMULs | 1,024 | 0 | -| REGULAR MATMULs | 1,024 | 1,024 | -| gap at peak TE | 27.3 us | 0 | +| Metric | V0 | V1 | +| ----------------- | -------------------- | ------- | +| hw_flops | 19.33 G | 17.18 G | +| transpose_flops | 2.15 G (11.1% of hw) | 0 | +| useful_flops | 17.18 G | 17.18 G | +| TRANSPOSE MATMULs | 1,024 | 0 | +| REGULAR MATMULs | 1,024 | 1,024 | +| gap at peak TE | 27.3 us | 0 | V0 has equal TRANSPOSE and REGULAR MATMULs — one `nc_transpose` per `nc_matmul`. The 2.15 G of transpose FLOPs accounts for 11.1% of total @@ -174,26 +173,25 @@ hardware FLOPs. V1 eliminates all transposes; `useful_flops` is unchanged. ### Step 2 -| Source line | count | GFLOPS | gap (us) | % of transpose | -|-------------|-------|--------|----------|----------------| -| v0\_with\_transpose.py:34 | 1,024 | 2.15 | 27.3 | 100% | +| Source line | count | GFLOPS | gap (us) | % of transpose | +| ----------------------- | ----- | ------ | -------- | -------------- | +| v0_with_transpose.py:34 | 1,024 | 2.15 | 27.3 | 100% | All 1,024 TRANSPOSE MATMULs in V0 come from the `nisa.nc_transpose` call at line 34. V1 has no transposes — Step 2 produces no output. ### Bounds comparison -| Bound | V0 (us) | V1 (us) | Change | -|-------|---------|---------|--------| -| total\_time | 802 | 608 | -194 us (1.3x) | -| compute\_bound (TE active) | 503 | 420 | -83 us | -| compute\_bound\_ideal | 246 | 218 | -28 us | -| compute\_bound\_ideal\_useful | 218 | 218 | — | +| Bound | V0 (us) | V1 (us) | Change | +| -------------------------- | ------- | ------- | -------------- | +| total_time | 802 | 608 | -194 us (1.3x) | +| compute_bound (TE active) | 503 | 420 | -83 us | +| compute_bound_ideal | 246 | 218 | -28 us | +| compute_bound_ideal_useful | 218 | 218 | — | `compute_bound_ideal` equals `compute_bound_ideal_useful` in V1 — all TensorE FLOPs are useful. Total kernel time reduced by 194 us (24%). - ## Known issues - **`Summary.transpose_flops`**: Returns NaN when 0 transposes exist diff --git a/skills/neuron-nki-profile-querying/references/investigations/te_inefficiency.md b/skills/neuron-nki-profile-querying/references/investigations/te_inefficiency.md index b22f41b..60ce24c 100644 --- a/skills/neuron-nki-profile-querying/references/investigations/te_inefficiency.md +++ b/skills/neuron-nki-profile-querying/references/investigations/te_inefficiency.md @@ -7,7 +7,7 @@ This investigation follows from a large [performance bounds](../performance-bounds.md). That gap measures TensorE active time not producing FLOPs at peak rate — the hardware is executing but underutilized. We will investigate tile sizes of TE -matmul instructions in this investigation but not all the gap may be +matmul instructions in this investigation but not all the gap may be explained. In some cases of poor utilization or poor pipelining, throttling can inflate instruction durations. @@ -95,8 +95,7 @@ print(f"Achieved: {achieved_tflops:.1f} TFLOPS ({utilization:.0%})") `adjusted_flops` normalizes each instruction to bf16-equivalent FLOPs (1× for bf16/fp16/tf32, 4× for fp32, 0.5× for fp8 on trn2), matching `te_peak` which is the bf16 peak. `utilization` is the fraction of TE -active time producing bf16-equivalent FLOPs at peak rate. - +active time producing bf16-equivalent FLOPs at peak rate. ## Step 2: Localize — which source lines use undersized tiles? @@ -132,7 +131,7 @@ minimal impact on streaming throughput because the initiation interval shortens proportionally. Before pointing out a gap. Verify that the user is not doing complicated PE -tiling intentionally, this is not in scope. +tiling intentionally, this is not in scope. ## Worked example @@ -175,11 +174,11 @@ for _ in nl.affine_range(512): ### Step 1 -| Metric | V0 (mixed) | V1 (peak) | -|--------|-----------|-----------| -| TE peak | 78.6 TFLOPS | 78.6 TFLOPS | +| Metric | V0 (mixed) | V1 (peak) | +| -------- | ----------------- | ----------------- | +| TE peak | 78.6 TFLOPS | 78.6 TFLOPS | | Achieved | 30.8 TFLOPS (39%) | 77.2 TFLOPS (98%) | -| MATMULs | 2,048 | 2,048 | +| MATMULs | 2,048 | 2,048 | V0's overall utilization is 39% — a weighted average across the four sections. V1 achieves 98% with all peak tiles. @@ -188,25 +187,24 @@ sections. V1 achieves 98% with all peak tiles. **V0** — four source lines with different tile dimensions: -| Source line | Tiles | % of flops | K/128 | M/128 | N/512 | Undersized | -|-------------|-------|-----------|-------|-------|-------|------------| -| v0\_mixed\_tiles.py:30 | 512 | 15% | 128/128 | 32/128 | 512/512 | M | -| v0\_mixed\_tiles.py:35 | 512 | 8% | 128/128 | 128/128 | 64/512 | N | -| v0\_mixed\_tiles.py:40 | 512 | 15% | 32/128 | 128/128 | 512/512 | K | -| v0\_mixed\_tiles.py:45 | 512 | 62% | 128/128 | 128/128 | 512/512 | none | +| Source line | Tiles | % of flops | K/128 | M/128 | N/512 | Undersized | +| -------------------- | ----- | ---------- | ------- | ------- | ------- | ---------- | +| v0_mixed_tiles.py:30 | 512 | 15% | 128/128 | 32/128 | 512/512 | M | +| v0_mixed_tiles.py:35 | 512 | 8% | 128/128 | 128/128 | 64/512 | N | +| v0_mixed_tiles.py:40 | 512 | 15% | 32/128 | 128/128 | 512/512 | K | +| v0_mixed_tiles.py:45 | 512 | 62% | 128/128 | 128/128 | 512/512 | none | **V1** — single source line, all peak: -| Source line | Tiles | % of flops | K/128 | M/128 | N/512 | Undersized | -|-------------|-------|-----------|-------|-------|-------|------------| -| v1\_peak.py:28 | 2,048 | 100% | 128/128 | 128/128 | 512/512 | none | +| Source line | Tiles | % of flops | K/128 | M/128 | N/512 | Undersized | +| ------------- | ----- | ---------- | ------- | ------- | ------- | ---------- | +| v1_peak.py:28 | 2,048 | 100% | 128/128 | 128/128 | 512/512 | none | Each undersized section contributes 512 tiles but different fractions of total FLOPs. The N=64 section produces only 8% of FLOPs (64/512 = 12.5% of peak per tile). The peak section produces 62% despite being only 512 of 2048 tiles. - ## Known issues - **`Instruction.operands` format**: The `K*M` trailing pair and `src` diff --git a/skills/neuron-nki-profile-querying/references/performance-bounds.md b/skills/neuron-nki-profile-querying/references/performance-bounds.md index 2c97b21..713b82e 100644 --- a/skills/neuron-nki-profile-querying/references/performance-bounds.md +++ b/skills/neuron-nki-profile-querying/references/performance-bounds.md @@ -84,16 +84,16 @@ Each bound is a formula over a small set of quantities. For the whole kernel, these come from Summary. For a time window `[t0, t1]`, they are recomputed from the raw tables as shown below. -| Quantity | Whole kernel | Time window `[t0, t1]` | -|----------|-------------|------------------------| -| **total_time** | `Summary.total_time` | `t1 - t0` | -| **dma_active_time** | `Summary.dma_active_time` | Interval-merge `DmaPacket` rows overlapping `[t0, t1]`, clipped to window edges | -| **dma_transfer_bytes** | `Summary.dma_transfer_total_bytes` | `SUM(transfer_bytes)` from `DmaPacket` where `queue_type != 'instruction'` and `transfer_bytes > 4`, starting in `[t0, t1)` | -| **necessary_bytes** | `Summary.inputs_outputs_weights_size_bytes` | Sum of `TensorInfo.size` for tensors with `type` in `['IN', 'OUT']` | -| **te_active_time** | `Summary.tensor_engine_active_time` | Sum `ActiveTime.duration_ns` where `engine = 'tensor'`, clipped to `[t0, t1]` | -| **hw_flops** | `Summary.hardware_flops` | `SUM(adjusted_flops)` from MATMUL `Instruction` rows in `[t0, t1)` | -| **transpose_flops** | `Summary.transpose_flops` | Same, filtered to `tensor_instruction_type = 'TRANSPOSE'` | -| **engine_active_times** | `Summary.*_engine_active_time` | Per-engine interval merge from `ActiveTime` clipped to window | +| Quantity | Whole kernel | Time window `[t0, t1]` | +| ----------------------- | ------------------------------------------- | --------------------------------------------------------------------------------------------------------------------------- | +| **total_time** | `Summary.total_time` | `t1 - t0` | +| **dma_active_time** | `Summary.dma_active_time` | Interval-merge `DmaPacket` rows overlapping `[t0, t1]`, clipped to window edges | +| **dma_transfer_bytes** | `Summary.dma_transfer_total_bytes` | `SUM(transfer_bytes)` from `DmaPacket` where `queue_type != 'instruction'` and `transfer_bytes > 4`, starting in `[t0, t1)` | +| **necessary_bytes** | `Summary.inputs_outputs_weights_size_bytes` | Sum of `TensorInfo.size` for tensors with `type` in `['IN', 'OUT']` | +| **te_active_time** | `Summary.tensor_engine_active_time` | Sum `ActiveTime.duration_ns` where `engine = 'tensor'`, clipped to `[t0, t1]` | +| **hw_flops** | `Summary.hardware_flops` | `SUM(adjusted_flops)` from MATMUL `Instruction` rows in `[t0, t1)` | +| **transpose_flops** | `Summary.transpose_flops` | Same, filtered to `tensor_instruction_type = 'TRANSPOSE'` | +| **engine_active_times** | `Summary.*_engine_active_time` | Per-engine interval merge from `ActiveTime` clipped to window | **Interval merging.** `dma_active_time` is the wall-clock time when any of the 16 DMA engines was active. It equals the interval-merge of all @@ -115,7 +115,7 @@ time = dma_active_time ``` Reflects the actual DMA workload: real transfer sizes, real -packet efficiency, real bandwidth utilization. Assuming all other work is +packet efficiency, real bandwidth utilization. Assuming all other work is pipelined behind data transfers. ```python @@ -132,9 +132,9 @@ time = dma_transfer_bytes / dma_bw_peak Same total DMA traffic, but at peak HBM bandwidth with zero per-transfer overhead. The gap `memory_bound → memory_bound_ideal` is the cost of DMA inefficiency: small packets, low per-transfer throughput, inefficient dma - transposes. Note that peak bandwidth may not be achievable for a given +transposes. Note that peak bandwidth may not be achievable for a given transfer pattern — per-transfer overhead and packet size constraints set a -practical ceiling below the theoretical peak. +practical ceiling below the theoretical peak. ```python kernel_pkts = dma_pkts[(dma_pkts['queue_type'] != 'instruction') @@ -252,6 +252,7 @@ non-improved engine is slowest: ``` bn_speedup = total_time / max(bound_time, max_non_impacted_engine_time) ``` + In practice, engine execution times can be intertwined. For example, improving memory efficiency might reduce Tensor Engine time as it is less often throttled after idle gaps. @@ -265,20 +266,20 @@ how much of the engine's time is attributable to that category of overhead. ### Memory family gaps ``` -1. memory_bound ─→ total_time : DMA idle gaps +1. memory_bound ─→ total_time : DMA idle gaps 2. memory_bound ─→ memory_bound_ideal : DMA inefficiency (BW utilization) 3. memory_bound_ideal ─→ no_reloads : excess HBM traffic (reloads + spills) ``` All three of these gaps can be interpreted as lost time under the assumption that DMA is the bottleneck engine. Gap 1 tells us time lost to DMA idle gaps. Gap 2 tells us time -lost to low bandwidth utilization but is theoretical and may not be achievable. Gap 3 -should now be intepreted / approximated as a proportional loss (since it the gap is counted at -ideal bw utilization). Excess inefficient loads potentially contribute more to gap 2 then 3. +lost to low bandwidth utilization but is theoretical and may not be achievable. Gap 3 +should now be intepreted / approximated as a proportional loss (since it the gap is counted at +ideal bw utilization). Excess inefficient loads potentially contribute more to gap 2 then 3. -Technically, gap 1 would need to go to 0 for gap 2 to be interpreted exactly as lost time and so on. +Technically, gap 1 would need to go to 0 for gap 2 to be interpreted exactly as lost time and so on. This is because reducing gap 2, even in a memory bound kernel, may not lead to an improvement in a poorly pipelined case -as another engine might be active anyways at that time. +as another engine might be active anyways at that time. ### Compute family gaps @@ -290,24 +291,25 @@ as another engine might be active anyways at that time. As in the memory case, these gaps can be intepreted as lost time under the assumption that Tensor Engine is the bottleneck. Gap 1 and 2 can be interpreted as raw values with the same caveats as before. Gap 3 -should be intepreted as a proportion of TE engine time. +should be intepreted as a proportion of TE engine time. ### Pipeline gap In the case where neither DMA nor Tensor Engine is the bottleneck, we are interested -in the gap to the bottleneck engine. +in the gap to the bottleneck engine. + ``` total_time ─→ perfect_pipeline : total serialization across all engines ``` -If this gap is small, we will need to look at ineffiency and redundant computation on that engine +If this gap is small, we will need to look at ineffiency and redundant computation on that engine but the calculation is non-trivial. If the gap is large, pipelining across engines should be the priority. This skill does not yet support it but you may extrapolate from the existing investigations -if EXPLICITELY permitted. +if EXPLICITELY permitted. ### Gaps are not improvement deltas -The gaps measure overhead *within a category*. They do not predict how much +The gaps measure overhead _within a category_. They do not predict how much faster the kernel will be if that overhead is eliminated. Three reasons: **1. Gaps are computed under idealized assumptions.** The memory family gaps @@ -330,7 +332,7 @@ TensorE's active time is inflated because individual MATMUL instructions take longer when weight data arrives late. Reducing DMA overhead (a memory optimization) feeds TensorE faster, which shrinks TE active time — even though no compute optimization was applied. This means the compute family -bounds computed from the *current* TE active time overstate the compute +bounds computed from the _current_ TE active time overstate the compute overhead that would remain after fixing the memory side. The reverse can also occur: improving TE pipelining can reduce the DMA working set if it changes how tiles are scheduled. @@ -356,21 +358,22 @@ gap to an optimization group and lists the available investigations. ### Inefficiency groups -For all engines there are the following *groups* of inefficiencies. For -(1) Tensor Engine and (2) DMA, and (3) other bottleneck engine, we defin the following -groups: (a) Idle gaps -> (b) engine underutilization -> (c) redundant instructions. +For all engines there are the following _groups_ of inefficiencies. For +(1) Tensor Engine and (2) DMA, and (3) other bottleneck engine, we defin the following +groups: (a) Idle gaps -> (b) engine underutilization -> (c) redundant instructions. **Group 1a: DMA engine idle gaps** Gap: `total_time → memory_bound` -DMA idle gaps are tough to interpret, do not make a statement about the cause of this gap. +DMA idle gaps are tough to interpret, do not make a statement about the cause of this gap. **Group 1b: DMA inefficiency** Gap: `memory_bound → memory_bound_ideal` [DMA Efficiency investigation](investigations/dma_efficiency.md) -Covers: +Covers: + - Dma transfer sizes **Group 1c: Redundant DMA transfers from HBM** @@ -378,7 +381,8 @@ Gap: `memory_bound_ideal → memory_bound_ideal_no_reloads` [Redundant dma transfers investigation](investigations/redundant_dma_transfers.md) -Covers: +Covers: + - Excess Input Reloads - Intermediate Data Spilling @@ -386,7 +390,7 @@ Covers: Gap: `total_time → compute_bound` This gap can be difficult to interpret due to the currently available information -regarding dependencies and anti-dependencies. Coming soon! +regarding dependencies and anti-dependencies. Coming soon! **Group 2b: Compute engine underutilization** Gap: `compute_bound → compute_bound_ideal_flops` @@ -396,12 +400,14 @@ cannot separate — insufficient tile sizes, instruction placement, fast weight throttling all live in the same gap. Investigations: + - [TE Inefficiency](investigations/te_inefficiency.md) **Group 2c: Redundant TE engine instructions** Gap: `compute_bound_ideal_flops → compute_bound_ideal_useful_flops` Investigations: + - [Redundant TE Transposes](investigations/redundant_te_transposes.md) **Group 3: Efficiency gaps when Vector/Scalar/Gpsimd is the bottleneck** @@ -411,17 +417,17 @@ When the bottleneck is neither DMA nor Tensor Engine, analysis is less straightforward since Vector, Scalar and Gpsimd play diverse roles within the kernel. Gap 3a (idle gaps) can be found as `total_time → perfect_pipeline` but defining efficiency and "minimal" workload is instruction implementation - specific. +specific. ### Reading the gap structure If an engine is by far the bottleneck, focus on that family of gaps. In practice multiple engines may be quite close, and in this case, focus on all bottlenecks. In -such a case, logically section the kernel and run windowed analysis. +such a case, logically section the kernel and run windowed analysis. Within a family, the relative sizes of the gaps indicate where overhead concentrates. However, gap c should not be read as an absolute value, it is inherently -proportional to the engine workload. +proportional to the engine workload. Gaps are not mutually exclusive (even across families) — a kernel can have significant overhead in multiple gaps simultaneously, and addressing one diff --git a/skills/neuron-nki-profiling/SKILL.md b/skills/neuron-nki-profiling/SKILL.md index 744e64a..8ee85df 100644 --- a/skills/neuron-nki-profiling/SKILL.md +++ b/skills/neuron-nki-profiling/SKILL.md @@ -36,6 +36,7 @@ neuron-explorer view --output-format summary-json -n $NEFF_PATH -s profile.ntff ``` The workflow generates two key artifacts: + - **NEFF file**: Compiled kernel binary, generated during execution - **NTFF file**: Execution trace captured by neuron-explorer @@ -48,6 +49,7 @@ Before profiling kernels, resolve the NKI virtual environment path: 3. If still not found, report: "NKI_VENV_PATH not configured. Set the environment variable or create .claude/nki-dev-suite.local.md with nki_venv_path in frontmatter." Activate before running any profiling commands: + ```bash source $NKI_VENV_PATH/bin/activate ``` @@ -69,18 +71,18 @@ os.environ['NEURON_RT_INSPECT_DEVICE_PROFILE'] = '1' os.environ['NEURON_RT_INSPECT_OUTPUT_DIR'] = './output' # Compiler flags for target hardware -os.environ['NEURON_CC_FLAGS'] = '--target trn2 --lnc 1' # use lnc=2 if explicitely told to. +os.environ['NEURON_CC_FLAGS'] = '--target trn2 --lnc 1' # use lnc=2 if explicitely told to. # Pin to a specific neuron core(s) to avoid conflicts with concurrent sessions os.environ['NEURON_RT_VISIBLE_CORES'] = '0' # '0,1', '0-1' ``` -| Environment Variable | Description | -|---------------------|-------------| -| `NEURON_RT_INSPECT_ENABLE` | Enable runtime inspection | -| `NEURON_RT_INSPECT_DEVICE_PROFILE` | Enable device-level profiling | -| `NEURON_RT_INSPECT_OUTPUT_DIR` | Directory for NEFF output | -| `NEURON_RT_VISIBLE_CORES` | Pin to specific core(s) — prevents contention when multiple agents profile concurrently | +| Environment Variable | Description | +| ---------------------------------- | --------------------------------------------------------------------------------------- | +| `NEURON_RT_INSPECT_ENABLE` | Enable runtime inspection | +| `NEURON_RT_INSPECT_DEVICE_PROFILE` | Enable device-level profiling | +| `NEURON_RT_INSPECT_OUTPUT_DIR` | Directory for NEFF output | +| `NEURON_RT_VISIBLE_CORES` | Pin to specific core(s) — prevents contention when multiple agents profile concurrently | ### Step 2: Execute Kernel @@ -103,6 +105,7 @@ result = my_nki_kernel(lhs, rhs) # Another NEFF ``` The runtime creates a subdirectory with instance and process ID naming: + ``` ./output/ └── i-0823210096b01e7ec_pid_1187583/ @@ -118,6 +121,7 @@ mkdir -p ./profiles/run_001 ``` Organize profile iterations: + ``` ./profiles/ ├── run_001/ # Baseline profiling @@ -145,12 +149,12 @@ neuron-explorer capture \ --enable-dge-notifs ``` -| Flag | Description | -|------|-------------| -| `-n` | Path to NEFF file | -| `-s` | Output path for NTFF trace file | -| `--profile-nth-exec=2` | Profile the 2nd execution (skip warmup) | -| `--enable-dge-notifs` | Enable DMA engine notifications for detailed analysis | +| Flag | Description | +| ---------------------- | ----------------------------------------------------- | +| `-n` | Path to NEFF file | +| `-s` | Output path for NTFF trace file | +| `--profile-nth-exec=2` | Profile the 2nd execution (skip warmup) | +| `--enable-dge-notifs` | Enable DMA engine notifications for detailed analysis | ### Step 5: View Results with neuron-explorer (JSON) @@ -200,7 +204,7 @@ neuron-explorer view \ ### Step 6: Querying the profile and/or profile analysis (optional) For detailed analysis of the kernel profile, use the /neuron-nki-profile-querying skill. It allows for high level performance bounds analysis, as well as zoomed in, instruction level -investigation of specific inefficiencies through python on parquet. +investigation of specific inefficiencies through python on parquet. ## Output Directory Structure @@ -248,13 +252,13 @@ NEFF_PATH=$(python3 scripts/identify-neffs.py ./output/i-*_pid_*/ matmul_relu) ## Key Metrics Quick Reference -| Metric | Description | Target | -|--------|-------------|--------| -| `latency` | Total kernel execution time (ms) | Lower is better | -| `tensor_engine_active_time_percent` | TensorE utilization | >90% for compute-bound | -| `hbm_read_bytes` | HBM read traffic | Minimize | -| `hbm_write_bytes` | HBM write traffic | Minimize | -| `mm_arithmetic_intensity` | FLOPs per byte of memory traffic | Compare to peak ratio | +| Metric | Description | Target | +| ----------------------------------- | -------------------------------- | ---------------------- | +| `latency` | Total kernel execution time (ms) | Lower is better | +| `tensor_engine_active_time_percent` | TensorE utilization | >90% for compute-bound | +| `hbm_read_bytes` | HBM read traffic | Minimize | +| `hbm_write_bytes` | HBM write traffic | Minimize | +| `mm_arithmetic_intensity` | FLOPs per byte of memory traffic | Compare to peak ratio | ## Comparing Optimization Iterations @@ -276,11 +280,11 @@ echo "Optimized: $(jq .latency ./profiles/optimized/metrics.json)" **Optimization tracking table:** -| Iteration | Change | Latency (ms) | TensorE (%) | -|-----------|--------|--------------|-------------| -| Baseline | - | 1.23 | 45% | -| Larger tiles | Increased tile 64→128 | 0.95 | 72% | -| Double buffer | Added prefetching | 0.78 | 89% | +| Iteration | Change | Latency (ms) | TensorE (%) | +| ------------- | --------------------- | ------------ | ----------- | +| Baseline | - | 1.23 | 45% | +| Larger tiles | Increased tile 64→128 | 0.95 | 72% | +| Double buffer | Added prefetching | 0.78 | 89% | Keep notes on what changed between iterations to correlate optimizations with metric improvements. @@ -292,54 +296,59 @@ See `examples/basic-profiling-workflow.py` for a complete end-to-end profiling s **Required settings:** -| Setting | Source | Description | -|---------|--------|-------------| +| Setting | Source | Description | +| --------------- | --------------------------------------------------- | --------------------------------- | | `nki_venv_path` | `.claude/nki-dev-suite.local.md` or `NKI_VENV_PATH` | Python venv with neuronx packages | **Environment variables (set in kernel script):** -| Variable | Value | Purpose | -|----------|-------|---------| -| `NEURON_RT_INSPECT_ENABLE` | `1` | Enable runtime inspection | -| `NEURON_RT_INSPECT_DEVICE_PROFILE` | `1` | Enable device profiling | -| `NEURON_RT_INSPECT_OUTPUT_DIR` | Path | NEFF output directory | +| Variable | Value | Purpose | +| ---------------------------------- | ----- | ------------------------- | +| `NEURON_RT_INSPECT_ENABLE` | `1` | Enable runtime inspection | +| `NEURON_RT_INSPECT_DEVICE_PROFILE` | `1` | Enable device profiling | +| `NEURON_RT_INSPECT_OUTPUT_DIR` | Path | NEFF output directory | ## Related Skills -| Skill | Purpose | -|-------|---------| -| `/neuron-nki-profile-querying` | Detailed profile querying and analysis | +| Skill | Purpose | +| --------------------------------- | -------------------------------------------------------------- | +| `/neuron-nki-profile-querying` | Detailed profile querying and analysis | | `/neuron-explorer-profile-schema` | Reference for the parquet schema produced by `neuron-explorer` | -| `/neuron-nki-debugging` | Debug compilation errors | -| `/neuron-nki-docs` | Look up API documentation | -| `/neuron-nki-writing` | Write NKI kernels | +| `/neuron-nki-debugging` | Debug compilation errors | +| `/neuron-nki-docs` | Look up API documentation | +| `/neuron-nki-writing` | Write NKI kernels | ## Troubleshooting **Multiple NEFFs generated (can't tell which is the NKI kernel):** + - **Primary fix**: Compute reference operations (e.g., `torch.matmul`) on CPU, not on the XLA device. Each on-device XLA graph generates its own NEFF. - **Identify NEFFs**: Use `python3 scripts/identify-neffs.py ./output` to list all NEFFs with their type (NKI vs XLA) and kernel names. See the [NEFF Identification](#neff-identification) section for details. - **Match NEFF to NTFF**: Each NEFF `neff__vnc_0.neff` has a matching trace `_vnc_0.ntff` in the same directory. **No NEFF file generated:** + - Verify `NEURON_RT_INSPECT_ENABLE=1` is set before imports - Check `NEURON_RT_INSPECT_OUTPUT_DIR` path exists and is writable - Ensure kernel actually executed (print forces XLA compilation) - Confirm you are on Trainium/Inferentia hardware: `neuron-ls` **neuron-explorer capture fails:** + - Verify running on Trainium/Inferentia hardware - Check NEFF file path is correct with `ls -la ` - Ensure neuronx packages are installed in venv - Check sufficient disk space for NTFF file **Empty or minimal profile data:** + - Use `--profile-nth-exec=2` to skip warmup execution - Add `--enable-dge-notifs` for detailed DMA analysis - Verify kernel ran successfully before profiling - Check NTFF file size is non-trivial: `ls -lh profile.ntff` **Latency varies between runs:** + - Use `--profile-nth-exec=2` or higher to skip warmup - Ensure system is not under other load - Run multiple iterations and average results diff --git a/skills/neuron-nki-writing/SKILL.md b/skills/neuron-nki-writing/SKILL.md index 1ac9f24..3162b89 100644 --- a/skills/neuron-nki-writing/SKILL.md +++ b/skills/neuron-nki-writing/SKILL.md @@ -35,7 +35,7 @@ so do not use it as the free-dim limit. Authoritative source: the `nc_matmul`, PSUM free-dim number inlined in this skill's reference files as illustrative; confirm against those API blocks. -`tile_size` *does* authoritatively report other limits — use it for those: +`tile_size` _does_ authoritatively report other limits — use it for those: `nl.tile_size.pmax` (128), `nl.tile_size.psum_num_banks` (bank cycling), `nl.tile_size.gemm_moving_fmax` (matmul moving-operand SBUF free dim), and `nl.tile_size.sbuf_fmax` / `sbuf_fmax_bytes` (SBUF capacity). @@ -74,17 +74,20 @@ def my_kernel(input_hbm: nl.ndarray) -> nl.ndarray: Before reading references, classify the task to avoid unnecessary overhead: **Simple** (element-wise op, single reduction, activation, layernorm, add/multiply): + - Use the Quick Start template and Step 4 API table directly - Skip utility library references entirely - **Start writing code immediately** — consult references only when stuck - Target: working kernel within 5 minutes **Medium** (matmul, softmax, multi-step fusion, transpose with tiling): + - Read `references/common-patterns.md`, `references/api-translation.md` and `references/memory-patterns.md` - Skip utility library references unless tiling is complex - Target: working kernel within 15 minutes **Complex** (multi-head attention, transformer blocks, state-space models, MoE): + - Full reference loading appropriate - Read utility selection guide and relevant patterns - Target: working kernel within 30 minutes @@ -101,12 +104,12 @@ Map PyTorch/NumPy operations to NKI equivalents using `references/api-translatio NKI operates on tiles with hardware constraints: -| Constraint | Limit | Notes | -|------------|-------|-------| -| Partition dimension (P) | ≤ 128 | First dimension of SBUF tensor | -| PSUM free dimension | 512 (gen2/3); gen4: 4096 fp32 / 8192 bf16 | For matrix multiply results — authoritative: `nc_matmul` doc block (see "Hardware limits" above) | -| SBUF free dimension | ≤ 32767 | Second+ dimensions | -| MatMul K dimension | ≤ 2048 | Contraction dimension | +| Constraint | Limit | Notes | +| ----------------------- | ----------------------------------------- | ------------------------------------------------------------------------------------------------ | +| Partition dimension (P) | ≤ 128 | First dimension of SBUF tensor | +| PSUM free dimension | 512 (gen2/3); gen4: 4096 fp32 / 8192 bf16 | For matrix multiply results — authoritative: `nc_matmul` doc block (see "Hardware limits" above) | +| SBUF free dimension | ≤ 32767 | Second+ dimensions | +| MatMul K dimension | ≤ 2048 | Contraction dimension | For tensors exceeding limits, use explicit tiling with `TiledRange` for remainder-safe iteration (see Utility Selection Guide below). @@ -156,21 +159,21 @@ Use multiple complementary checks (atol/rtol, max absolute difference, tensor no ## Hardware Constraints Quick Reference -| Buffer | Max P | Max F | Use Case | -|--------|-------|-------|----------| -| `nl.sbuf` | 128 | 32767 | General compute | -| `nl.psum` | 128 | 512 (gen2/3); gen4: 4096 fp32 / 8192 bf16 | MatMul accumulation | -| `nl.shared_hbm` | - | - | Input/output tensors | +| Buffer | Max P | Max F | Use Case | +| --------------- | ----- | ----------------------------------------- | -------------------- | +| `nl.sbuf` | 128 | 32767 | General compute | +| `nl.psum` | 128 | 512 (gen2/3); gen4: 4096 fp32 / 8192 bf16 | MatMul accumulation | +| `nl.shared_hbm` | - | - | Input/output tensors | ## Loop Types -| Loop Type | Use Case | Unrolling | -|-----------|----------|-----------| -| `nl.affine_range(N)` | Parallel iterations, no dependencies | Full unroll | -| `nl.sequential_range(N)` | Loop-carried dependencies (cumsum) | No unroll | -| `nl.static_range(N)` | Compile-time constant iterations | Partial unroll | -| `nl.fori_loop(lower, upper, body_fun, step=1)` | Counted loop with **runtime bound** | Structured on-chip loop (not unrolled) | -| `nl.while_loop(init, body_fun)` | **Data-dependent** condition-driven loop | Structured on-chip loop (not unrolled) | +| Loop Type | Use Case | Unrolling | +| ---------------------------------------------- | ---------------------------------------- | -------------------------------------- | +| `nl.affine_range(N)` | Parallel iterations, no dependencies | Full unroll | +| `nl.sequential_range(N)` | Loop-carried dependencies (cumsum) | No unroll | +| `nl.static_range(N)` | Compile-time constant iterations | Partial unroll | +| `nl.fori_loop(lower, upper, body_fun, step=1)` | Counted loop with **runtime bound** | Structured on-chip loop (not unrolled) | +| `nl.while_loop(init, body_fun)` | **Data-dependent** condition-driven loop | Structured on-chip loop (not unrolled) | **Runtime/data-dependent loops (NKI 0.6.0+):** use `nl.fori_loop` (counted, runtime bound) or `nl.while_loop` (condition-driven), NOT `for i in nl.dynamic_range(...)` or bare `while reg:` — @@ -213,6 +216,7 @@ For detailed code examples, anti-patterns, and production patterns (cumsum, rmsn References are tiered to minimize overhead on simple tasks. Load only what you need based on the Complexity Assessment above. ### Always load (core references): + - `references/nki-language-constraint.md` - **MANDATORY**: Required and forbidden API patterns for NKI 0.4.0, reference kernel template - `references/common-patterns.md` - Full code examples: matmul PSUM accumulation, fused ScalarE, associative scan, production patterns - `references/api-translation.md` - PyTorch/NumPy to NKI operation mapping @@ -220,19 +224,23 @@ References are tiered to minimize overhead on simple tasks. Load only what you n - `references/indexing-patterns.md` - **Complete indexing guide**: memory-type rules (HBM/SBUF/PSUM), operation constraints (matmul/transpose/reduce), dynamic indexing with DGE modes ### Load when tiling or DMA patterns are needed (medium+ complexity): + - `references/memory-patterns.md` - DMA and tiling patterns with code examples - `references/nkilib/core/tiled-range.md` - TiledRange: dimension tiling with remainder handling - `references/nkilib/core/kernel-helpers.md` - Math helpers, SPMD, dtype utilities ### Load when layout manipulation is needed: + - `references/transpose-and-layout.md` - **Transpose and layout transformation guide**: nc_transpose, TensorView, array patterns, strided DMA, decision trees for layout operations - `references/nkilib/core/tensor-view.md` - TensorView: zero-copy tensor manipulation ### API features documented in `/neuron-nki-docs` (query it, not duplicated here): + - **Native `NkiTensor` view methods** — composable, zero-copy views callable directly on a tensor: `slice`, `select`, `permute`, `broadcast`, `expand_dim`, `squeeze_dim`, `reshape_dim`, `flatten_dims`, `rearrange`, `reshape`, `view`, `vector_select`, plus query methods `is_contiguous` / `is_indirect` and low-level `ap` / `get_pattern`. These are the native tensor methods (e.g. `t.slice(1, 0, 256)`); the `TensorView` helper above wraps the same ops for composing complex 3D+ `.ap()` patterns. Look up signatures/examples via `/neuron-nki-docs` → `api-nki-tensor.md`. - **Tensor indirection on compute ops (`.indirect()`)** — on NeuronCore-v4+, gather/scatter for **on-chip compute** (not just DMA): pass a `.indirect(index)` view as `dst` or `data` to `nc_matmul`, `nc_matmul_mx`, `tensor_tensor`, `tensor_scalar`, `tensor_reduce`, `tensor_copy`, `tensor_copy_predicated`, `tensor_scalar_reduce`, `tensor_scalar_cumulative`, `activation`, `activation_reduce`, `activate2`, `exponential`. Subject to quadrant/partition-alignment rules (group size 16 for vector/scalar/gpsimd, 32 for tensor engine). This extends the DMA-only `vector_offset` indirection to compute. Look up `NkiTensor.indirect` and the per-op notes via `/neuron-nki-docs`. ### Load when advanced patterns are needed (complex kernels only): + - `references/performance-basics.md` - Optimization patterns (fusion, double buffering) - `references/nkilib/core/allocator.md` - SbufManager: stack/heap SBUF allocation - `references/nkilib/core/tile-info.md` - TiledDimInfo: tile tracking with subtile support @@ -256,21 +264,21 @@ Full source for nkilib/core utilities and subkernels: ### Always Use -| Utility | Adopt When | -|---------|-----------| -| `div_ceil(n, d)` | Any tile count computation. **Never** write `(n + d - 1) // d` inline. | -| `kernel_assert()` | Any input validation. **Never** use Python `assert`. | +| Utility | Adopt When | +| ----------------- | ---------------------------------------------------------------------- | +| `div_ceil(n, d)` | Any tile count computation. **Never** write `(n + d - 1) // d` inline. | +| `kernel_assert()` | Any input validation. **Never** use Python `assert`. | ### Use When Pattern Matches -| Utility | Adopt When | Reference | -|---------|-----------|-----------| -| `TiledRange` | Tiled dimension iteration with remainder handling | `references/nkilib/core/tiled-range.md` | -| `TensorView` | Strided/interleaved DMA, broadcasting, reshape without copy, dynamic selection | `references/nkilib/core/tensor-view.md` | -| `stream_shuffle_broadcast` | Replicate partition-0 value (bias, scale) to all 128 partitions | `references/nkilib/ops/stream-shuffle-broadcast.md` | -| `SbufManager` | 4+ SBUF tensors or sub-functions sharing SBUF | `references/nkilib/core/allocator.md` | -| `NkiTensor` view methods | Zero-copy reshape/slice/permute/broadcast directly on a tensor (`t.slice(...)`, `t.reshape_dim(...)`, etc.) | `/neuron-nki-docs` → `api-nki-tensor.md` | -| `.indirect()` on compute ops | On-chip gather/scatter (NeuronCore-v4+) for matmul/tensor/activation ops via an index tensor | `/neuron-nki-docs` → `NkiTensor.indirect` | +| Utility | Adopt When | Reference | +| ---------------------------- | ----------------------------------------------------------------------------------------------------------- | --------------------------------------------------- | +| `TiledRange` | Tiled dimension iteration with remainder handling | `references/nkilib/core/tiled-range.md` | +| `TensorView` | Strided/interleaved DMA, broadcasting, reshape without copy, dynamic selection | `references/nkilib/core/tensor-view.md` | +| `stream_shuffle_broadcast` | Replicate partition-0 value (bias, scale) to all 128 partitions | `references/nkilib/ops/stream-shuffle-broadcast.md` | +| `SbufManager` | 4+ SBUF tensors or sub-functions sharing SBUF | `references/nkilib/core/allocator.md` | +| `NkiTensor` view methods | Zero-copy reshape/slice/permute/broadcast directly on a tensor (`t.slice(...)`, `t.reshape_dim(...)`, etc.) | `/neuron-nki-docs` → `api-nki-tensor.md` | +| `.indirect()` on compute ops | On-chip gather/scatter (NeuronCore-v4+) for matmul/tensor/activation ops via an index tensor | `/neuron-nki-docs` → `NkiTensor.indirect` | **Specialized:** `TiledDimInfo` (subtile metadata), `tp_broadcast` (P→F broadcast, very rare). @@ -325,6 +333,7 @@ If the input/output tensor layout would make the kernel significantly harder to > kernel and improve performance." Layout changes that typically help: + - Putting the reduction dimension last (contiguous in memory) - Aligning dimensions to 128 (partition) and 512 (PSUM free) - Transposing to avoid strided DMA patterns @@ -338,7 +347,7 @@ in memory. Target ≥2KB contiguous free dimension to saturate memory bandwidth. hardware partitions. The free dimension F (second dim onwards) should be large and contiguous. | Data Type | Minimum Free Dimension (Contiguous) | -|-----------|-------------------------------------| +| --------- | ----------------------------------- | | float32 | 512 elements (2KB) | | bfloat16 | 1024 elements (2KB) | | float8 | 2048 elements (2KB) | @@ -346,6 +355,7 @@ hardware partitions. The free dimension F (second dim onwards) should be large a **What "contiguous" means:** The free dimension elements are adjacent in HBM memory with stride=1. **Production example** from `mlp_tkg_gate_up_projection.py:169-181`: + ```python # Weight layout: [H, I] where I is the contiguous free dimension # Load weight tile [HTile=2048, I] where I is large and contiguous @@ -374,6 +384,7 @@ If your tensor layout requires strided access, consider asking the user to chang Avoid unnecessary HBM round-trips by keeping intermediate results in SBUF between operations. **Common pattern: MatMul → Element-wise → HBM** + ```python # MatMul result in PSUM psum_result = nl.ndarray((P, F), dtype=nl.float32, buffer=nl.psum) @@ -391,6 +402,7 @@ nisa.dma_copy(dst=output_hbm, src=sbuf_result) ``` **Anti-pattern to avoid:** + ```python # BAD: Writing matmul result to HBM, then reading back for activation nisa.dma_copy(dst=hbm_temp, src=psum_result) # Unnecessary write @@ -403,6 +415,7 @@ nisa.activation(dst=sbuf_for_act, data=sbuf_for_act, op=nl.gelu) Always try to use the full partition dimension (128) for hardware parallelism. **Production example** from `mlp_tkg_constants.py:156`: + ```python # Hardware partition dimension constraint - always use 128 _pmax = nl.tile_size.pmax # Max partition dimension in SBUF = 128 @@ -417,12 +430,13 @@ tile = nl.ndarray((H0, free_dim), dtype=dtype, buffer=nl.sbuf) # [128, ...] ### 5. Minimum Tile Sizes -| Operation | Minimum Tile Size | Rationale | -|-----------|------------------|-----------| -| MatMul (nc_matmul) | (128, 512) | Partition=128, PSUM free=512 for pipelining | -| Vector/Scalar ops | (128, 64) | Partition=128, free dim ≥64 for efficiency | +| Operation | Minimum Tile Size | Rationale | +| ------------------ | ----------------- | ------------------------------------------- | +| MatMul (nc_matmul) | (128, 512) | Partition=128, PSUM free=512 for pipelining | +| Vector/Scalar ops | (128, 64) | Partition=128, free dim ≥64 for efficiency | **Production MatMul example** from `mlp_tkg_gate_up_projection.py:188-204`: + ```python # Standard matmul tile: stationary [128, T], moving [128, 512] for i_tiles in TiledRange(I, dims._psum_fmax): # _psum_fmax = 512 @@ -444,6 +458,7 @@ for i_tiles in TiledRange(I, dims._psum_fmax): # _psum_fmax = 512 ``` **Vector/Scalar tile sizing** from `mlp_tkg_constants.py:194-206`: + ```python # column_tiling_dim sets the free dimension for vector/scalar ops # (e.g., activation functions, element-wise ops after matmul) @@ -456,12 +471,11 @@ else: column_tiling_dim = 128 # Large T: use 128 ``` - ## Related Skills -| Skill | Use When | -|-------|----------| -| `/neuron-nki-docs` | Look up specific API documentation | -| `/neuron-nki-debugging` | Debug compiler errors on device | -| `/neuron-nki-profiling` | Profile kernel performance | +| Skill | Use When | +| ------------------------------ | ------------------------------------- | +| `/neuron-nki-docs` | Look up specific API documentation | +| `/neuron-nki-debugging` | Debug compiler errors on device | +| `/neuron-nki-profiling` | Profile kernel performance | | `/neuron-nki-profile-querying` | Query and analyze kernel profile data | diff --git a/skills/neuron-nki-writing/references/api-translation.md b/skills/neuron-nki-writing/references/api-translation.md index 14141a5..53f8dda 100644 --- a/skills/neuron-nki-writing/references/api-translation.md +++ b/skills/neuron-nki-writing/references/api-translation.md @@ -8,47 +8,47 @@ This reference maps common PyTorch and NumPy operations to their NKI equivalents ### Arithmetic -| PyTorch/NumPy | NKI | Notes | -|---------------|-----|-------| -| `a + b` | `nisa.tensor_tensor(dst=result, data1=a, data2=b, op=nl.add)` | Both tensors in SBUF | -| `a - b` | `nisa.tensor_tensor(dst=result, data1=a, data2=b, op=nl.subtract)` | | -| `a * b` | `nisa.tensor_tensor(dst=result, data1=a, data2=b, op=nl.multiply)` | | -| `a / b` | `nisa.tensor_tensor(dst=result, data1=a, data2=b, op=nl.divide)` | | -| `a + scalar` | `nisa.tensor_scalar(dst=result, data=a, op0=nl.add, operand0=value)` | | -| `a * scalar` | `nisa.tensor_scalar(dst=result, data=a, op0=nl.multiply, operand0=value)` | | +| PyTorch/NumPy | NKI | Notes | +| ------------- | ------------------------------------------------------------------------- | -------------------- | +| `a + b` | `nisa.tensor_tensor(dst=result, data1=a, data2=b, op=nl.add)` | Both tensors in SBUF | +| `a - b` | `nisa.tensor_tensor(dst=result, data1=a, data2=b, op=nl.subtract)` | | +| `a * b` | `nisa.tensor_tensor(dst=result, data1=a, data2=b, op=nl.multiply)` | | +| `a / b` | `nisa.tensor_tensor(dst=result, data1=a, data2=b, op=nl.divide)` | | +| `a + scalar` | `nisa.tensor_scalar(dst=result, data=a, op0=nl.add, operand0=value)` | | +| `a * scalar` | `nisa.tensor_scalar(dst=result, data=a, op0=nl.multiply, operand0=value)` | | ### Comparison -| PyTorch/NumPy | NKI | Notes | -|---------------|-----|-------| -| `a > b` | `nisa.tensor_tensor(dst=result, data1=a, data2=b, op=nl.greater)` | Returns 0 or 1 | -| `a < b` | `nisa.tensor_tensor(dst=result, data1=a, data2=b, op=nl.less)` | | -| `a >= b` | `nisa.tensor_tensor(dst=result, data1=a, data2=b, op=nl.greater_equal)` | | -| `torch.maximum(a, b)` | `nisa.tensor_tensor(dst=result, data1=a, data2=b, op=nl.maximum)` | | -| `torch.minimum(a, b)` | `nisa.tensor_tensor(dst=result, data1=a, data2=b, op=nl.minimum)` | | +| PyTorch/NumPy | NKI | Notes | +| --------------------- | ----------------------------------------------------------------------- | -------------- | +| `a > b` | `nisa.tensor_tensor(dst=result, data1=a, data2=b, op=nl.greater)` | Returns 0 or 1 | +| `a < b` | `nisa.tensor_tensor(dst=result, data1=a, data2=b, op=nl.less)` | | +| `a >= b` | `nisa.tensor_tensor(dst=result, data1=a, data2=b, op=nl.greater_equal)` | | +| `torch.maximum(a, b)` | `nisa.tensor_tensor(dst=result, data1=a, data2=b, op=nl.maximum)` | | +| `torch.minimum(a, b)` | `nisa.tensor_tensor(dst=result, data1=a, data2=b, op=nl.minimum)` | | ## Activation Functions -| PyTorch/NumPy | NKI | Notes | -|---------------|-----|-------| -| `torch.exp(x)` | `nisa.activation(dst=result, data=x, op=nl.exp)` | | -| `torch.relu(x)` | `nisa.activation(dst=result, data=x, op=nl.relu)` | | -| `torch.sigmoid(x)` | `nisa.activation(dst=result, data=x, op=nl.sigmoid)` | | -| `torch.tanh(x)` | `nisa.activation(dst=result, data=x, op=nl.tanh)` | | -| `F.gelu(x)` | `nisa.activation(dst=result, data=x, op=nl.gelu)` | | -| `torch.sqrt(x)` | `nisa.activation(dst=result, data=x, op=nl.sqrt)` | | -| `torch.rsqrt(x)` | `nisa.activation(dst=result, data=x, op=nl.rsqrt)` | 1/sqrt(x) | -| `1/x` | `nisa.reciprocal(dst=result, data=x)` | | -| `-x` | `nisa.tensor_scalar(dst=result, data=x, op0=nl.multiply, operand0=-1.0)` | | +| PyTorch/NumPy | NKI | Notes | +| ------------------ | ------------------------------------------------------------------------ | --------- | +| `torch.exp(x)` | `nisa.activation(dst=result, data=x, op=nl.exp)` | | +| `torch.relu(x)` | `nisa.activation(dst=result, data=x, op=nl.relu)` | | +| `torch.sigmoid(x)` | `nisa.activation(dst=result, data=x, op=nl.sigmoid)` | | +| `torch.tanh(x)` | `nisa.activation(dst=result, data=x, op=nl.tanh)` | | +| `F.gelu(x)` | `nisa.activation(dst=result, data=x, op=nl.gelu)` | | +| `torch.sqrt(x)` | `nisa.activation(dst=result, data=x, op=nl.sqrt)` | | +| `torch.rsqrt(x)` | `nisa.activation(dst=result, data=x, op=nl.rsqrt)` | 1/sqrt(x) | +| `1/x` | `nisa.reciprocal(dst=result, data=x)` | | +| `-x` | `nisa.tensor_scalar(dst=result, data=x, op0=nl.multiply, operand0=-1.0)` | | ## Reduction Operations -| PyTorch/NumPy | NKI | Notes | -|---------------|-----|-------| -| `torch.sum(x, dim=axis)` | `nisa.tensor_reduce(dst=result, data=x, op=nl.add, axis=axis)` | | -| `torch.max(x, dim=axis)` | `nisa.tensor_reduce(dst=result, data=x, op=nl.maximum, axis=axis)` | | -| `torch.min(x, dim=axis)` | `nisa.tensor_reduce(dst=result, data=x, op=nl.minimum, axis=axis)` | | -| `torch.mean(x, dim=axis)` | Sum then divide by count | No direct mean op | +| PyTorch/NumPy | NKI | Notes | +| ------------------------- | ------------------------------------------------------------------ | ----------------- | +| `torch.sum(x, dim=axis)` | `nisa.tensor_reduce(dst=result, data=x, op=nl.add, axis=axis)` | | +| `torch.max(x, dim=axis)` | `nisa.tensor_reduce(dst=result, data=x, op=nl.maximum, axis=axis)` | | +| `torch.min(x, dim=axis)` | `nisa.tensor_reduce(dst=result, data=x, op=nl.minimum, axis=axis)` | | +| `torch.mean(x, dim=axis)` | Sum then divide by count | No direct mean op | ### Reduction Example (Sum over axis=1) @@ -63,10 +63,10 @@ nisa.tensor_reduce(dst=result, data=x, op=nl.add, axis=1) ## Matrix Operations -| PyTorch/NumPy | NKI | Notes | -|---------------|-----|-------| -| `a @ b` | `nisa.nc_matmul(dst=psum_result, stationary=a, moving=b)` | Result in PSUM | -| `a.T` | `nisa.nc_transpose(dst=result, data=a)` | | +| PyTorch/NumPy | NKI | Notes | +| ------------- | --------------------------------------------------------- | -------------- | +| `a @ b` | `nisa.nc_matmul(dst=psum_result, stationary=a, moving=b)` | Result in PSUM | +| `a.T` | `nisa.nc_transpose(dst=result, data=a)` | | ### Matrix Multiply Pattern @@ -85,46 +85,47 @@ nisa.tensor_copy(dst=sbuf_result, src=psum_result) ## Data Type Mapping -| PyTorch | NKI | Notes | -|---------|-----|-------| -| `torch.float32` | `nl.float32` | Full precision | -| `torch.float16` | `nl.float16` | Half precision | -| `torch.bfloat16` | `nl.bfloat16` | Brain float | -| `torch.int32` | `nl.int32` | Integer | -| `torch.int8` | `nl.int8` | Quantized | +| PyTorch | NKI | Notes | +| --------------------- | ---------------- | ---------------- | +| `torch.float32` | `nl.float32` | Full precision | +| `torch.float16` | `nl.float16` | Half precision | +| `torch.bfloat16` | `nl.bfloat16` | Brain float | +| `torch.int32` | `nl.int32` | Integer | +| `torch.int8` | `nl.int8` | Quantized | | `torch.float8_e4m3fn` | `nl.float8_e4m3` | FP8 (gen3+ only) | -| `torch.float8_e5m2` | `nl.float8_e5m2` | FP8 (gen3+ only) | +| `torch.float8_e5m2` | `nl.float8_e5m2` | FP8 (gen3+ only) | ## Memory Operations -| Operation | NKI | Notes | -|-----------|-----|-------| -| Load from HBM | `nisa.dma_copy(dst=sbuf_tile, src=hbm_tensor[slice])` | | -| Store to HBM | `nisa.dma_copy(dst=hbm_tensor[slice], src=sbuf_tile)` | | -| Copy SBUF to SBUF | `nisa.tensor_copy(dst=dest, src=src)` | | -| Copy PSUM to SBUF | `nisa.tensor_copy(dst=sbuf, src=psum)` | Type conversion | -| Initialize to value | `nisa.memset(dst=tensor, value=0.0)` | | +| Operation | NKI | Notes | +| ------------------- | ----------------------------------------------------- | --------------- | +| Load from HBM | `nisa.dma_copy(dst=sbuf_tile, src=hbm_tensor[slice])` | | +| Store to HBM | `nisa.dma_copy(dst=hbm_tensor[slice], src=sbuf_tile)` | | +| Copy SBUF to SBUF | `nisa.tensor_copy(dst=dest, src=src)` | | +| Copy PSUM to SBUF | `nisa.tensor_copy(dst=sbuf, src=psum)` | Type conversion | +| Initialize to value | `nisa.memset(dst=tensor, value=0.0)` | | ## Shape Operations -| PyTorch/NumPy | NKI | Notes | -|---------------|-----|-------| -| `x.reshape(shape)` | `x.reshape(shape)` | Zero-copy reshape | -| `x.view(shape)` | `x.reshape(shape)` | Same as reshape | -| Broadcasting | Manual expansion | Explicit broadcast required | +| PyTorch/NumPy | NKI | Notes | +| ------------------ | ------------------ | --------------------------- | +| `x.reshape(shape)` | `x.reshape(shape)` | Zero-copy reshape | +| `x.view(shape)` | `x.reshape(shape)` | Same as reshape | +| Broadcasting | Manual expansion | Explicit broadcast required | ## Not Directly Supported These operations require manual implementation: - `torch.softmax()` - Implement as: exp(x - max(x)) / sum(exp(x - max(x))) -- `torch.layer_norm()` - Implement as: (x - mean) / sqrt(var + eps) * gamma + beta +- `torch.layer_norm()` - Implement as: (x - mean) / sqrt(var + eps) \* gamma + beta - `torch.gather()` - Use dynamic access patterns (use `/neuron-nki-docs dynamic access` for details) - `torch.scatter()` - Use dynamic access patterns ## Need More APIs? Use `/neuron-nki-docs ` for: + - Complete API signatures and parameters - APIs not listed here - Hardware generation support details diff --git a/skills/neuron-nki-writing/references/common-patterns.md b/skills/neuron-nki-writing/references/common-patterns.md index 3224bde..1d6c3da 100644 --- a/skills/neuron-nki-writing/references/common-patterns.md +++ b/skills/neuron-nki-writing/references/common-patterns.md @@ -5,6 +5,7 @@ Detailed code examples, anti-patterns, and production patterns for common NKI op ## Element-wise Operations Element-wise operations use VectorE/ScalarE and work on any buffer type. For basic pattern: + - Reshape to 2D for simpler tiling (collapse batch dimensions) - Tile partition dimension (≤128) and free dimension as needed - Use `nisa.activation()` for element-wise functions (exp, sigmoid, tanh) @@ -15,6 +16,7 @@ Element-wise operations use VectorE/ScalarE and work on any buffer type. For bas **PERFORMANCE REQUIREMENT**: Matrix multiplication accumulation over the K (contraction) dimension requires careful loop structure to trigger efficient hardware PSUM accumulation. Using `nl.sequential_range` will serialize execution and prevent PSUM accumulation, causing severe performance degradation. **Use the `accumulate=(k_idx > 0)` flag to control PSUM initialization.** The `accumulate` parameter of `nisa.nc_matmul` makes the overwrite-vs-accumulate behavior explicit: + - `accumulate=False` — overwrite the `dst` PSUM tile. The **first** matmul targeting a PSUM location must overwrite to initialize it. - `accumulate=True` — add the result onto existing PSUM content. Valid only after the location was initialized with `accumulate=False`. @@ -47,6 +49,7 @@ nisa.tensor_copy(dst=result_sbuf, src=result_psum) ``` **Key points:** + - **PSUM allocation**: Uninitialized `nl.ndarray(..., buffer=nl.psum)` is correct — **never `nisa.memset` the PSUM tile before the K loop**; the `accumulate=False` first write initializes it - **Loop type**: Use `nl.affine_range()` or `range()` for K-dimension loops, NEVER `nl.sequential_range` - **Accumulation mechanism**: Pass `accumulate=(k_idx > 0)` so the first K tile overwrites (initializes) PSUM and subsequent tiles accumulate. This is the explicit, preferred form; omitting `accumulate` (default `None`) lets the compiler auto-infer the same behavior, but the explicit flag documents intent and avoids ambiguity @@ -112,6 +115,7 @@ nisa.dma_copy(dst=hbm_output, src=sbuf_temp) ``` **Why this matters**: + - `accumulate=(k_idx > 0)` makes PSUM initialization explicit: the first matmul overwrites (no separate memset), later matmuls accumulate - `nl.sequential_range` forces serialization, preventing efficient pipelining of the accumulation group - Hardware PSUM accumulation is performed in FP32 with very low overhead @@ -124,6 +128,7 @@ nisa.dma_copy(dst=hbm_output, src=sbuf_temp) ScalarE supports "pipelined multiply-add" before applying non-linear functions, allowing two operations at the cost of one. This is useful when translating PyTorch operations like `torch.exp(x * scale)` or `torch.sigmoid(x + bias)`. **Pattern from Mamba (combine multiplication and exponential):** + ```python # PyTorch: torch.exp(delta * A) # @@ -144,6 +149,7 @@ nisa.activation(dst=deltaA, op=nl.exp, data=temp) ``` **Available fused patterns:** + - `nisa.activation(op=nl.exp, data=x, scale=s)` → `exp(x * s)` - `nisa.activation(op=nl.sigmoid, data=x, bias=b)` → `sigmoid(x + b)` - `nisa.activation(op=nl.tanh, data=x, scale=s, bias=b)` → `tanh(x * s + b)` @@ -155,6 +161,7 @@ nisa.activation(dst=deltaA, op=nl.exp, data=temp) For operations with loop-carried dependencies (e.g., cumulative sum, RNN, state space models), use `nisa.tensor_tensor_scan` instead of explicit sequential loops. **PyTorch pattern:** + ```python # torch.cumsum equivalent, or RNN-like: out[i] = f(out[i-1], x[i]) out = torch.empty_like(x) @@ -166,6 +173,7 @@ for i in range(seq_len): **NKI translation using associative scan:** See `examples/associative_scan.py` for complete pattern demonstrating: + - Single-instruction sequential operations (no explicit loops) - Internal caching of intermediate scan results in VectorE - Initial state handling for multi-tile sequences @@ -194,12 +202,14 @@ for i in nl.sequential_range(seq_len - 1): ``` **Common use cases:** + - `torch.cumsum(x)` → `tensor_tensor_scan(ones, x, initial=0, op0=nl.multiply, op1=nl.add)` - `torch.cumprod(x)` → `tensor_tensor_scan(x, zeros, initial=1, op0=nl.multiply, op1=nl.add)` - RNN cell: `h[t] = tanh(W_h @ h[t-1] + W_x @ x[t])` → Use scan with matmul in loop body - State space models (Mamba, S4) → Associative scan over sequence dimension **Multi-tile sequences with loop-carried dependencies:** + ```python scan_init = nl.zeros((channels, 1), dtype=deltaA.dtype, buffer=nl.sbuf) diff --git a/skills/neuron-nki-writing/references/indexing-patterns.md b/skills/neuron-nki-writing/references/indexing-patterns.md index 6d2c974..47bd81a 100644 --- a/skills/neuron-nki-writing/references/indexing-patterns.md +++ b/skills/neuron-nki-writing/references/indexing-patterns.md @@ -4,13 +4,13 @@ Production-proven indexing patterns. Use these patterns in order of preference. ## Pattern Summary -| Pattern | When to Use | Reference | -|---------|-------------|-----------| -| `tensor[start:end, :]` | Contiguous access, tiling | Examples below | -| `TensorView.slice(step=N)` | Strided/interleaved | [tensor-view.md](nkilib/core/tensor-view.md) | -| `.ap(pattern=...)` | Complex layouts, dynamic | Examples below | -| `nl.ds(start, size)` | Runtime-computed bounds | Examples below | -| **`nl.mgrid` - NOT USED** | N/A - avoid | 0 occurrences in production | +| Pattern | When to Use | Reference | +| -------------------------- | ------------------------- | -------------------------------------------- | +| `tensor[start:end, :]` | Contiguous access, tiling | Examples below | +| `TensorView.slice(step=N)` | Strided/interleaved | [tensor-view.md](nkilib/core/tensor-view.md) | +| `.ap(pattern=...)` | Complex layouts, dynamic | Examples below | +| `nl.ds(start, size)` | Runtime-computed bounds | Examples below | +| **`nl.mgrid` - NOT USED** | N/A - avoid | 0 occurrences in production | ## Memory Type Indexing Rules @@ -32,6 +32,7 @@ every_other = input_hbm[::2, :, :] ``` **HBM characteristics:** + - No partition dimension restriction - Any dimension ordering allowed - Strided access permitted (DMA handles it) @@ -59,6 +60,7 @@ nisa.dma_copy( ``` **SBUF constraints:** + - Partition dimension (shape[0]): **≤ 128** - Free dimension (shape[1:]): **≤ 32767** - First dimension IS the partition dimension @@ -91,6 +93,7 @@ nisa.nc_transpose(dst=result_sb, data=psum_tile) ``` **PSUM constraints:** + - Partition dimension (shape[0]): **≤ 128** - Free dimension (shape[1]): **512 (gen2/3); gen4: 4096 for fp32 dst, 8192 for bf16 dst** (generation- and dtype-gated; authoritative: `nc_matmul` doc block via `/neuron-nki-docs`) - Used primarily for matrix multiply accumulation @@ -98,15 +101,15 @@ nisa.nc_transpose(dst=result_sb, data=psum_tile) ### Memory Type Comparison Table -| Feature | HBM | SBUF | PSUM | -|---------|-----|------|------| -| Max partition (P) | Unlimited | 128 | 128 | -| Max free (F) | Unlimited | 32767 | 512 (gen2/3); gen4: 4096 fp32 / 8192 bf16 | -| Direct DMA to HBM | - | Yes | No | -| MatMul destination | No | No | Yes | -| General compute | No | Yes | No | -| Strided access | Yes | Via `.ap()` | No | -| Multi-dimensional | Yes (N-D) | 2D logical | 2D only | +| Feature | HBM | SBUF | PSUM | +| ------------------ | --------- | ----------- | ----------------------------------------- | +| Max partition (P) | Unlimited | 128 | 128 | +| Max free (F) | Unlimited | 32767 | 512 (gen2/3); gen4: 4096 fp32 / 8192 bf16 | +| Direct DMA to HBM | - | Yes | No | +| MatMul destination | No | No | Yes | +| General compute | No | Yes | No | +| Strided access | Yes | Via `.ap()` | No | +| Multi-dimensional | Yes (N-D) | 2D logical | 2D only | ### Memory Type Index Examples @@ -142,6 +145,7 @@ The partition dimension is the most constrained aspect of NKI programming. These ### What Cannot Be Done with Partitions **1. Reshape partition dimension:** + ```python # INVALID: Cannot reshape 128x512 to 64x1024 (changes partition count) sbuf_tile = nl.ndarray((128, 512), dtype=nl.float32, buffer=nl.sbuf) @@ -153,6 +157,7 @@ sbuf_3d = nl.ndarray((128, 32, 16), dtype=nl.float32, buffer=nl.sbuf) ``` **2. Strided partition access:** + ```python # INVALID: Cannot stride across partitions # tile = sbuf[::2, :] # Error: strided partition access @@ -163,6 +168,7 @@ tile = sbuf[64:128, :] # Last 64 partitions ``` **3. Flatten partition with free dims:** + ```python # INVALID: Cannot flatten partition into free # flat = sbuf.flatten() # Error! @@ -173,6 +179,7 @@ for p in nl.affine_range(num_p_tiles): ``` **4. Negative strides on partition:** + ```python # INVALID: Cannot reverse partition order # reversed_p = sbuf[::-1, :] # Error! @@ -183,6 +190,7 @@ for p in range(num_tiles - 1, -1, -1): ``` **5. Transpose without explicit instruction:** + ```python # INVALID: Cannot transpose via indexing # transposed = sbuf.T # Error! @@ -195,6 +203,7 @@ nisa.nc_transpose(dst=transposed, data=sbuf) ### Valid Partition Operations **Contiguous slicing:** + ```python # Full partition range full = sbuf[0:128, :] @@ -207,6 +216,7 @@ single = sbuf[0:1, :] ``` **Dynamic offset with nl.ds:** + ```python # Dynamic partition offset (must still be contiguous) for i in nl.affine_range(num_tiles): @@ -216,12 +226,12 @@ for i in nl.affine_range(num_tiles): ### Common Partition Errors and Fixes -| Error | Cause | Fix | -|-------|-------|-----| -| `Partition dimension exceeds 128` | Shape[0] > 128 | Tile with outer loop | -| `Strided partition access` | Using `::stride` on dim 0 | Use contiguous slice + loop | -| `Cannot reshape partition` | Reshape changes dim 0 | Reshape only free dimensions | -| `Invalid transpose` | Using `.T` or `np.transpose` | Use `nisa.nc_transpose()` | +| Error | Cause | Fix | +| --------------------------------- | ---------------------------- | ---------------------------- | +| `Partition dimension exceeds 128` | Shape[0] > 128 | Tile with outer loop | +| `Strided partition access` | Using `::stride` on dim 0 | Use contiguous slice + loop | +| `Cannot reshape partition` | Reshape changes dim 0 | Reshape only free dimensions | +| `Invalid transpose` | Using `.T` or `np.transpose` | Use `nisa.nc_transpose()` | ```python # Fix: Partition exceeds 128 @@ -264,6 +274,7 @@ tile = tensor[:, f_start:f_end] ``` **Key points:** + - Works for contiguous memory regions - Use `min()` at boundaries, NOT deprecated `mask=` parameter - Variables in slices are resolved at compile time @@ -305,6 +316,7 @@ nisa.tensor_tensor( ``` **TensorView operations:** + - `.slice(dim, start, end, step)` - Slice with optional stride - `.expand_dim(dim)` - Add dimension of size 1 - `.broadcast(dim, size)` - Broadcast size-1 dim @@ -332,6 +344,7 @@ tensor.ap( ``` **Pattern structure:** + - Each `[stride, size]` pair describes one dimension - `stride` = elements to skip between consecutive indices - `size` = number of elements in this dimension @@ -362,6 +375,7 @@ nisa.tensor_copy( ``` **Key points:** + - `nl.ds(start, size)` not `nl.ds(start, end)` - Works in both source and destination positions - Commonly used with loop-computed offsets @@ -394,6 +408,7 @@ nisa.nc_matmul(dst=c_psum, stationary=a_sb, moving=b_sb) ``` **Stationary vs Moving operand:** + ```python # Stationary operand: held in place, free dim ≤ 128 # Moving operand: streamed through, free dim ≤ nl.tile_size.gemm_moving_fmax @@ -437,11 +452,11 @@ DMA operations support three addressing modes for different use cases. **DGE (Data Gather Engine) Modes:** -| Mode | Parameter | When to Use | Performance | -|------|-----------|-------------|-------------| -| None | (default) | Compile-time known indices | Fastest | -| SWDGE | `dge_mode=dge_mode.swdge` | Loop-variable indices, small iteration count | Medium | -| HWDGE | `dge_mode=dge_mode.hwdge` | Runtime-computed indices, large iteration count | Flexible | +| Mode | Parameter | When to Use | Performance | +| ----- | ------------------------- | ----------------------------------------------- | ----------- | +| None | (default) | Compile-time known indices | Fastest | +| SWDGE | `dge_mode=dge_mode.swdge` | Loop-variable indices, small iteration count | Medium | +| HWDGE | `dge_mode=dge_mode.hwdge` | Runtime-computed indices, large iteration count | Flexible | ```python # Required import for DGE modes @@ -472,6 +487,7 @@ nisa.dma_copy( ``` **dma_transpose specific:** + ```python # dma_transpose: loads with transpose in single operation # More efficient than dma_copy + nc_transpose for HBM→SBUF @@ -510,6 +526,7 @@ nisa.tensor_reduce( ``` **Multi-axis reduction:** + ```python # For 3D tensor, can reduce multiple free axes input_3d = nl.ndarray((128, 64, 32), dtype=nl.float32, buffer=nl.sbuf) @@ -525,16 +542,16 @@ nisa.tensor_reduce(dst=final, data=partial, op=nl.add, axis=1) ### Constraint Quick Reference Table -| Operation | Constraint | Limit | Notes | -|-----------|------------|-------|-------| -| `nc_matmul` | K (contraction) | ≤ 2048 | Tile K for larger | -| `nc_matmul` | M (lhs partition) | ≤ 128 | Standard partition limit | -| `nc_matmul` | N (dst free) | 512 (gen2/3); gen4: 4096 fp32 / 8192 bf16 | Tile N for larger | -| `nc_transpose` | Tile size | ≤ 128×128 | Tile both dims | -| `tensor_reduce` | Axis | ≥ 1 (free only) | Cannot reduce partition | -| `dma_copy` | Partition | ≤ 128 | Standard | -| `dma_copy` | Free | ≤ 32767 | SBUF limit | -| Any PSUM op | Free | 512 (gen2/3); gen4: 4096 fp32 / 8192 bf16 | PSUM limit | +| Operation | Constraint | Limit | Notes | +| --------------- | ----------------- | ----------------------------------------- | ------------------------ | +| `nc_matmul` | K (contraction) | ≤ 2048 | Tile K for larger | +| `nc_matmul` | M (lhs partition) | ≤ 128 | Standard partition limit | +| `nc_matmul` | N (dst free) | 512 (gen2/3); gen4: 4096 fp32 / 8192 bf16 | Tile N for larger | +| `nc_transpose` | Tile size | ≤ 128×128 | Tile both dims | +| `tensor_reduce` | Axis | ≥ 1 (free only) | Cannot reduce partition | +| `dma_copy` | Partition | ≤ 128 | Standard | +| `dma_copy` | Free | ≤ 32767 | SBUF limit | +| Any PSUM op | Free | 512 (gen2/3); gen4: 4096 fp32 / 8192 bf16 | PSUM limit | ## Dynamic Indexing Deep Dive @@ -542,17 +559,18 @@ Understanding when indices are resolved is critical for correct kernel design. ### Static vs Dynamic Index Resolution -| Index Type | Resolution Time | Example | Use Case | -|------------|-----------------|---------|----------| -| Literal | Compile | `tensor[0:128, 0:256]` | Fixed-size tiles | -| Python variable | Compile | `tensor[0:p_size, 0:f_size]` | Variable from Python scope | -| `nl.affine_range` var | Compile (unrolled) | `tensor[i*128:(i+1)*128, :]` | Parallel tiling | -| `nl.ds()` | Runtime | `tensor[:, nl.ds(offset, size)]` | Dynamic bounds | -| `nl.fori_loop` / `nl.while_loop` body var | Runtime | Structured on-chip loop | Data-dependent iteration (NKI 0.6.0+; replaces legacy `nl.dynamic_range` / bare `while reg:`) | +| Index Type | Resolution Time | Example | Use Case | +| ----------------------------------------- | ------------------ | -------------------------------- | --------------------------------------------------------------------------------------------- | +| Literal | Compile | `tensor[0:128, 0:256]` | Fixed-size tiles | +| Python variable | Compile | `tensor[0:p_size, 0:f_size]` | Variable from Python scope | +| `nl.affine_range` var | Compile (unrolled) | `tensor[i*128:(i+1)*128, :]` | Parallel tiling | +| `nl.ds()` | Runtime | `tensor[:, nl.ds(offset, size)]` | Dynamic bounds | +| `nl.fori_loop` / `nl.while_loop` body var | Runtime | Structured on-chip loop | Data-dependent iteration (NKI 0.6.0+; replaces legacy `nl.dynamic_range` / bare `while reg:`) | ### nl.ds() Usage Patterns **Basic usage:** + ```python # nl.ds(start, size) creates a dynamic slice # The size must be a compile-time constant @@ -564,6 +582,7 @@ tile = tensor[:, nl.ds(offset, size)] ``` **In nested loops:** + ```python for p_idx in nl.affine_range(num_p_tiles): for f_idx in nl.affine_range(num_f_tiles): @@ -575,6 +594,7 @@ for p_idx in nl.affine_range(num_p_tiles): ``` **With conditional sizing:** + ```python # Handle edge tiles with min() for i in nl.affine_range(num_tiles): @@ -617,14 +637,15 @@ Index pattern? ### Performance Implications -| Mode | Overhead | When Optimal | -|------|----------|--------------| -| No DGE | Lowest | Fixed access patterns | -| SWDGE | Low | Small, unrolled loops (< 16 iterations) | -| HWDGE | Medium | Large loops, runtime indices | -| `.ap()` indirect | Higher | True gather/scatter | +| Mode | Overhead | When Optimal | +| ---------------- | -------- | --------------------------------------- | +| No DGE | Lowest | Fixed access patterns | +| SWDGE | Low | Small, unrolled loops (< 16 iterations) | +| HWDGE | Medium | Large loops, runtime indices | +| `.ap()` indirect | Higher | True gather/scatter | **Example: Choosing between SWDGE and HWDGE:** + ```python # SWDGE: better for small, unrolled loops # Loop is unrolled, each iteration becomes separate instruction @@ -648,11 +669,11 @@ nl.fori_loop(0, runtime_count, body) ## Compile-Time vs Runtime -| Evaluated At | Constructs | Notes | -|--------------|------------|-------| -| Compile-time | `range()`, `tensor.shape`, `print()`, slice literals | Loop unrolled, values baked in | -| Unrolled at compile | `nl.affine_range()`, `nl.sequential_range()` | Loop body replicated N times | -| Runtime (on-device) | `nl.fori_loop()`, `nl.while_loop()`, registers | Structured on-chip iteration (replaces legacy `nl.dynamic_range()` / bare `while reg:`) | +| Evaluated At | Constructs | Notes | +| ------------------- | ---------------------------------------------------- | --------------------------------------------------------------------------------------- | +| Compile-time | `range()`, `tensor.shape`, `print()`, slice literals | Loop unrolled, values baked in | +| Unrolled at compile | `nl.affine_range()`, `nl.sequential_range()` | Loop body replicated N times | +| Runtime (on-device) | `nl.fori_loop()`, `nl.while_loop()`, registers | Structured on-chip iteration (replaces legacy `nl.dynamic_range()` / bare `while reg:`) | ```python # Compile-time: shape known, loop unrolled @@ -674,13 +695,13 @@ nl.fori_loop(0, runtime_count, body) ## Common Mistakes to Avoid -| Mistake | Why It's Wrong | Correct Approach | -|---------|---------------|------------------| -| Using `nl.mgrid[]` | Not used in production (0 occurrences) | Use slicing or `.ap()` | -| Using `nl.load()`/`nl.store()` | Deprecated in Beta 2 | Use `nisa.dma_copy()` | -| Using `nl.arange()` | Deprecated | Use slicing or `nl.ds()` | -| Using `mask=` for bounds | Deprecated | Use `min()` for edge cases | -| Compile-time `print()` confusion | `print()` runs at compile time | Use `nl.device_print()` for runtime | +| Mistake | Why It's Wrong | Correct Approach | +| -------------------------------- | -------------------------------------- | ----------------------------------- | +| Using `nl.mgrid[]` | Not used in production (0 occurrences) | Use slicing or `.ap()` | +| Using `nl.load()`/`nl.store()` | Deprecated in Beta 2 | Use `nisa.dma_copy()` | +| Using `nl.arange()` | Deprecated | Use slicing or `nl.ds()` | +| Using `mask=` for bounds | Deprecated | Use `min()` for edge cases | +| Compile-time `print()` confusion | `print()` runs at compile time | Use `nl.device_print()` for runtime | ## Decision Tree @@ -779,9 +800,9 @@ Index type in DMA? ## Further Reading -| Pattern | Self-Contained Reference | -|---------|------------------------| -| TiledRange for tiling loops | [tiled-range.md](nkilib/core/tiled-range.md) | -| TensorView strided access | [tensor-view.md](nkilib/core/tensor-view.md) | -| Layout conversion (.ap()) | [layout-conversion.md](nkilib/patterns/layout-conversion.md) | -| div_ceil, dtype helpers | [kernel-helpers.md](nkilib/core/kernel-helpers.md) | +| Pattern | Self-Contained Reference | +| --------------------------- | ------------------------------------------------------------ | +| TiledRange for tiling loops | [tiled-range.md](nkilib/core/tiled-range.md) | +| TensorView strided access | [tensor-view.md](nkilib/core/tensor-view.md) | +| Layout conversion (.ap()) | [layout-conversion.md](nkilib/patterns/layout-conversion.md) | +| div_ceil, dtype helpers | [kernel-helpers.md](nkilib/core/kernel-helpers.md) | diff --git a/skills/neuron-nki-writing/references/memory-patterns.md b/skills/neuron-nki-writing/references/memory-patterns.md index bf358ad..f023ef4 100644 --- a/skills/neuron-nki-writing/references/memory-patterns.md +++ b/skills/neuron-nki-writing/references/memory-patterns.md @@ -4,11 +4,11 @@ This reference covers DMA patterns and tiling strategies for NKI kernels. ## Buffer Types -| Buffer | Syntax | Max P | Max F | Use Case | -|--------|--------|-------|-------|----------| -| SBUF | `buffer=nl.sbuf` | 128 | 32767 | General compute storage | -| PSUM | `buffer=nl.psum` | 128 | 512 (gen2/3); gen4: 4096 fp32 / 8192 bf16 | MatMul accumulation | -| HBM | `buffer=nl.shared_hbm` | - | - | Input/output tensors | +| Buffer | Syntax | Max P | Max F | Use Case | +| ------ | ---------------------- | ----- | ----------------------------------------- | ----------------------- | +| SBUF | `buffer=nl.sbuf` | 128 | 32767 | General compute storage | +| PSUM | `buffer=nl.psum` | 128 | 512 (gen2/3); gen4: 4096 fp32 / 8192 bf16 | MatMul accumulation | +| HBM | `buffer=nl.shared_hbm` | - | - | Input/output tensors | ## Basic Contiguous DMA @@ -34,6 +34,7 @@ nisa.dma_copy( ``` **Key points:** + - Source and destination slices must have matching shapes - Use `min()` to handle edge cases: `f_end = min(f_start + F_TILE_SIZE, total_f)` - Contiguous access maximizes DMA bandwidth @@ -67,6 +68,7 @@ nisa.dma_copy( ``` **Key points:** + - Use `TensorView` helper for strided access patterns - Less efficient than contiguous DMA - use when necessary - Useful for interleaved data layouts @@ -103,6 +105,7 @@ for p_tile in TiledRange(outer_dim, P_MAX): ``` **TiledRange attributes:** + - `p_tile.start_offset` - Starting index in partition dimension - `p_tile.size` - Size of current tile (handles edge cases) - `p_tile.end_offset` - End index (start_offset + size) diff --git a/skills/neuron-nki-writing/references/nki-language-constraint.md b/skills/neuron-nki-writing/references/nki-language-constraint.md index 18c3d3f..2119b42 100644 --- a/skills/neuron-nki-writing/references/nki-language-constraint.md +++ b/skills/neuron-nki-writing/references/nki-language-constraint.md @@ -31,71 +31,72 @@ def example_kernel(input_tensor): ## Hard Rules — Beta 1 → Beta 2 (Violating ANY is a compilation failure) -| NEVER use (Beta 1) | ALWAYS use (Beta 2+) | -|---|---| -| `import neuronxcc.nki` | `import nki` | -| `nl.load(tensor[...])` | `nisa.dma_copy(dst=sbuf_tile, src=tensor[0:128, 0:512])` | -| `nl.store(tensor[...], value=x)` | `nisa.dma_copy(dst=tensor[0:128, 0:512], src=sbuf_tile)` | -| `result = nisa.func(...)` | `nisa.func(dst=result, ...)` | -| `nl.mgrid[...]` or `nl.arange(...)` | `tensor[0:128, 0:512]` or `nl.ds(offset, size)` | -| `mask=` on any ISA call | `min()` for boundary clamping | -| `nl.max` / `nl.min` | `nl.maximum` / `nl.minimum` | -| `nisa.activation(op=nl.reciprocal)` | `nisa.reciprocal(dst=..., data=...)` | -| `nisa.activation(op=nl.rsqrt)` | `nisa.rsqrt(dst=..., data=...)` | -| `negate=`, `reverse0=`, `dtype=` in reduce | Remove these parameters | -| `np.float32`, `np.add` inside kernels | `nl.float32`, `nl.add` | -| `@nki.jit` on sub-functions | Remove decorator from helpers | +| NEVER use (Beta 1) | ALWAYS use (Beta 2+) | +| ------------------------------------------ | -------------------------------------------------------- | +| `import neuronxcc.nki` | `import nki` | +| `nl.load(tensor[...])` | `nisa.dma_copy(dst=sbuf_tile, src=tensor[0:128, 0:512])` | +| `nl.store(tensor[...], value=x)` | `nisa.dma_copy(dst=tensor[0:128, 0:512], src=sbuf_tile)` | +| `result = nisa.func(...)` | `nisa.func(dst=result, ...)` | +| `nl.mgrid[...]` or `nl.arange(...)` | `tensor[0:128, 0:512]` or `nl.ds(offset, size)` | +| `mask=` on any ISA call | `min()` for boundary clamping | +| `nl.max` / `nl.min` | `nl.maximum` / `nl.minimum` | +| `nisa.activation(op=nl.reciprocal)` | `nisa.reciprocal(dst=..., data=...)` | +| `nisa.activation(op=nl.rsqrt)` | `nisa.rsqrt(dst=..., data=...)` | +| `negate=`, `reverse0=`, `dtype=` in reduce | Remove these parameters | +| `np.float32`, `np.add` inside kernels | `nl.float32`, `nl.add` | +| `@nki.jit` on sub-functions | Remove decorator from helpers | **Mutable tensor annotations:** Use `import neuronxcc.nki.typing as nt` ONLY for annotating mutable output tensors in function signatures (caller allocates, kernel writes). ## Hard Rules — Beta 2 → NKI 0.3.0 (Violating ANY is a compilation failure) -| NEVER use (Beta 2) | ALWAYS use (NKI 0.3.0) | -|---|---| -| `@nki.jit(platform_target=...)` | Set `NEURON_PLATFORM_TARGET_OVERRIDE` env var instead | -| `@nki.jit(mode=...)` | Remove `mode=`; compiler auto-detects framework from arguments | -| `nisa.dma_copy(dst=hbm, src=psum)` | Copy PSUM→SBUF with `nisa.tensor_copy` first, then `nisa.dma_copy` from SBUF | -| `nisa.dma_copy(..., dst_rmw_op=...)` | Use `nisa.dma_compute(dst, srcs=[...], reduce_op=...)` | -| `nisa.dma_copy(..., unique_indices=...)` | Move `unique_indices` to `nisa.dma_compute(...)` | -| `buffer='sbuf'`, `buffer='psum'`, `buffer='hbm'` | Use objects: `buffer=nl.sbuf`, `buffer=nl.psum`, `buffer=nl.hbm` | -| `dge_mode=2` (integer enum constants) | Use named enums: `dge_mode=nisa.dge_mode.hwdge` | -| `buffer=nl.hbm` for kernel output tensors | `buffer=nl.shared_hbm` for all output tensors | -| `nisa.register_move(dst, imm=42)` | `src = nisa.register_alloc(x=42)` then `nisa.register_move(dst, src=src)` | -| `nisa.sendrecv(..., use_gpsimd_dma=True)` | `nisa.sendrecv(..., dma_engine=nisa.dma_engine.gpsimd_dma)` | -| deprecated dynamic-source tensor copy API | `nisa.tensor_copy()` with `.ap()` and `scalar_offset` | -| deprecated dynamic-destination tensor copy API | `nisa.tensor_copy()` with `.ap()` and `scalar_offset` | -| `nisa.memset(dst=int_buf, value=2.0)` | `nisa.memset(dst=int_buf, value=2)` — value dtype must match dst dtype | -| `def kernel(X, *, flag=True):` | `def kernel(X, flag=True):` — no `*` keyword-only separator | -| identity operators in conditionals (`is` / `is not`) | equality operators `==` / `!=` — e.g. `if x != None:` instead of the identity form | -| list-typed default arguments for kernel collection params | tuple defaults — e.g. `stride=(1, 1)` — use tuples, not lists, for kernel arguments | -| `num_channels=N` in collectives | Use `channel_ids=[0, 1, ...]` list in `collective_permute_implicit` | -| `nisa.dma_copy(dst=f4, src=ui16, dge_mode=hwdge)` (mismatched types) | Use `.view()` to match types: `src=src.view(nl.float4_e2m1fn_x4)` | -| `nisa.tensor_reduce(..., axis=1)` on 3D/4D tensors (wrong axis) | Use correct axis for actual tensor dimension (Beta 2 axis handling was buggy) | -| `nisa.dma_compute(dst, srcs, scales, reduce_op)` (Beta 2 order) | `nisa.dma_compute(dst, srcs, reduce_op, scales=None, unique_indices=True)` | -| `nisa.affine_select(dst, pattern, offset, ch_mul, ...)` (positional offset) | `nisa.affine_select(dst, pattern, ch_mul, on_true, on_false, offset=offset)` | +| NEVER use (Beta 2) | ALWAYS use (NKI 0.3.0) | +| --------------------------------------------------------------------------- | ----------------------------------------------------------------------------------- | +| `@nki.jit(platform_target=...)` | Set `NEURON_PLATFORM_TARGET_OVERRIDE` env var instead | +| `@nki.jit(mode=...)` | Remove `mode=`; compiler auto-detects framework from arguments | +| `nisa.dma_copy(dst=hbm, src=psum)` | Copy PSUM→SBUF with `nisa.tensor_copy` first, then `nisa.dma_copy` from SBUF | +| `nisa.dma_copy(..., dst_rmw_op=...)` | Use `nisa.dma_compute(dst, srcs=[...], reduce_op=...)` | +| `nisa.dma_copy(..., unique_indices=...)` | Move `unique_indices` to `nisa.dma_compute(...)` | +| `buffer='sbuf'`, `buffer='psum'`, `buffer='hbm'` | Use objects: `buffer=nl.sbuf`, `buffer=nl.psum`, `buffer=nl.hbm` | +| `dge_mode=2` (integer enum constants) | Use named enums: `dge_mode=nisa.dge_mode.hwdge` | +| `buffer=nl.hbm` for kernel output tensors | `buffer=nl.shared_hbm` for all output tensors | +| `nisa.register_move(dst, imm=42)` | `src = nisa.register_alloc(x=42)` then `nisa.register_move(dst, src=src)` | +| `nisa.sendrecv(..., use_gpsimd_dma=True)` | `nisa.sendrecv(..., dma_engine=nisa.dma_engine.gpsimd_dma)` | +| deprecated dynamic-source tensor copy API | `nisa.tensor_copy()` with `.ap()` and `scalar_offset` | +| deprecated dynamic-destination tensor copy API | `nisa.tensor_copy()` with `.ap()` and `scalar_offset` | +| `nisa.memset(dst=int_buf, value=2.0)` | `nisa.memset(dst=int_buf, value=2)` — value dtype must match dst dtype | +| `def kernel(X, *, flag=True):` | `def kernel(X, flag=True):` — no `*` keyword-only separator | +| identity operators in conditionals (`is` / `is not`) | equality operators `==` / `!=` — e.g. `if x != None:` instead of the identity form | +| list-typed default arguments for kernel collection params | tuple defaults — e.g. `stride=(1, 1)` — use tuples, not lists, for kernel arguments | +| `num_channels=N` in collectives | Use `channel_ids=[0, 1, ...]` list in `collective_permute_implicit` | +| `nisa.dma_copy(dst=f4, src=ui16, dge_mode=hwdge)` (mismatched types) | Use `.view()` to match types: `src=src.view(nl.float4_e2m1fn_x4)` | +| `nisa.tensor_reduce(..., axis=1)` on 3D/4D tensors (wrong axis) | Use correct axis for actual tensor dimension (Beta 2 axis handling was buggy) | +| `nisa.dma_compute(dst, srcs, scales, reduce_op)` (Beta 2 order) | `nisa.dma_compute(dst, srcs, reduce_op, scales=None, unique_indices=True)` | +| `nisa.affine_select(dst, pattern, offset, ch_mul, ...)` (positional offset) | `nisa.affine_select(dst, pattern, ch_mul, on_true, on_false, offset=offset)` | ## Hard Rules — NKI 0.3.0 → NKI 0.4.0 (Violating ANY is a compilation failure) -| NEVER use (NKI 0.3.0) | ALWAYS use (NKI 0.4.0) | -|---|---| -| `nisa.dma_transpose` with mismatched src/dst ranks | `dst.shape` must match transposed `src.shape` exactly including rank | -| `nisa.tensor_copy_dynamic_src(...)` | Removed. Use `nisa.tensor_copy()` with `.ap()` and `scalar_offset` | -| `nisa.tensor_copy_dynamic_dst(...)` | Removed. Use `nisa.tensor_copy()` with `.ap()` and `scalar_offset` | -| `import neuronxcc.nki` inside kernels | Now a **compilation error** (was warning). Use `import nki` | -| `nl.tile_size.total_available_sbuf_size` | Deprecated. Use `nl.tile_size.sbuf_fmax_bytes` (per-partition) or `nl.tile_size.sbuf_size_bytes` (total) | +| NEVER use (NKI 0.3.0) | ALWAYS use (NKI 0.4.0) | +| -------------------------------------------------- | -------------------------------------------------------------------------------------------------------- | +| `nisa.dma_transpose` with mismatched src/dst ranks | `dst.shape` must match transposed `src.shape` exactly including rank | +| `nisa.tensor_copy_dynamic_src(...)` | Removed. Use `nisa.tensor_copy()` with `.ap()` and `scalar_offset` | +| `nisa.tensor_copy_dynamic_dst(...)` | Removed. Use `nisa.tensor_copy()` with `.ap()` and `scalar_offset` | +| `import neuronxcc.nki` inside kernels | Now a **compilation error** (was warning). Use `import nki` | +| `nl.tile_size.total_available_sbuf_size` | Deprecated. Use `nl.tile_size.sbuf_fmax_bytes` (per-partition) or `nl.tile_size.sbuf_size_bytes` (total) | ### NKI 0.4.0 tile_size Bytes-Aware Constants New properties on `nl.tile_size` for SBUF/PSUM capacity checks: -| Constant | Description | -|----------|-------------| -| `nl.tile_size.sbuf_size_bytes` | Total SBUF capacity across all 128 partitions, in bytes | -| `nl.tile_size.sbuf_fmax` | Per-partition usable SBUF free dimension in FP32 elements | -| `nl.tile_size.sbuf_fmax_bytes` | Per-partition usable SBUF free dimension in bytes | -| `nl.tile_size.psum_fmax_bytes` | PSUM bank size in bytes | +| Constant | Description | +| ------------------------------ | --------------------------------------------------------- | +| `nl.tile_size.sbuf_size_bytes` | Total SBUF capacity across all 128 partitions, in bytes | +| `nl.tile_size.sbuf_fmax` | Per-partition usable SBUF free dimension in FP32 elements | +| `nl.tile_size.sbuf_fmax_bytes` | Per-partition usable SBUF free dimension in bytes | +| `nl.tile_size.psum_fmax_bytes` | PSUM bank size in bytes | **CORRECT — use bytes-aware constants for capacity checks:** + ```python assert F * 4 <= nl.tile_size.sbuf_fmax_bytes # Check per-partition SBUF capacity in bytes assert F <= nl.tile_size.sbuf_fmax # Check per-partition SBUF capacity in elements @@ -106,6 +107,7 @@ assert F <= nl.tile_size.sbuf_fmax # Check per-partition SBUF capacit `nisa.dma_transpose` now enforces that `dst.shape` rank matches the transposed `src.shape` rank exactly. **FORBIDDEN — do NOT generate this code:** + ```python # FORBIDDEN: 3D dst with 4D src — rank mismatch src_4d = nl.ndarray((128, 1, 1, 4096), dtype=nl.float32, buffer=nl.sbuf) @@ -114,6 +116,7 @@ nisa.dma_transpose(dst=dst_3d, src=src_4d, axes=(3, 1, 2, 0)) # FORBIDDEN: rank ``` **CORRECT — match dst rank to src rank:** + ```python src_4d = nl.ndarray((128, 1, 1, 4096), dtype=nl.float32, buffer=nl.sbuf) dst_4d = nl.ndarray((4096, 1, 1, 128), dtype=nl.float32, buffer=nl.sbuf) @@ -125,11 +128,13 @@ nisa.dma_transpose(dst=dst_4d, src=src_4d, axes=(3, 1, 2, 0)) PSUM cannot be directly DMA-copied to HBM. Always copy through SBUF first. **FORBIDDEN — do NOT generate this code:** + ```python nisa.dma_copy(dst=hbm_tensor, src=psum_tensor[0:TILE, 0:N]) # FORBIDDEN: direct PSUM→HBM ``` **CORRECT — always use this pattern:** + ```python sbuf_temp = nl.ndarray((TILE, N), dtype=nl.float32, buffer=nl.sbuf) nisa.tensor_copy(dst=sbuf_temp[0:TILE, 0:N], src=psum_tensor[0:TILE, 0:N]) @@ -142,16 +147,19 @@ For ANY scatter-add, accumulation, or read-modify-write on HBM tensors, use `nis with `reduce_op`. Do NOT use manual load+add+store as a workaround. **FORBIDDEN — do NOT generate this code:** + ```python nisa.dma_copy(dst=hbm_dst, src=sbuf_src, dst_rmw_op=nl.add) # FORBIDDEN: dst_rmw_op removed ``` **CORRECT — simple scatter-add:** + ```python nisa.dma_compute(dst=hbm_dst, srcs=[sbuf_src], reduce_op=nl.add) ``` **CORRECT — accumulation loop with indirect indexing:** + ```python for k_idx in range(K): src_access = input_tensor.ap(...) @@ -173,12 +181,14 @@ used for **dynamic loop boundaries and while loop conditions** — they control branching. They are NOT for adding constants to 2D tensors. **FORBIDDEN — do NOT generate this code:** + ```python loop_reg = nisa.register_alloc() nisa.register_move(loop_reg, imm=10) # FORBIDDEN: imm= removed ``` **CORRECT — allocate register with initial value directly:** + ```python src_reg = nisa.register_alloc(x=10) nisa.register_move(dst=loop_reg, src=src_reg) @@ -194,6 +204,7 @@ has **no value at trace/build time**, so `for i in nl.dynamic_range(...)` and ba migrating is safe and works on either frontend. **DEPRECATED — bare `while reg:` (parser-only; removed under tracing / 0.8.0):** + ```python reg = nisa.register_alloc(5) cond = nl.ndarray((1, 1), buffer=nl.sbuf, dtype=nl.int32) @@ -207,6 +218,7 @@ while reg: # DEPRECATED: register has no value at trace time under the tracer **CORRECT — data-dependent loop with `nl.while_loop(init, body_fun)`:** The body returns the next condition register. It is a true while (skips if `init` is zero). + ```python reg = nisa.register_alloc(5) cond = nl.ndarray((1, 1), buffer=nl.sbuf, dtype=nl.int32) @@ -224,6 +236,7 @@ nl.while_loop(reg, body) **CORRECT — counted loop with a runtime bound via `nl.fori_loop(lower, upper, body_fun, step=1)`:** Replaces `for i in nl.dynamic_range(reg): BODY`. The body receives the iteration value as a `VirtualRegister`; read it as a runtime offset via `.ap(scalar_offset=...)` or `nl.ds(offset, size)`. + ```python def body(i): nisa.dma_copy(dst=temp, src=data.ap(scalar_offset=i, indirect_dim=1)) @@ -232,6 +245,7 @@ nl.fori_loop(0, reg, body) # replaces: for i in nl.dynamic_range(reg): ... ``` **Rules for structured dynamic loops:** + - **No loop-carried dependencies.** `fori_loop` follows Pallas semantics; keep cross-iteration state (running max/sum, accumulators, counters) in SBUF/HBM and mutate in place. - **The loop variable is a frozen register.** No writes to it and no raw Python arithmetic on it; @@ -251,12 +265,14 @@ Use named enums for all enum parameters. For hardware DGE mode, use the `dge_mod with `nisa.dge_mode.hwdge`. **FORBIDDEN — do NOT generate this code:** + ```python nisa.dma_copy(src=src_tensor, dst=dst_tensor, dge_mode=2) # FORBIDDEN: integer enum nisa.dma_copy(src=src_tensor, dst=dst_tensor, engine=nisa.dge) # FORBIDDEN: wrong param name ``` **CORRECT — use dge_mode parameter with named enum:** + ```python nisa.dma_copy(dst=dst_tensor, src=src_tensor, dge_mode=nisa.dge_mode.hwdge) ``` @@ -271,6 +287,7 @@ tile overwrites (initializes) PSUM, subsequent tiles accumulate. Do NOT `nisa.me before the loop — the `accumulate=False` first write initializes it. **FORBIDDEN — do NOT generate this code:** + ```python # FORBIDDEN: memset PSUM to zero, then accumulate every iteration psum = nl.ndarray((M, N), dtype=nl.float32, buffer=nl.psum) @@ -280,6 +297,7 @@ for k_idx in nl.affine_range(num_k_tiles): ``` **CORRECT — use accumulate=(k_idx > 0), no memset:** + ```python psum = nl.ndarray((M, N), dtype=nl.float32, buffer=nl.psum) # no memset needed for k_idx in nl.affine_range(num_k_tiles): # affine_range, NOT sequential_range @@ -292,11 +310,13 @@ Use `nisa.tensor_copy()` with `.ap()` (access pattern) and `scalar_offset` for d Do NOT use legacy dynamic-copy APIs or invent helper APIs — they do not exist. **FORBIDDEN — do NOT generate this code:** + ```python nisa.tensor_copy_dynamic_src(dst=dst_tile, src=src_tile, offset=dyn_offset) # FORBIDDEN: deprecated API ``` **CORRECT — use .ap() with scalar_offset:** + ```python nisa.tensor_copy(dst=dst_tile, src=src_tile.ap(scalar_offset=dyn_offset)) ``` diff --git a/skills/neuron-nki-writing/references/nkilib/core/allocator.md b/skills/neuron-nki-writing/references/nkilib/core/allocator.md index dad35cb..8b1996b 100644 --- a/skills/neuron-nki-writing/references/nkilib/core/allocator.md +++ b/skills/neuron-nki-writing/references/nkilib/core/allocator.md @@ -7,6 +7,7 @@ SbufManager is a user-space stack/heap allocator for SBUF memory on NeuronCore. ## When to Use Adopt SbufManager when: + - **4+ SBUF tensors** are allocated in the kernel — centralized allocation prevents address conflicts - **Sub-functions share SBUF space** — pass the allocator as a parameter for composable allocation across call sites - **Multi-buffering / double-buffering** — `open_scope(interleave_degree=2)` + `increment_section()` enables ping-pong buffers @@ -18,29 +19,29 @@ Used in 23 production kernels including attention, QKV projection, MLP, MoE, nor ## Quick Reference -| Method / Function | Signature | Description | -|-------------------|-----------|-------------| -| `SbufManager.__init__` | `(sb_lower_bound, sb_upper_bound, logger=None, use_auto_alloc=False, default_stack_alloc=True)` | Create allocator for an SBUF region | -| `open_scope` | `(interleave_degree=1, name="")` | Push a new stack scope | -| `close_scope` | `()` | Pop scope and free its stack allocations | -| `increment_section` | `()` | Advance multi-buffer section (modular within scope) | -| `alloc` | `(shape, dtype, buffer=nl.sbuf, name=None, base_partition=0, align=None)` | Allocate on default (stack or heap) | -| `alloc_stack` | `(shape, dtype, buffer=nl.sbuf, name=None, base_partition=0, align=None)` | Allocate on stack (auto-freed with scope) | -| `alloc_heap` | `(shape, dtype, buffer=nl.sbuf, name=None, base_partition=0, align=None)` | Allocate on heap (manual free) | -| `pop_heap` | `()` | Free most recent heap allocation | -| `get_total_space` | `() -> int` | Total managed SBUF bytes | -| `get_free_space` | `() -> int` | Available SBUF bytes | -| `get_used_space` | `() -> int` | Used SBUF bytes | -| `get_stack_curr_addr` | `() -> int` | Current stack pointer | -| `get_heap_curr_addr` | `() -> int` | Current heap pointer | -| `align_stack_curr_addr` | `(align=32)` | Align stack pointer to boundary | -| `set_name_prefix` | `(prefix)` | Set prefix for tensor names | -| `get_name_prefix` | `() -> str` | Get current name prefix | -| `flush_logs` | `()` | Print buffered allocation tree | -| `create_auto_alloc_manager` | `(logger=None) -> SbufManager` | Create auto-alloc manager (function) | -| `sizeinbytes` | `(dtype) -> int` | Bytes per element for dtype | -| `align_to` | `(value, alignment) -> int` | Align value up to boundary | -| `num_elts` | `(shape) -> int` | Product of shape elements | +| Method / Function | Signature | Description | +| --------------------------- | ----------------------------------------------------------------------------------------------- | --------------------------------------------------- | +| `SbufManager.__init__` | `(sb_lower_bound, sb_upper_bound, logger=None, use_auto_alloc=False, default_stack_alloc=True)` | Create allocator for an SBUF region | +| `open_scope` | `(interleave_degree=1, name="")` | Push a new stack scope | +| `close_scope` | `()` | Pop scope and free its stack allocations | +| `increment_section` | `()` | Advance multi-buffer section (modular within scope) | +| `alloc` | `(shape, dtype, buffer=nl.sbuf, name=None, base_partition=0, align=None)` | Allocate on default (stack or heap) | +| `alloc_stack` | `(shape, dtype, buffer=nl.sbuf, name=None, base_partition=0, align=None)` | Allocate on stack (auto-freed with scope) | +| `alloc_heap` | `(shape, dtype, buffer=nl.sbuf, name=None, base_partition=0, align=None)` | Allocate on heap (manual free) | +| `pop_heap` | `()` | Free most recent heap allocation | +| `get_total_space` | `() -> int` | Total managed SBUF bytes | +| `get_free_space` | `() -> int` | Available SBUF bytes | +| `get_used_space` | `() -> int` | Used SBUF bytes | +| `get_stack_curr_addr` | `() -> int` | Current stack pointer | +| `get_heap_curr_addr` | `() -> int` | Current heap pointer | +| `align_stack_curr_addr` | `(align=32)` | Align stack pointer to boundary | +| `set_name_prefix` | `(prefix)` | Set prefix for tensor names | +| `get_name_prefix` | `() -> str` | Get current name prefix | +| `flush_logs` | `()` | Print buffered allocation tree | +| `create_auto_alloc_manager` | `(logger=None) -> SbufManager` | Create auto-alloc manager (function) | +| `sizeinbytes` | `(dtype) -> int` | Bytes per element for dtype | +| `align_to` | `(value, alignment) -> int` | Align value up to boundary | +| `num_elts` | `(shape) -> int` | Product of shape elements | ## Import Options @@ -48,6 +49,7 @@ Used in 23 production kernels including attention, QKV projection, MLP, MoE, nor Source: `references/nkilib/core/utils/allocator.py` **If nkilib is installed** in the user's environment: + ```python from nkilib.core.utils.allocator import SbufManager, create_auto_alloc_manager ``` @@ -59,6 +61,7 @@ from nkilib.core.utils.allocator import SbufManager, create_auto_alloc_manager Create an SBUF memory manager. **Args:** + - `sb_lower_bound` (`int`): Lower bound of available SBUF region (stack starts here) - `sb_upper_bound` (`int`): Upper bound of available SBUF region (heap starts here) - `logger` (`Logger`, optional): Logger instance; creates default "SBM" logger if None @@ -76,6 +79,7 @@ sbm = SbufManager(0, 128 * 1024) # 128KB SBUF region Push a new stack scope. All stack allocations within this scope are freed when `close_scope()` is called. **Args:** + - `interleave_degree` (`int`): Number of multi-buffer sections (default: 1 = no interleaving) - `name` (`str`): Optional scope name for debug logging @@ -114,6 +118,7 @@ sbm.close_scope() Allocate on default target (stack or heap based on `default_stack_alloc`). **Args:** + - `shape` (`tuple`): Tensor shape; first dim is partition, rest are free dims - `dtype`: Data type (`nl.bfloat16`, `nl.float32`, etc.) - `buffer`: Buffer type (only `nl.sbuf` supported) @@ -134,6 +139,7 @@ Allocate on the stack. Requires an open scope. Freed automatically when scope cl **Returns:** `nl.ndarray` allocated on the stack. **Constraints:** + - Must have an open scope - `buffer` must be `nl.sbuf` - Must have sufficient free space (stack_addr + size <= heap_addr) @@ -196,13 +202,13 @@ sbm = create_auto_alloc_manager() Return byte size per element for a given NKI data type. -| dtype | Size | -|-------|------| -| `nl.float32`, `nl.int32`, `nl.uint32` | 4 | -| `nl.float8_e4m3fn_x4`, `nl.float8_e5m2_x4` | 4 | -| `nl.bfloat16`, `nl.float16`, `nl.uint16`, `nl.int16` | 2 | -| `nl.float4_e2m1fn_x4` | 2 | -| `nl.int8`, `nl.uint8`, `float8_e4m3`, `float8_e5m2`, `float8e4`, `float8e5` | 1 | +| dtype | Size | +| --------------------------------------------------------------------------- | ---- | +| `nl.float32`, `nl.int32`, `nl.uint32` | 4 | +| `nl.float8_e4m3fn_x4`, `nl.float8_e5m2_x4` | 4 | +| `nl.bfloat16`, `nl.float16`, `nl.uint16`, `nl.int16` | 2 | +| `nl.float4_e2m1fn_x4` | 2 | +| `nl.int8`, `nl.uint8`, `float8_e4m3`, `float8_e5m2`, `float8e4`, `float8e5` | 1 | --- diff --git a/skills/neuron-nki-writing/references/nkilib/core/kernel-helpers.md b/skills/neuron-nki-writing/references/nkilib/core/kernel-helpers.md index 3f623e0..4f96150 100644 --- a/skills/neuron-nki-writing/references/nkilib/core/kernel-helpers.md +++ b/skills/neuron-nki-writing/references/nkilib/core/kernel-helpers.md @@ -7,32 +7,34 @@ Kernel helpers provide commonly-used utility functions for NKI kernels: ceiling/ ## When to Use **Always use:** + - `div_ceil(n, d)` — for any tile count computation. Never write `(n + d - 1) // d` inline. Used in 60+ call sites across 20+ production kernels. - `kernel_assert()` — for all input validation. Never use Python `assert` in NKI kernels. **Use when needed:** + - `get_ceil_aligned_size()` / `get_floor_aligned_size()` — when allocating buffers that must be aligned to hardware boundaries - `is_launched_as_spmd()` / `get_program_sharding_info()` — for SPMD-aware kernels that shard across NeuronCores - `get_max_positive_value_for_dtype()` — when computing softmax masks or clamping to dtype range ## Quick Reference -| Function | Signature | Description | -|----------|-----------|-------------| -| `is_hbm_buffer` | `(tensor: nl.ndarray) -> bool` | Check if tensor buffer is HBM | -| `get_ceil_quotient` | `(numerator, denominator) -> int` | Ceiling division | -| `div_ceil` | `(n, d) -> int` | Ceiling division (alias) | -| `get_ceil_aligned_size` | `(size, alignment_multiple) -> int` | Round up to alignment boundary | -| `get_floor_quotient` | `(numerator, denominator) -> int` | Floor division | -| `get_floor_aligned_size` | `(size, alignment_multiple) -> int` | Round down to alignment boundary | -| `get_nl_act_fn_from_type` | `(act_fn: ActFnType) -> function` | Map enum to NKI activation function | -| `is_launched_as_spmd` | `() -> bool` | Check if running in SPMD mode | -| `get_program_sharding_info` | `() -> Tuple[int, int, int]` | Get (grid_ndim, n_prgs, prg_id) | -| `get_verified_program_sharding_info` | `(kernel_name, allowed_ndims, max_sharding) -> Tuple` | Get sharding info with validation | -| `is_rms_normalization` | `(norm_type: NormType) -> bool` | Check if norm type is RMS | -| `normalization_uses_weights` | `(norm_type: NormType) -> bool` | Check if norm uses weight params | -| `get_max_positive_value_for_dtype` | `(dtype) -> float` | Max positive value for FP8 types | -| `reduce` | `(op, input, initial_value) -> result` | Reduce list with mul/add/min/max | +| Function | Signature | Description | +| ------------------------------------ | ----------------------------------------------------- | ----------------------------------- | +| `is_hbm_buffer` | `(tensor: nl.ndarray) -> bool` | Check if tensor buffer is HBM | +| `get_ceil_quotient` | `(numerator, denominator) -> int` | Ceiling division | +| `div_ceil` | `(n, d) -> int` | Ceiling division (alias) | +| `get_ceil_aligned_size` | `(size, alignment_multiple) -> int` | Round up to alignment boundary | +| `get_floor_quotient` | `(numerator, denominator) -> int` | Floor division | +| `get_floor_aligned_size` | `(size, alignment_multiple) -> int` | Round down to alignment boundary | +| `get_nl_act_fn_from_type` | `(act_fn: ActFnType) -> function` | Map enum to NKI activation function | +| `is_launched_as_spmd` | `() -> bool` | Check if running in SPMD mode | +| `get_program_sharding_info` | `() -> Tuple[int, int, int]` | Get (grid_ndim, n_prgs, prg_id) | +| `get_verified_program_sharding_info` | `(kernel_name, allowed_ndims, max_sharding) -> Tuple` | Get sharding info with validation | +| `is_rms_normalization` | `(norm_type: NormType) -> bool` | Check if norm type is RMS | +| `normalization_uses_weights` | `(norm_type: NormType) -> bool` | Check if norm uses weight params | +| `get_max_positive_value_for_dtype` | `(dtype) -> float` | Max positive value for FP8 types | +| `reduce` | `(op, input, initial_value) -> result` | Reduce list with mul/add/min/max | **Constants:** | Constant | Value | Description | @@ -46,6 +48,7 @@ Kernel helpers provide commonly-used utility functions for NKI kernels: ceiling/ Source: `references/nkilib/core/utils/kernel_helpers.py` **If nkilib is installed** in the user's environment: + ```python from nkilib.core.utils.kernel_helpers import ( is_hbm_buffer, @@ -66,6 +69,7 @@ from nkilib.core.utils.kernel_helpers import ( Check if the tensor's buffer is any HBM type (hbm, shared_hbm, or private_hbm). **Args:** + - `tensor` (`nl.ndarray`): NKI tensor to check **Returns:** `True` if the tensor buffer is `nl.hbm`, `nl.shared_hbm`, or `nl.private_hbm`. @@ -85,6 +89,7 @@ is_hbm_buffer(sbuf_tensor) # False Compute ceiling division using integer arithmetic. **Args:** + - `numerator` (`int`): Dividend - `denominator` (`int`): Divisor (must be non-zero) @@ -112,6 +117,7 @@ num_tiles = div_ceil(seq_len, tile_size) Round `size` up to the nearest multiple of `alignment_multiple`. **Args:** + - `size` (`int`): Value to align - `alignment_multiple` (`int`): Alignment boundary @@ -149,6 +155,7 @@ get_floor_aligned_size(100, 64) # 64 Map an `ActFnType` enum to the corresponding `nki.language` activation function. **Args:** + - `act_fn` (`ActFnType`): One of `SiLU`, `GELU`, `GELU_Tanh_Approx`, `Swish` **Returns:** NKI function (`nl.silu`, `nl.gelu`, `nl.gelu_apprx_tanh`, `nl.gelu_apprx_sigmoid`) @@ -184,6 +191,7 @@ grid_ndim, n_prgs, prg_id = get_program_sharding_info() Same as `get_program_sharding_info` with optional validation. **Args:** + - `kernel_name` (`str`): Kernel name for error messages - `allowed_ndims` (`Tuple[int, ...]`, optional): Allowed grid dimensions - `max_sharding` (`int`, optional): Maximum sharding degree @@ -207,6 +215,7 @@ Check if the normalization type uses weight parameters (`RMS_NORM` or `LAYER_NOR Get maximum positive representable value for FP8 data types. **Args:** + - `dtype`: `nl.float8_e4m3` or `nl.float8_e5m2` **Returns:** `240.0` for e4m3, `57344.0` for e5m2, `None` for other types. @@ -218,6 +227,7 @@ Get maximum positive representable value for FP8 data types. Perform a reduction operation over a list. **Args:** + - `op` (`str`): One of `'mul'`, `'add'`, `'min'`, `'max'` - `input` (`List`): Values to reduce - `initial_value`: Starting accumulator value diff --git a/skills/neuron-nki-writing/references/nkilib/core/tensor-view.md b/skills/neuron-nki-writing/references/nkilib/core/tensor-view.md index 18f91d9..b435312 100644 --- a/skills/neuron-nki-writing/references/nkilib/core/tensor-view.md +++ b/skills/neuron-nki-writing/references/nkilib/core/tensor-view.md @@ -7,6 +7,7 @@ TensorView is a high-level wrapper around NKI tensors that provides PyTorch-like ## When to Use Adopt TensorView when the kernel needs any of: + - **Strided/interleaved DMA**: `slice(dim, start, end, step=2)` gathers even/odd elements without loops - **Broadcasting**: `expand_dim(d).broadcast(d, size)` replicates across a dimension (e.g., cos/sin across heads in RoPE) - **Reshape without copy**: `reshape_dim()` / `flatten_dims()` to reshape multi-dimensional tensors for DMA or matmul @@ -19,24 +20,24 @@ Used in 17+ production kernels including attention, RoPE, MLP projections, MoE e ## Quick Reference -| Method | Signature | Description | -|--------|-----------|-------------| -| `__init__` | `(base_tensor: nl.ndarray)` | Create a view from an NKI tensor | -| `get_view` | `() -> nl.ndarray` | Generate the actual NKI tensor with array pattern applied | -| `slice` | `(dim, start, end, step=1) -> TensorView` | Slice along a dimension | -| `permute` | `(dims: List[int]) -> TensorView` | Reorder dimensions | -| `broadcast` | `(dim, size) -> TensorView` | Expand a size-1 dimension | -| `reshape_dim` | `(dim, shape: List[int]) -> TensorView` | Split one dimension into multiple | -| `flatten_dims` | `(start_dim, end_dim) -> TensorView` | Flatten contiguous dimensions into one | -| `expand_dim` | `(dim) -> TensorView` | Insert a size-1 dimension | -| `squeeze_dim` | `(dim) -> TensorView` | Remove a size-1 dimension | -| `select` | `(dim, index) -> TensorView` | Select a single index along a dimension | -| `rearrange` | `(src_pattern, dst_pattern, fixed_sizes=None) -> TensorView` | Einops-style dimension rearrangement | -| `get_dim` | `() -> int` | Return number of dimensions | -| `is_sbuf` | `() -> bool` | Check if base tensor is in SBUF | -| `is_hbm` | `() -> bool` | Check if base tensor is in HBM | -| `has_dynamic_access` | `() -> bool` | Check if view uses dynamic (indirect) indexing | -| `get_trivial_strides` | `(shape, base_stride=1) -> Tuple[int, ...]` | Compute row-major strides (static) | +| Method | Signature | Description | +| --------------------- | ------------------------------------------------------------ | --------------------------------------------------------- | +| `__init__` | `(base_tensor: nl.ndarray)` | Create a view from an NKI tensor | +| `get_view` | `() -> nl.ndarray` | Generate the actual NKI tensor with array pattern applied | +| `slice` | `(dim, start, end, step=1) -> TensorView` | Slice along a dimension | +| `permute` | `(dims: List[int]) -> TensorView` | Reorder dimensions | +| `broadcast` | `(dim, size) -> TensorView` | Expand a size-1 dimension | +| `reshape_dim` | `(dim, shape: List[int]) -> TensorView` | Split one dimension into multiple | +| `flatten_dims` | `(start_dim, end_dim) -> TensorView` | Flatten contiguous dimensions into one | +| `expand_dim` | `(dim) -> TensorView` | Insert a size-1 dimension | +| `squeeze_dim` | `(dim) -> TensorView` | Remove a size-1 dimension | +| `select` | `(dim, index) -> TensorView` | Select a single index along a dimension | +| `rearrange` | `(src_pattern, dst_pattern, fixed_sizes=None) -> TensorView` | Einops-style dimension rearrangement | +| `get_dim` | `() -> int` | Return number of dimensions | +| `is_sbuf` | `() -> bool` | Check if base tensor is in SBUF | +| `is_hbm` | `() -> bool` | Check if base tensor is in HBM | +| `has_dynamic_access` | `() -> bool` | Check if view uses dynamic (indirect) indexing | +| `get_trivial_strides` | `(shape, base_stride=1) -> Tuple[int, ...]` | Compute row-major strides (static) | ## Import Options @@ -44,6 +45,7 @@ Used in 17+ production kernels including attention, RoPE, MLP projections, MoE e Source: `references/nkilib/core/utils/tensor_view.py` **If nkilib is installed** in the user's environment: + ```python from nkilib.core.utils.tensor_view import TensorView ``` @@ -55,11 +57,13 @@ from nkilib.core.utils.tensor_view import TensorView Create a TensorView wrapping an NKI tensor. **Args:** + - `base_tensor` (`nl.ndarray`): The underlying NKI tensor. Must not be None. **Returns:** None (constructor) **Constraints:** + - `base_tensor` must not be `None` ```python @@ -88,6 +92,7 @@ result = sliced.get_view() # nl.ndarray usable in NKI ops Create a sliced view along a specific dimension. **Args:** + - `dim` (`int`): Dimension to slice - `start` (`int`): Start index (inclusive), must be >= 0 - `end` (`int`): End index (exclusive), must be > start and <= shape[dim] @@ -96,6 +101,7 @@ Create a sliced view along a specific dimension. **Returns:** New `TensorView` with sliced dimension. **Constraints:** + - `dim < get_dim()` - `0 <= start < end <= shape[dim]` @@ -112,11 +118,13 @@ sliced = view.slice(1, 0, 256) Create a permuted view by reordering dimensions. **Args:** + - `dims` (`List[int]`): New order of dimensions. Must be a valid permutation. **Returns:** New `TensorView` with permuted dimensions. **Constraints:** + - Length of `dims` must equal number of dimensions - No duplicate indices - For SBUF tensors, `dims[0]` must be `0` (partition dimension stays outermost) @@ -134,12 +142,14 @@ permuted = view.permute([0, 2, 1]) Expand a size-1 dimension by broadcasting (stride set to 0, same element repeated). **Args:** + - `dim` (`int`): Dimension to broadcast (must currently have size 1) - `size` (`int`): New size for the dimension **Returns:** New `TensorView` with broadcasted dimension. **Constraints:** + - `shape[dim]` must be 1 - For SBUF tensors, partition dim cannot be broadcast beyond `nl.tile_size.pmax` @@ -156,12 +166,14 @@ broadcasted = view.broadcast(1, 8) Split a single dimension into multiple dimensions. Supports `-1` for one inferred dimension. **Args:** + - `dim` (`int`): Dimension to reshape - `shape` (`List[int]`): New sizes (product must equal original dim size; at most one `-1`) **Returns:** New `TensorView` with reshaped dimension. **Constraints:** + - Product of `shape` must equal `self.shape[dim]` - For SBUF, partition dim (dim 0) cannot be reshaped (except trivially) @@ -178,12 +190,14 @@ reshaped = view.reshape_dim(1, [2, -1, 4]) # -1 inferred as 3 Flatten a contiguous range of dimensions into a single dimension. **Args:** + - `start_dim` (`int`): First dimension to flatten (inclusive) - `end_dim` (`int`): Last dimension to flatten (inclusive) **Returns:** New `TensorView` with flattened dimensions. **Constraints:** + - `start_dim < end_dim` - Dimensions must be contiguous in memory - For SBUF, `start_dim > 0` (cannot flatten partition dim) @@ -201,11 +215,13 @@ flat = view.flatten_dims(1, 3) Insert a new dimension of size 1 at the specified position. **Args:** + - `dim` (`int`): Position to insert (0 to get_dim() inclusive) **Returns:** New `TensorView` with additional dimension. **Constraints:** + - For SBUF, `dim > 0` (cannot expand before partition dim) ```python @@ -221,11 +237,13 @@ expanded = view.expand_dim(1) Remove a dimension that has size 1. **Args:** + - `dim` (`int`): Dimension to remove (must have size 1) **Returns:** New `TensorView` with dimension removed. **Constraints:** + - `shape[dim]` must be 1 - For SBUF, `dim > 0` @@ -242,6 +260,7 @@ squeezed = view.squeeze_dim(1) Select a single element along a dimension, reducing dimensionality by one. **Args:** + - `dim` (`int`): Dimension to select from - `index` (`int` or `nl.ndarray`): Static integer index, or a scalar NKI tensor for dynamic indexing @@ -264,6 +283,7 @@ dynamic_selected = view.select(0, idx_tensor) Einops-style dimension rearrangement combining reshape, permute, and flatten. **Args:** + - `src_pattern` (`Tuple[Union[str, Tuple[str]]]`): Source dimension names. Tuples indicate grouped dims to split. - `dst_pattern` (`Tuple[Union[str, Tuple[str]]]`): Destination dimension names. Tuples indicate dims to flatten. - `fixed_sizes` (`Dict[str, int]`, optional): Known sizes for dimensions used in reshaping. @@ -287,6 +307,7 @@ rearranged = view.rearrange( Compute row-major (C-style) strides for a given shape. **Args:** + - `shape` (`List[int]`): Dimension sizes - `base_stride` (`int`): Stride of innermost dimension (default: 1) @@ -321,16 +342,19 @@ Reshape the tensor to new dimensions without copying data. The total number of e For non-HBM tensors (SBUF/PSUM), the partition dimension (dim 0) size must be preserved in the new shape. For HBM tensors, all dimensions participate in reshape. The algorithm has three phases: + 1. **Remove unit dims** -- strip size-1 dims whose strides are irrelevant 2. **Collapse contiguous** -- merge adjacent dims with contiguous strides into blocks 3. **Repartition** -- assign new strides by splitting/merging blocks to match `new_shape` **Args:** + - `new_shape` (`Tuple[int, ...]`): New dimension sizes (total elements must match) **Returns:** New `TensorView` with reshaped dimensions. **Constraints:** + - Total element count must match between old and new shapes - For non-HBM tensors, `new_shape[0]` must equal current `shape[0]` - Layout must be compatible (fails if reshape would require a data copy) diff --git a/skills/neuron-nki-writing/references/nkilib/core/tile-info.md b/skills/neuron-nki-writing/references/nkilib/core/tile-info.md index a6d7d6f..d23faf0 100644 --- a/skills/neuron-nki-writing/references/nkilib/core/tile-info.md +++ b/skills/neuron-nki-writing/references/nkilib/core/tile-info.md @@ -7,29 +7,30 @@ TiledDimInfo is a dataclass that encapsulates tiling metadata for a single dimen ## When to Use Adopt TiledDimInfo when: + - **CTE-style kernel** with precomputed tile metadata that multiple functions query (tile counts, last-block sizes, subtile bounds) - **Two-level tiling** with subtiles nested inside tiles — `build_with_subtiling()` precomputes both levels - **Parameter structs** that carry tiling config — store a `TiledDimInfo` per tiled dimension instead of loose integers -**Skip when**: `TiledRange` is sufficient for iteration. TiledDimInfo is for *metadata storage and querying*, TiledRange is for *iteration*. +**Skip when**: `TiledRange` is sufficient for iteration. TiledDimInfo is for _metadata storage and querying_, TiledRange is for _iteration_. Used in 4 production kernels (output projection CTE, RMSNorm quant, MLP CTE tile info, MLP CTE transpose) where tiling metadata is built once and queried across multiple kernel phases. ## Quick Reference -| Method | Signature | Description | -|--------|-----------|-------------| -| `build` | `(tiled_dim_size, tile_size, subtile_info=None) -> TiledDimInfo` | Factory: create from dimension size and tile size | -| `build_with_subtiling` | `(tiled_dim_size, tile_size, subtile_size) -> TiledDimInfo` | Factory: create with two-level tiling | -| `is_subtiled` | `() -> bool` | Check if subtile info is present | -| `get_tile_indices` | `(tile_num, tile_offset) -> nl.ds` | Get NKI index slice for a tile | -| `get_subtile_indices` | `(tile_num, subtile_num, subtile_offset) -> nl.ds` | Get NKI index slice for a subtile | -| `get_subtile_start` | `(tile_idx, subtile_idx) -> int` | Absolute start position of a subtile | -| `get_local_subtile_start` | `(subtile_idx) -> int` | Local start position within a loaded tile | -| `get_subtile_bound` | `(tile_idx, subtile_idx) -> int` | Valid size of a subtile (handles remainder) | -| `get_local_subtile_bound` | `(tile_idx, subtile_idx) -> int` | Valid local size within loaded tile | -| `get_tile_bound` | `(tile_idx) -> int` | Valid size of a tile (handles remainder) | -| `get_actual_subtile_num` | `(tile_idx) -> int` | Number of subtiles in a given tile | +| Method | Signature | Description | +| ------------------------- | ---------------------------------------------------------------- | ------------------------------------------------- | +| `build` | `(tiled_dim_size, tile_size, subtile_info=None) -> TiledDimInfo` | Factory: create from dimension size and tile size | +| `build_with_subtiling` | `(tiled_dim_size, tile_size, subtile_size) -> TiledDimInfo` | Factory: create with two-level tiling | +| `is_subtiled` | `() -> bool` | Check if subtile info is present | +| `get_tile_indices` | `(tile_num, tile_offset) -> nl.ds` | Get NKI index slice for a tile | +| `get_subtile_indices` | `(tile_num, subtile_num, subtile_offset) -> nl.ds` | Get NKI index slice for a subtile | +| `get_subtile_start` | `(tile_idx, subtile_idx) -> int` | Absolute start position of a subtile | +| `get_local_subtile_start` | `(subtile_idx) -> int` | Local start position within a loaded tile | +| `get_subtile_bound` | `(tile_idx, subtile_idx) -> int` | Valid size of a subtile (handles remainder) | +| `get_local_subtile_bound` | `(tile_idx, subtile_idx) -> int` | Valid local size within loaded tile | +| `get_tile_bound` | `(tile_idx) -> int` | Valid size of a tile (handles remainder) | +| `get_actual_subtile_num` | `(tile_idx) -> int` | Number of subtiles in a given tile | ## Import Options @@ -37,6 +38,7 @@ Used in 4 production kernels (output projection CTE, RMSNorm quant, MLP CTE tile Source: `references/nkilib/core/utils/tile_info.py` **If nkilib is installed** in the user's environment: + ```python from nkilib.core.utils.tile_info import TiledDimInfo ``` @@ -48,6 +50,7 @@ from nkilib.core.utils.tile_info import TiledDimInfo Factory method to create a TiledDimInfo from dimension size and tile size. **Args:** + - `tiled_dim_size` (`int`): Total size of the dimension being tiled - `tile_size` (`int`): Size of each tile - `subtile_info` (`TiledDimInfo`, optional): Nested subtile information @@ -66,6 +69,7 @@ info = TiledDimInfo.build(1024, 256) Factory method to create a TiledDimInfo with two-level tiling (tile + subtile). **Args:** + - `tiled_dim_size` (`int`): Total size of the dimension - `tile_size` (`int`): Size of each outer tile - `subtile_size` (`int`): Size of each inner subtile within a tile @@ -92,6 +96,7 @@ Check whether this dimension has subtile information. Get an NKI dynamic slice for a given tile. **Args:** + - `tile_num`: Tile number (0-based) - `tile_offset`: Offset size for the slice @@ -108,6 +113,7 @@ idx = info.get_tile_indices(2, 256) # nl.ds(512, 256) for tile_size=256 Get an NKI dynamic slice for a specific subtile within a tile. **Args:** + - `tile_num`: Outer tile number - `subtile_num`: Subtile number within the tile - `subtile_offset`: Offset size for the slice @@ -123,6 +129,7 @@ Get an NKI dynamic slice for a specific subtile within a tile. Calculate absolute start position for a subtile. **Args:** + - `tile_idx`: Outer tile index - `subtile_idx`: Subtile index within the tile @@ -137,6 +144,7 @@ Calculate absolute start position for a subtile. Calculate the local start position of a subtile within a loaded tile. **Args:** + - `subtile_idx`: Subtile index **Returns:** `subtile_idx * subtile_size` @@ -150,6 +158,7 @@ Calculate the local start position of a subtile within a loaded tile. Calculate valid size of a subtile, clamped to the dimension boundary. **Args:** + - `tile_idx`: Outer tile index - `subtile_idx`: Subtile index @@ -164,6 +173,7 @@ Calculate valid size of a subtile, clamped to the dimension boundary. Calculate valid local size of a subtile within a loaded tile. **Args:** + - `tile_idx`: Outer tile index - `subtile_idx`: Subtile index @@ -178,6 +188,7 @@ Calculate valid local size of a subtile within a loaded tile. Calculate valid size of a tile, clamped to the dimension boundary. **Args:** + - `tile_idx`: Tile index **Returns:** `min(tiled_dim_size - tile_start, tile_size)` @@ -195,6 +206,7 @@ info.get_tile_bound(2) # 44 (remainder) Calculate the actual number of subtiles in a given tile (handles partial tiles). **Args:** + - `tile_idx`: Tile index **Returns:** Ceiling division of `tile_bound / subtile_size` diff --git a/skills/neuron-nki-writing/references/nkilib/core/tiled-range.md b/skills/neuron-nki-writing/references/nkilib/core/tiled-range.md index eefdf21..5642ecc 100644 --- a/skills/neuron-nki-writing/references/nkilib/core/tiled-range.md +++ b/skills/neuron-nki-writing/references/nkilib/core/tiled-range.md @@ -7,6 +7,7 @@ TiledRange divides a dimension into fixed-size tiles, handling remainder logic f ## When to Use Adopt TiledRange when: + - **Tiling any dimension** where the size may not be evenly divisible by the tile size (remainder handling) - **Nested tiling**: pass an outer `TiledRangeIterator` as input to create subtiles — avoids manual two-level remainder logic - **Multiple tiled dimensions**: each `TiledRangeIterator` carries `.size`, `.start_offset`, `.index` so DMA copies use correct bounds @@ -17,11 +18,11 @@ Used in 8+ production kernels including cumsum, RMSNorm, router TopK, and MLP pr ## Quick Reference -| Name | Signature | Description | -|------|-----------|-------------| -| `TiledRange` | `(size, tile_size: int) -> Tuple[TiledRangeIterator, ...]` | Divide a dimension into tiles and return iterators | -| `TiledRangeIterator` | `(tile_size, tile_index, start_offset, end_offset)` | Single tile with `.size`, `.index`, `.start_offset`, `.end_offset` properties | -| `TiledRangeIterator.__repr__` | `() -> str` | String representation for debugging | +| Name | Signature | Description | +| ----------------------------- | ---------------------------------------------------------- | ----------------------------------------------------------------------------- | +| `TiledRange` | `(size, tile_size: int) -> Tuple[TiledRangeIterator, ...]` | Divide a dimension into tiles and return iterators | +| `TiledRangeIterator` | `(tile_size, tile_index, start_offset, end_offset)` | Single tile with `.size`, `.index`, `.start_offset`, `.end_offset` properties | +| `TiledRangeIterator.__repr__` | `() -> str` | String representation for debugging | ## Import Options @@ -29,6 +30,7 @@ Used in 8+ production kernels including cumsum, RMSNorm, router TopK, and MLP pr Source: `references/nkilib/core/utils/tiled_range.py` **If nkilib is installed** in the user's environment: + ```python from nkilib.core.utils.tiled_range import TiledRange, TiledRangeIterator ``` @@ -40,12 +42,14 @@ from nkilib.core.utils.tiled_range import TiledRange, TiledRangeIterator Divide a dimension into tiles and return a tuple of iterators. **Args:** + - `size` (`int` or `TiledRangeIterator`): Total size to tile, or a `TiledRangeIterator` for nested (sub)tiling - `tile_size` (`int`): Size of each tile **Returns:** Tuple of `TiledRangeIterator` objects. The last tile may be smaller than `tile_size` if the dimension is not evenly divisible. **Constraints:** + - `tile_size` should be > 0 - When `size` is a `TiledRangeIterator`, tiling operates on that tile's `.size` and offsets are computed relative to the parent tile's `.start_offset` @@ -64,6 +68,7 @@ tiles = TiledRange(300, 128) Represents a single tile in a tiled range. **Attributes:** + - `size` (`int`): Size of this tile (may be < tile_size for last tile) - `index` (`int`): 0-based index of this tile in the range - `start_offset` (`int`): Absolute starting offset in the original dimension diff --git a/skills/neuron-nki-writing/references/nkilib/ops/stream-shuffle-broadcast.md b/skills/neuron-nki-writing/references/nkilib/ops/stream-shuffle-broadcast.md index dd7a340..8f64846 100644 --- a/skills/neuron-nki-writing/references/nkilib/ops/stream-shuffle-broadcast.md +++ b/skills/neuron-nki-writing/references/nkilib/ops/stream-shuffle-broadcast.md @@ -7,6 +7,7 @@ Broadcasts the first partition of a source tensor across all partitions of a des ## When to Use Adopt stream_shuffle_broadcast when: + - **Bias/scale addition after DMA load**: a 1D vector (bias, quantization scale, affinity score) was loaded into partition 0 and must be replicated to all 128 partitions before element-wise operations - **Scalar broadcast**: any value that exists in a single partition but is needed across all PEs @@ -16,8 +17,8 @@ Used in 13+ production kernels including attention (RoPE positions, softmax stat ## Quick Reference -| Function | Description | -|----------|-------------| +| Function | Description | +| ------------------------------------ | ----------------------------------------------------------------------- | | `stream_shuffle_broadcast(src, dst)` | Broadcast `src[0:1, :]` to all partitions of `dst` using stream shuffle | ## Import Options @@ -26,6 +27,7 @@ Used in 13+ production kernels including attention (RoPE positions, softmax stat Source: `references/nkilib/core/utils/stream_shuffle_broadcast.py` **If nkilib is installed** in the user's environment: + ```python from nkilib.core.utils.stream_shuffle_broadcast import stream_shuffle_broadcast ``` @@ -38,14 +40,15 @@ Broadcasts the first partition (`src[0:1, :]`) onto every partition of `dst` usi **Args:** -| Parameter | Type | Description | -|-----------|------|-------------| -| `src` | `nl.ndarray` (2D) | Source tensor in SBUF. Only partition 0 is read. | -| `dst` | `nl.ndarray` (2D) | Destination tensor in SBUF. All partitions are written. | +| Parameter | Type | Description | +| --------- | ----------------- | ------------------------------------------------------- | +| `src` | `nl.ndarray` (2D) | Source tensor in SBUF. Only partition 0 is read. | +| `dst` | `nl.ndarray` (2D) | Destination tensor in SBUF. All partitions are written. | **Returns:** None (writes result into `dst`). **Constraints:** + - Both `src` and `dst` must be 2D tensors. - The free dimension (axis 1) of `src` must match that of `dst`: `src.shape[1] == dst.shape[1]`. - Both tensors must reside in SBUF. @@ -53,6 +56,7 @@ Broadcasts the first partition (`src[0:1, :]`) onto every partition of `dst` usi - Internally processes partitions in chunks of 32 (the stream shuffle hardware width). **Example:** + ```python import nki.language as nl from nkilib.core.utils.stream_shuffle_broadcast import stream_shuffle_broadcast @@ -66,6 +70,7 @@ stream_shuffle_broadcast(src=shared_params, dst=expanded_params) ## Usage Examples ### Pattern 1: Broadcasting a shared bias across partitions + ```python # Load bias into partition 0, then broadcast to all partitions bias_p0 = nl.ndarray((1, hidden_dim), dtype=nl.float32, buffer=nl.sbuf) @@ -76,6 +81,7 @@ stream_shuffle_broadcast(src=bias_p0, dst=bias_all) ``` ### Pattern 2: Replicating a scaling vector for element-wise ops + ```python # Single-partition scale factor replicated for parallel computation scale_single = nl.ndarray((1, seq_len), dtype=nl.float32, buffer=nl.sbuf) diff --git a/skills/neuron-nki-writing/references/nkilib/ops/tp-broadcast.md b/skills/neuron-nki-writing/references/nkilib/ops/tp-broadcast.md index 0b4d756..9767f7e 100644 --- a/skills/neuron-nki-writing/references/nkilib/ops/tp-broadcast.md +++ b/skills/neuron-nki-writing/references/nkilib/ops/tp-broadcast.md @@ -7,6 +7,7 @@ Transposes a column from the source tensor and broadcasts it across all partitio ## When to Use Adopt tp_broadcast when: + - **Partition-to-free dimension broadcast**: a value exists as a column vector along the partition dimension and needs to be transposed and replicated across all partitions in the free dimension (e.g., softmax max value in attention) **Skip when**: `stream_shuffle_broadcast` suffices (simpler, for same-dimension broadcast), or the value is already in the correct layout. @@ -15,8 +16,8 @@ Highly specialized: used in 1 production kernel (attention TKG — broadcasting ## Quick Reference -| Function | Description | -|----------|-------------| +| Function | Description | +| ------------------------------------------------------- | -------------------------------------------------------------------------------------- | | `tp_broadcast(src, dst, src_offset, psum_address=None)` | Transpose `src[0:1, :]` and broadcast to all partitions of `dst` via PSUM intermediate | ## Import Options @@ -25,6 +26,7 @@ Highly specialized: used in 1 production kernel (attention TKG — broadcasting Source: `references/nkilib/core/utils/tp_broadcast.py` **If nkilib is installed** in the user's environment: + ```python from nkilib.core.utils.tp_broadcast import tp_broadcast ``` @@ -37,16 +39,17 @@ Transposes then broadcasts `src[0:1, :]` onto all partitions of `dst`. Each part **Args:** -| Parameter | Type | Default | Description | -|-----------|------|---------|-------------| -| `src` | `nl.ndarray` (2D) | (required) | Source tensor in SBUF. Shape: `[P, F]`. Only partition 0 is read. | -| `dst` | `nl.ndarray` (2D) | (required) | Destination tensor in SBUF. Shape: `[B, P]` where `B` is the broadcast dimension. | -| `src_offset` | `int` | (required) | Offset in the free dimension (F) to select which column to transpose from. | -| `psum_address` | `int` or `None` | `None` | Optional explicit PSUM bank address for the intermediate transpose buffer. | +| Parameter | Type | Default | Description | +| -------------- | ----------------- | ---------- | --------------------------------------------------------------------------------- | +| `src` | `nl.ndarray` (2D) | (required) | Source tensor in SBUF. Shape: `[P, F]`. Only partition 0 is read. | +| `dst` | `nl.ndarray` (2D) | (required) | Destination tensor in SBUF. Shape: `[B, P]` where `B` is the broadcast dimension. | +| `src_offset` | `int` | (required) | Offset in the free dimension (F) to select which column to transpose from. | +| `psum_address` | `int` or `None` | `None` | Optional explicit PSUM bank address for the intermediate transpose buffer. | **Returns:** None (writes result into `dst`). **Constraints:** + - `src` must be 2D with shape `[P, F]`. - `dst` must be 2D with shape `[B, P]` where `P` matches `src.shape[0]`. - The transposed dimension of `dst` (`dst.shape[1]`) must equal `src.shape[0]` (the partition dimension). @@ -55,6 +58,7 @@ Transposes then broadcasts `src[0:1, :]` onto all partitions of `dst`. Each part - PSUM free dimension limit applies: `B` must be within the PSUM limit — 512 (gen2/3); gen4: 4096 for fp32, 8192 for bf16. **Example:** + ```python import nki.language as nl from nkilib.core.utils.tp_broadcast import tp_broadcast @@ -68,6 +72,7 @@ tp_broadcast(src=src, dst=dst, src_offset=0) ## Usage Examples ### Pattern 1: Broadcasting a column for attention score computation + ```python # Transpose a column from the key tensor and broadcast across query heads key_buf = nl.ndarray((128, seq_len), dtype=nl.float16, buffer=nl.sbuf) @@ -78,6 +83,7 @@ tp_broadcast(src=key_buf, dst=broadcast_key, src_offset=col_idx) ``` ### Pattern 2: Using explicit PSUM address to avoid conflicts + ```python # When other operations are using PSUM, specify a non-conflicting address tp_broadcast( diff --git a/skills/neuron-nki-writing/references/nkilib/patterns/layout-conversion.md b/skills/neuron-nki-writing/references/nkilib/patterns/layout-conversion.md index 9e757ed..19910e4 100644 --- a/skills/neuron-nki-writing/references/nkilib/patterns/layout-conversion.md +++ b/skills/neuron-nki-writing/references/nkilib/patterns/layout-conversion.md @@ -1,16 +1,17 @@ # Layout Conversion ## Overview + Layout conversion patterns for transforming between interleaved and contiguous memory layouts in the partition dimension, primarily used for Rotary Position Embedding (RoPE). These patterns use permutation matrices and SBUF matmul for efficient in-SBUF layout changes when tensor sizes are small enough. ## Quick Reference -| Function | Signature | Description | -|----------|-----------|-------------| -| `_compute_convert_to_interleaved_mat` | `(x_sb) -> nl.ndarray` | Generate permutation matrix for layout conversion | -| `_convert_from_interleaved` | `(x_sb, mat) -> nl.ndarray` | Interleaved to contiguous: `[e0,o0,e1,o1,...] -> [e0,e1,...,o0,o1,...]` | -| `_convert_to_interleaved` | `(x_sb, mat) -> nl.ndarray` | Contiguous to interleaved: `[e0,e1,...,o0,o1,...] -> [e0,o0,e1,o1,...]` | -| `RoPE_sbuf` | `(x_in_sb, cos_sb, sin_sb, x_out_sb, convert_from_interleaved) -> nl.ndarray` | Apply RoPE rotation entirely in SBUF | +| Function | Signature | Description | +| ------------------------------------- | ----------------------------------------------------------------------------- | ----------------------------------------------------------------------- | +| `_compute_convert_to_interleaved_mat` | `(x_sb) -> nl.ndarray` | Generate permutation matrix for layout conversion | +| `_convert_from_interleaved` | `(x_sb, mat) -> nl.ndarray` | Interleaved to contiguous: `[e0,o0,e1,o1,...] -> [e0,e1,...,o0,o1,...]` | +| `_convert_to_interleaved` | `(x_sb, mat) -> nl.ndarray` | Contiguous to interleaved: `[e0,e1,...,o0,o1,...] -> [e0,o0,e1,o1,...]` | +| `RoPE_sbuf` | `(x_in_sb, cos_sb, sin_sb, x_out_sb, convert_from_interleaved) -> nl.ndarray` | Apply RoPE rotation entirely in SBUF | ## Import Options @@ -18,6 +19,7 @@ Layout conversion patterns for transforming between interleaved and contiguous m See the "Full Source Implementation" section below, or the bundled source files in `references/nkilib/core/`. **If nkilib is installed** in the user's environment: + ```python from nkilib.core.embeddings.rope import ( RoPE_sbuf, @@ -37,21 +39,26 @@ Generate a permutation matrix P for converting between contiguous and interleave - `P^T @ X`: interleaved to contiguous: `[e0,o0,e1,o1,...] -> [e0,e1,...,o0,o1,...]` **Args:** + - `x_sb` (nl.ndarray): SBUF tensor with shape `[d_head, B, n_heads, S]` (used only for shape information) **Returns:** + - `nl.ndarray`: Permutation matrix of shape `[d_head, d_head]` in SBUF **Constraints:** + - `d_head` must be even - `B * n_heads * S <= nl.tile_size.gemm_moving_fmax` (required for nc_matmul) **Implementation Notes:** + - Builds the permutation matrix by applying strided access on an identity matrix - Uses `nisa.tensor_copy` with `scalar_engine` and strided access patterns - For d_head=4, the matrix maps: row 0->col 0, row 1->col 2, row 2->col 1, row 3->col 3 **Example:** + ```python import nki.language as nl @@ -67,17 +74,21 @@ convert_mat = _compute_convert_to_interleaved_mat(x_sb) Convert interleaved to contiguous layout using matrix multiplication: `P^T @ x_sb`. **Args:** + - `x_sb` (nl.ndarray): Input tensor in SBUF with shape `[d_head, B, n_heads, S]` in interleaved layout - `convert_to_interleaved_mat` (nl.ndarray): Permutation matrix from `_compute_convert_to_interleaved_mat` **Returns:** + - `nl.ndarray`: New SBUF tensor with shape `[d_head, B, n_heads, S]` in contiguous layout **Constraints:** + - Input must be in SBUF - `B * n_heads * S <= nl.tile_size.gemm_moving_fmax` **Implementation Notes:** + - Uses `nisa.nc_matmul` with the permutation matrix as stationary and x_sb reshaped to 2D as moving - Copies PSUM result back to SBUF via `nisa.activation` with `nl.copy` - Returns a new buffer (does not modify input) @@ -89,17 +100,21 @@ Convert interleaved to contiguous layout using matrix multiplication: `P^T @ x_s Convert contiguous to interleaved layout using matrix multiplication: `P @ x_sb`. **Args:** + - `x_sb` (nl.ndarray): Input tensor in SBUF with shape `[d_head, B, n_heads, S]` in contiguous layout - `convert_to_interleaved_mat` (nl.ndarray): Permutation matrix from `_compute_convert_to_interleaved_mat` **Returns:** + - `nl.ndarray`: Same buffer with interleaved layout applied in-place **Constraints:** + - Input must be in SBUF - `B * n_heads * S <= nl.tile_size.gemm_moving_fmax` **Implementation Notes:** + - Pre-transposes the permutation matrix (via `nisa.nc_transpose`) to compensate for `nc_matmul`'s implicit transpose of the stationary operand - Modifies input buffer in-place @@ -110,12 +125,14 @@ Convert contiguous to interleaved layout using matrix multiplication: `P @ x_sb` Apply Rotary Position Embedding entirely in SBUF, for megakernel fusion scenarios where data is already in SBUF. **RoPE Formula:** + ``` out[even] = x[even] * cos - x[odd] * sin out[odd] = x[odd] * cos + x[even] * sin ``` **Args:** + - `x_in_sb` (nl.ndarray): `[d_head, B, n_heads, S]` in SBUF - input embeddings - `cos_sb` (nl.ndarray): `[d_head//2, B, S]` in SBUF - cosine frequencies - `sin_sb` (nl.ndarray): `[d_head//2, B, S]` in SBUF - sine frequencies @@ -123,9 +140,11 @@ out[odd] = x[odd] * cos + x[even] * sin - `convert_from_interleaved` (bool): Convert from interleaved to contiguous layout before computation. Default: False **Returns:** + - `nl.ndarray`: `x_out_sb` with RoPE applied (modified in-place) **Constraints:** + - `d_head` must be 64 or 128 - `B` must be in (0, 64] - `S` must be in (0, 512] @@ -136,6 +155,7 @@ out[odd] = x[odd] * cos + x[even] * sin - For `convert_from_interleaved=True`: `B * n_heads * S <= nl.tile_size.gemm_moving_fmax` **Example:** + ```python import nki.language as nl @@ -154,6 +174,7 @@ RoPE_sbuf(x_in_sb, cos_sb, sin_sb, x_out_sb) ## Usage Examples ### Pattern 1: RoPE in a fused attention kernel + ```python import nki.isa as nisa import nki.language as nl @@ -172,6 +193,7 @@ def apply_rope_in_attention(q_sb, k_sb, cos_sb, sin_sb): ``` ### Pattern 2: Layout conversion for interleaved-format models + ```python import nki.isa as nisa import nki.language as nl @@ -195,6 +217,7 @@ def convert_layout_for_rope(x_sb): ``` ### Pattern 3: Standalone RoPE kernel with strided DMA fallback + ```python import nki.language as nl diff --git a/skills/neuron-nki-writing/references/nkilib/patterns/moe-patterns.md b/skills/neuron-nki-writing/references/nkilib/patterns/moe-patterns.md index 9aecbf4..472a3d9 100644 --- a/skills/neuron-nki-writing/references/nkilib/patterns/moe-patterns.md +++ b/skills/neuron-nki-writing/references/nkilib/patterns/moe-patterns.md @@ -1,19 +1,20 @@ # MoE Patterns ## Overview + Mixture of Experts (MoE) utility patterns for expert affinity computation, token index loading, block-expert mapping, and expert affinity gathering/broadcasting. These patterns are used across both CTE (Continuous Tensor Engine) and TKG (Token Generation Kernel) MoE implementations. ## Quick Reference -| Function | Module | Signature | Description | -|----------|--------|-----------|-------------| -| `load_block_expert` | CTE | `(block_to_expert, block_idx) -> nl.ndarray` | Load expert ID for current block | -| `load_token_indices` | CTE | `(token_position_to_id, block_idx, B, NUM_TILES) -> nl.ndarray` | Load and transpose token indices (static block) | -| `load_token_indices_dynamic_block` | CTE | `(token_position_to_id, block_idx, B, NUM_TILES, skip_dma) -> nl.ndarray` | Load token indices (dynamic block) | -| `calculate_expert_affinities` | CTE | `(expert_affinities_masked, token_indices, block_expert, E, NUM_TILES, dtype, ...) -> List` | Compute expert affinity scores per token | -| `stream_shuffle_broadcast` | CTE | `(src, dst) -> None` | Broadcast first partition across all partitions | -| `gather_expert_affinities` | TKG | `(expert_affinities_sb, expert_idx, dims, sbm) -> nl.ndarray` | Gather affinities via local_gather | -| `broadcast_token_affinity` | TKG | `(dst, gathered_affinities_sb, token_index, dims, sbm) -> nl.ndarray` | Broadcast per-token affinities across partitions | +| Function | Module | Signature | Description | +| ---------------------------------- | ------ | ------------------------------------------------------------------------------------------- | ------------------------------------------------ | +| `load_block_expert` | CTE | `(block_to_expert, block_idx) -> nl.ndarray` | Load expert ID for current block | +| `load_token_indices` | CTE | `(token_position_to_id, block_idx, B, NUM_TILES) -> nl.ndarray` | Load and transpose token indices (static block) | +| `load_token_indices_dynamic_block` | CTE | `(token_position_to_id, block_idx, B, NUM_TILES, skip_dma) -> nl.ndarray` | Load token indices (dynamic block) | +| `calculate_expert_affinities` | CTE | `(expert_affinities_masked, token_indices, block_expert, E, NUM_TILES, dtype, ...) -> List` | Compute expert affinity scores per token | +| `stream_shuffle_broadcast` | CTE | `(src, dst) -> None` | Broadcast first partition across all partitions | +| `gather_expert_affinities` | TKG | `(expert_affinities_sb, expert_idx, dims, sbm) -> nl.ndarray` | Gather affinities via local_gather | +| `broadcast_token_affinity` | TKG | `(dst, gathered_affinities_sb, token_index, dims, sbm) -> nl.ndarray` | Broadcast per-token affinities across partitions | ## Import Options @@ -21,6 +22,7 @@ Mixture of Experts (MoE) utility patterns for expert affinity computation, token See the "Full Source Implementation" section below, or the bundled source files in `references/nkilib/core/`. **If nkilib is installed** in the user's environment: + ```python # CTE MoE utilities from nkilib.core.moe.moe_cte.moe_cte_utils import ( @@ -46,18 +48,22 @@ from nkilib.core.moe.moe_tkg.moe_tkg_utils import ( Load the expert ID assigned to the current block from the block-to-expert mapping tensor. **Args:** + - `block_to_expert` (nl.ndarray): Mapping tensor of shape `[N, 1]` where N is number of blocks, containing expert indices - `block_idx` (int or nl.ndarray): Block index to load, either static integer or dynamic tensor value **Returns:** + - `nl.ndarray`: Expert ID tensor of shape `[1, 1]` in SBUF (int32) **Notes:** + - Handles both static (int) and dynamic (tensor) block indices - Uses `scalar_offset` for dynamic indices via temporary tensor - Result stored in SBUF for efficient access in subsequent operations **Example:** + ```python import nki.language as nl @@ -75,15 +81,18 @@ block_expert = load_block_expert(block_to_expert, block_idx=dynamic_idx_tensor) Load and transpose token indices for the current block using static block indexing. **Args:** + - `token_position_to_id` (nl.ndarray): Token position mapping of shape `[N*B]` - `block_idx` (int): Current block index (static) - `B` (int): Block size (number of tokens per block) - `NUM_TILES` (int): Number of tiles (`B // TILE_SIZE`) **Returns:** + - `nl.ndarray`: Transposed token indices of shape `[TILE_SIZE, NUM_TILES]` in SBUF (int32) **Notes:** + - Uses `dma_transpose` for efficient layout transformation - Tokens are distributed across the partition dimension for efficient vector DGE @@ -94,6 +103,7 @@ Load and transpose token indices for the current block using static block indexi Load token indices when block_idx is a dynamic tensor value (runtime-determined). **Args:** + - `token_position_to_id` (nl.ndarray): Token position mapping tensor - `block_idx` (nl.ndarray): Dynamic block index tensor - `B` (int): Block size (number of tokens per block) @@ -101,9 +111,11 @@ Load token indices when block_idx is a dynamic tensor value (runtime-determined) - `skip_dma` (SkipMode): DMA skip configuration **Returns:** + - `nl.ndarray`: Token indices of shape `[TILE_SIZE, NUM_TILES]` in SBUF **Notes:** + - Reshapes `token_position_to_id` to `[total_size//B, B]` for indexing - Uses `scalar_offset` with indirect_dim for dynamic block addressing - Memsets to zero when `skip_dma.skip_token` is True (for out-of-bounds tokens) @@ -115,17 +127,21 @@ Load token indices when block_idx is a dynamic tensor value (runtime-determined) Broadcast the first partition of src onto all partitions of dst. **Args:** + - `src` (nl.ndarray): 2D input tensor in SBUF - `dst` (nl.ndarray): 2D output tensor in SBUF (final dim must match src) **Returns:** + - None: Broadcasts src to dst in-place **Notes:** + - Uses `nisa.nc_stream_shuffle` with a zero shuffle mask to replicate partition 0 - Processes in banks of 32 partitions **Example:** + ```python import nki.language as nl @@ -142,6 +158,7 @@ stream_shuffle_broadcast(src=scalar, dst=broadcasted) Calculate expert affinity scores for tokens in the current block using indirect addressing. **Args:** + - `expert_affinities_masked` (nl.ndarray): Expert affinities tensor of shape `[(T+1)*E, 1]` - `token_indices` (nl.ndarray): Token indices of shape `[TILE_SIZE, NUM_TILES]` in SBUF - `block_expert` (nl.ndarray): Expert ID of shape `[1, 1]` in SBUF @@ -152,9 +169,11 @@ Calculate expert affinity scores for tokens in the current block using indirect - `token_indices_offset` (int): Offset for block tiling. Default: 0 **Returns:** + - `List[nl.ndarray]`: List of expert affinity tensors in SBUF, one per tile, each shape `[TILE_SIZE, 1]` in float32 **Notes:** + - Uses pointer arithmetic: `addr = token_indices * E + block_expert` - Broadcasts `block_expert` to all partitions via `stream_shuffle_broadcast` - Performs indirect load from `expert_affinities_masked` using `vector_offset` @@ -167,19 +186,23 @@ Calculate expert affinity scores for tokens in the current block using indirect Gather expert affinities based on expert indices using `local_gather` operation (TKG path). **Args:** + - `expert_affinities_sb` (nl.ndarray): `[_pmax, E]` expert affinities in SBUF - `expert_idx` (nl.ndarray): `[T, K]` expert indices for each token - `dims` (MLPTKGConstantsDimensionSizes): Dimension sizes object - `sbm` (SbufManager): SBUF memory manager **Returns:** + - `nl.ndarray`: `[_pmax, 16, 16]` gathered affinities tensor **Constraints:** + - `K <= 16` (PARTITIONS_PER_CORE) - `E > 1` (local_gather requires src_buffer_size > 1) **Notes:** + - Uses different strategies for `T <= 16` (nc_transpose path) vs `T > 16` (dma_transpose path) - Converts expert indices to uint16 for `local_gather` @@ -190,6 +213,7 @@ Gather expert affinities based on expert indices using `local_gather` operation Broadcast expert affinities for a specific token across all partitions (TKG path). **Args:** + - `dst` (nl.ndarray): Destination tensor for broadcasted affinities - `gathered_affinities_sb` (nl.ndarray): `[_pmax, 16, 16]` gathered affinities - `token_index` (int): Index of the current token @@ -197,9 +221,11 @@ Broadcast expert affinities for a specific token across all partitions (TKG path - `sbm` (SbufManager): SBUF memory manager **Returns:** + - `nl.ndarray`: `[_pmax, K]` broadcasted token affinities **Notes:** + - Computes partition and quadrant positions from token_index - Uses `nc_stream_shuffle` for partition alignment - Uses `stream_shuffle_broadcast` for final broadcast @@ -207,6 +233,7 @@ Broadcast expert affinities for a specific token across all partitions (TKG path ## Usage Examples ### Pattern 1: CTE MoE block processing loop + ```python import nki.language as nl @@ -237,6 +264,7 @@ def process_moe_block(block_to_expert, token_position_to_id, expert_affinities, ``` ### Pattern 2: Dynamic block processing with skip mode + ```python import nki.language as nl @@ -262,6 +290,7 @@ def process_dynamic_block(block_to_expert, token_position_to_id, ``` ### Pattern 3: TKG expert affinity gathering and broadcasting + ```python import nki.language as nl diff --git a/skills/neuron-nki-writing/references/nkilib/patterns/normalization-patterns.md b/skills/neuron-nki-writing/references/nkilib/patterns/normalization-patterns.md index d0e3cf8..856c691 100644 --- a/skills/neuron-nki-writing/references/nkilib/patterns/normalization-patterns.md +++ b/skills/neuron-nki-writing/references/nkilib/patterns/normalization-patterns.md @@ -1,14 +1,15 @@ # Normalization Patterns ## Overview + Reusable patterns for normalization kernel data loading and shape validation in token-generation mode. These utilities handle the HBM-to-SBUF data movement and layout transformations required for RMSNorm/LayerNorm operations, including support for hidden-dimension sharding and transpose loading. ## Quick Reference -| Function | Signature | Description | -|----------|-----------|-------------| -| `validate_shapes` | `(input_view, gamma_view, output_view) -> (BxS, H, H0, H1)` | Validate and extract normalization tensor dimensions | -| `load_input_to_sbuf` | `(input_hbm, input_sb, num_H_shards, hidden_dim_tp) -> TensorView` | Load input from HBM to SBUF with layout transformation | +| Function | Signature | Description | +| -------------------- | ------------------------------------------------------------------ | -------------------------------------------------------------- | +| `validate_shapes` | `(input_view, gamma_view, output_view) -> (BxS, H, H0, H1)` | Validate and extract normalization tensor dimensions | +| `load_input_to_sbuf` | `(input_hbm, input_sb, num_H_shards, hidden_dim_tp) -> TensorView` | Load input from HBM to SBUF with layout transformation | | `load_gamma_to_sbuf` | `(gamma_hbm, gamma_sb, num_H_shards, hidden_dim_tp) -> TensorView` | Load gamma weights from HBM to SBUF with layout transformation | ## Import Options @@ -17,6 +18,7 @@ Reusable patterns for normalization kernel data loading and shape validation in See the "Full Source Implementation" section below, or the bundled source files in `references/nkilib/core/`. **If nkilib is installed** in the user's environment: + ```python from nkilib.core.subkernels.norm_tkg_utils import validate_shapes, load_input_to_sbuf, load_gamma_to_sbuf ``` @@ -28,11 +30,13 @@ from nkilib.core.subkernels.norm_tkg_utils import validate_shapes, load_input_to Validate tensor shapes for normalization operations. Handles both HBM inputs (shape `[B, S, H]`) and SBUF inputs (shape `[H0, BxS, H1]`). **Args:** + - `input_view` (TensorView): Input tensor view. HBM shape `[B, S, H]` or SBUF shape `[H0, BxS, H1]` - `gamma_view` (TensorView): Gamma tensor view with shape `[1, H]` - `output_view` (TensorView): Output tensor view with expected shape `[H0, BxS, H1]` **Returns:** + - `Tuple[int, int, int, int]`: `(BxS, H, H0, H1)` dimensions where: - `BxS`: Flattened batch-times-sequence dimension - `H`: Full hidden dimension @@ -40,12 +44,14 @@ Validate tensor shapes for normalization operations. Handles both HBM inputs (sh - `H1`: Hidden dimension tiles (`H // H0`) **Constraints:** + - `H0` must equal `nl.tile_size.pmax` (128) - `H` must be divisible by `H0` - Output shape must be `[H0, BxS, H1]` - Gamma shape must be `[1, H]` **Example:** + ```python from nkilib.core.utils.tensor_view import TensorView @@ -64,25 +70,30 @@ BxS, H, H0, H1 = validate_shapes(input_view, gamma_view, output_view) Load input data from HBM to SBUF with appropriate layout transformation. Supports two loading strategies depending on the hidden dimension layout. **Args:** + - `input_hbm` (TensorView): Input tensor view in HBM with shape `[BxS, H]` - `input_sb` (TensorView): Destination buffer in SBUF with shape `[H0, BxS, H1]` - `num_H_shards` (int): Number of shards along the H dimension - `hidden_dim_tp` (bool): If True, use transpose load for `(H/128, 128)` layout. Default: False **Returns:** + - `TensorView`: Input tensor view in SBUF with shape `[H0, BxS, H1]` **Constraints:** + - `H0` is always `nl.tile_size.pmax` (128) - `H` must be divisible by `H0` - `H1` must be divisible by `num_H_shards` **Notes:** + - `hidden_dim_tp=True`: Uses `dma_transpose` for `(BxS, H) -> (BxS*H1, H0) -> (H0, BxS, H1)` transformation - `hidden_dim_tp=False`: Uses `dma_copy` with permutation for `(BxS, H) -> (BxS, num_H_shards, H0, H2) -> (H0, BxS, num_H_shards, H2)` transformation - Static DMA mode (`_DGE_MODE_NONE = 3`) is used for the non-transpose path **Example:** + ```python import nki.language as nl @@ -106,24 +117,29 @@ input_view = load_input_to_sbuf( Load gamma (scale) weights from HBM to SBUF with appropriate layout transformation. Follows the same layout strategy as `load_input_to_sbuf` but for 1D gamma vectors. **Args:** + - `gamma_hbm` (TensorView): Gamma tensor view in HBM with shape `[1, H]` - `gamma_sb` (TensorView): Destination buffer in SBUF with shape `[H0, H1]` - `num_H_shards` (int): Number of shards along the H dimension - `hidden_dim_tp` (bool): If True, use transpose load. Default: False **Returns:** + - `TensorView`: Gamma tensor view in SBUF with shape `[H0, H1]` **Constraints:** + - Gamma must have shape `[1, H]` in HBM - `H` must be divisible by `H0` (128) - `H1` must be divisible by `num_H_shards` **Notes:** + - `hidden_dim_tp=True`: Transpose load `(H) -> (H1, H0) -> (H0, H1)` - `hidden_dim_tp=False`: Standard layout `(H) -> (num_H_shards, H0, H2) -> (H0, num_H_shards, H2)` **Example:** + ```python import nki.language as nl @@ -142,6 +158,7 @@ gamma_view = load_gamma_to_sbuf( ## Usage Examples ### Pattern 1: Standard normalization data preparation + ```python import nki.language as nl from nkilib.core.utils.tensor_view import TensorView @@ -172,6 +189,7 @@ def prepare_norm_inputs(input_tensor, gamma_tensor, batch_size, seq_len, hidden_ ``` ### Pattern 2: Sharded normalization with LNC + ```python import nki.language as nl from nkilib.core.utils.tensor_view import TensorView @@ -210,6 +228,7 @@ def prepare_sharded_norm(input_tensor, gamma_tensor, hidden_size, num_shards=2): ## Source See `references/nkilib/core/subkernels/` for full implementations: + - `rmsnorm_tkg.py` — RMSNorm kernel - `layernorm_tkg.py` — LayerNorm kernel - `norm_tkg_utils.py` — Shared normalization utilities diff --git a/skills/neuron-nki-writing/references/nkilib/patterns/quantization-helpers.md b/skills/neuron-nki-writing/references/nkilib/patterns/quantization-helpers.md index b95b93c..f767cf0 100644 --- a/skills/neuron-nki-writing/references/nkilib/patterns/quantization-helpers.md +++ b/skills/neuron-nki-writing/references/nkilib/patterns/quantization-helpers.md @@ -1,16 +1,17 @@ # Quantization Helpers ## Overview + FP8 dtype detection and quantization-compatible dtype selection patterns extracted from attention and MoE kernel implementations. Use these when writing kernels that need to handle FP8 quantized inputs or select compute dtypes based on hardware generation. ## Quick Reference -| Function | Signature | Description | -|----------|-----------|-------------| -| `is_fp8_e4m3` | `(dtype) -> bool` | Check if dtype is FP8 E4M3 format | -| `is_fp8_e5m2` | `(dtype) -> bool` | Check if dtype is FP8 E5M2 format | +| Function | Signature | Description | +| ------------------ | ------------------------- | ------------------------------------------------- | +| `is_fp8_e4m3` | `(dtype) -> bool` | Check if dtype is FP8 E4M3 format | +| `is_fp8_e5m2` | `(dtype) -> bool` | Check if dtype is FP8 E5M2 format | | `compatible_dtype` | `(compute_type) -> dtype` | Return gen3+-compatible dtype or float32 fallback | -| `div_ceil` | `(n, d) -> int` | Ceiling division helper | +| `div_ceil` | `(n, d) -> int` | Ceiling division helper | ## Import Options @@ -18,6 +19,7 @@ FP8 dtype detection and quantization-compatible dtype selection patterns extract See the "Full Source Implementation" section below, or the bundled source files in `references/nkilib/core/`. **If nkilib is installed** in the user's environment: + ```python # FP8 detection (from attention utils) from nkilib.core.attention.attention_tkg_utils import is_fp8_e4m3, is_fp8_e5m2 @@ -33,15 +35,19 @@ from nkilib.core.moe.moe_cte.moe_cte_utils import compatible_dtype, div_ceil Check if a dtype is FP8 E4M3 format, handling both numpy dtype objects and compiler internal name strings. **Args:** + - `dtype`: A data type value (e.g., `nl.float8_e4m3` or compiler internal string) **Returns:** + - `bool`: True if dtype is FP8 E4M3 **Constraints:** + - Handles two representations: `nl.float8_e4m3` object equality and `"float8e4"` string comparison **Example:** + ```python import nki.language as nl @@ -57,15 +63,19 @@ if is_fp8_e4m3(input_tensor.dtype): Check if a dtype is FP8 E5M2 format, handling both numpy dtype objects and compiler internal name strings. **Args:** + - `dtype`: A data type value (e.g., `nl.float8_e5m2` or compiler internal string) **Returns:** + - `bool`: True if dtype is FP8 E5M2 **Constraints:** + - Handles two representations: `nl.float8_e5m2` object equality and `"float8e5"` string comparison **Example:** + ```python import nki.language as nl @@ -81,16 +91,20 @@ if is_fp8_e5m2(weight.dtype): Return a compute-compatible dtype based on the current NeuronCore version. On gen3+ hardware, returns the requested dtype directly. On gen2 (Trn1/Inf2), falls back to `nl.float32` since bfloat16 compute is not fully supported for all operations. **Args:** + - `compute_type`: Desired compute dtype (e.g., `nl.bfloat16`) **Returns:** + - `dtype`: `compute_type` on gen3+, `nl.float32` on gen2 **Constraints:** + - Requires `nki.isa` for `get_nc_version()` and `nc_version.gen3` - Only meaningful at trace time (compile time) **Example:** + ```python import nki.isa as nisa import nki.language as nl @@ -107,13 +121,16 @@ intermediate = nl.ndarray((128, 512), dtype=dtype, buffer=nl.sbuf) Integer ceiling division. Returns the smallest integer >= n/d. **Args:** + - `n` (int): Numerator - `d` (int): Denominator **Returns:** + - `int`: Ceiling of n/d **Example:** + ```python num_tiles = div_ceil(seq_len, 128) # Number of 128-element tiles needed ``` @@ -121,6 +138,7 @@ num_tiles = div_ceil(seq_len, 128) # Number of 128-element tiles needed ## Usage Examples ### Pattern 1: FP8-aware weight loading + ```python import nki.language as nl import nki.isa as nisa @@ -140,6 +158,7 @@ def load_weights_with_dequant(weights, scale, compute_dtype): ``` ### Pattern 2: Generation-aware compute dtype selection + ```python import nki.isa as nisa import nki.language as nl @@ -158,6 +177,7 @@ def setup_compute_buffers(hidden_size, seq_len): ``` ### Pattern 3: Conditional FP8 scaling in MoE kernels + ```python import nki.isa as nisa import nki.language as nl diff --git a/skills/neuron-nki-writing/references/nkilib/types/common-types.md b/skills/neuron-nki-writing/references/nkilib/types/common-types.md index 6e0474e..146dbbd 100644 --- a/skills/neuron-nki-writing/references/nkilib/types/common-types.md +++ b/skills/neuron-nki-writing/references/nkilib/types/common-types.md @@ -6,15 +6,15 @@ Enum types used across NKI kernel configurations for specifying output layouts, ## Quick Reference -| Enum | Description | -|------|-------------| -| `QKVOutputLayout` | Output tensor layout for QKV projections | -| `NormType` | Normalization type selection (none, RMS, LayerNorm) | -| `ActFnType` | Activation function type (SiLU, GELU, Swish) | -| `RouterActFnType` | Activation type for MoE router TopK kernel | -| `ExpertAffinityScaleMode` | Scaling mode for MoE expert affinity scores | -| `QuantizationType` | Quantization strategy (none, static, row, MX) | -| `GateUpDim` | Index selector for gate/up projection in MLP | +| Enum | Description | +| ------------------------- | --------------------------------------------------- | +| `QKVOutputLayout` | Output tensor layout for QKV projections | +| `NormType` | Normalization type selection (none, RMS, LayerNorm) | +| `ActFnType` | Activation function type (SiLU, GELU, Swish) | +| `RouterActFnType` | Activation type for MoE router TopK kernel | +| `ExpertAffinityScaleMode` | Scaling mode for MoE expert affinity scores | +| `QuantizationType` | Quantization strategy (none, static, row, MX) | +| `GateUpDim` | Index selector for gate/up projection in MLP | ## Import Options @@ -22,6 +22,7 @@ Enum types used across NKI kernel configurations for specifying output layouts, Source: `references/nkilib/core/utils/common_types.py` **If nkilib is installed** in the user's environment: + ```python from nkilib.core.utils.common_types import QKVOutputLayout, NormType, ActFnType from nkilib.core.utils.common_types import RouterActFnType, ExpertAffinityScaleMode @@ -34,13 +35,14 @@ from nkilib.core.utils.common_types import QuantizationType, GateUpDim Specifies the memory layout for QKV (Query/Key/Value) projection outputs. -| Value | Int | Layout Shape | Description | -|-------|-----|-------------|-------------| -| `BSD` | 0 | `(b, s, (n_q_heads + 2 * n_kv_heads) * d_head)` | Batch-Sequence-Dim interleaved layout | -| `NBSd` | 1 | `(num_heads, b, s, d_head)` | Heads-first with sequence-major inner layout | -| `NBdS` | 2 | `(num_heads, b, d_head, s)` | Heads-first with head-dim-major inner layout | +| Value | Int | Layout Shape | Description | +| ------ | --- | ----------------------------------------------- | -------------------------------------------- | +| `BSD` | 0 | `(b, s, (n_q_heads + 2 * n_kv_heads) * d_head)` | Batch-Sequence-Dim interleaved layout | +| `NBSd` | 1 | `(num_heads, b, s, d_head)` | Heads-first with sequence-major inner layout | +| `NBdS` | 2 | `(num_heads, b, d_head, s)` | Heads-first with head-dim-major inner layout | **Example:** + ```python from nkilib.core.utils.common_types import QKVOutputLayout @@ -54,14 +56,15 @@ if layout == QKVOutputLayout.BSD: Specifies the normalization method to apply. -| Value | Int | Description | -|-------|-----|-------------| -| `NO_NORM` | 0 | No normalization applied | -| `RMS_NORM` | 1 | Root Mean Square normalization | -| `LAYER_NORM` | 2 | Layer normalization (mean + variance) | -| `RMS_NORM_SKIP_GAMMA` | 3 | RMS normalization without the gamma scaling parameter | +| Value | Int | Description | +| --------------------- | --- | ----------------------------------------------------- | +| `NO_NORM` | 0 | No normalization applied | +| `RMS_NORM` | 1 | Root Mean Square normalization | +| `LAYER_NORM` | 2 | Layer normalization (mean + variance) | +| `RMS_NORM_SKIP_GAMMA` | 3 | RMS normalization without the gamma scaling parameter | **Example:** + ```python from nkilib.core.utils.common_types import NormType @@ -75,14 +78,15 @@ if norm == NormType.RMS_NORM_SKIP_GAMMA: Specifies the activation function for MLP/FFN layers. -| Value | Int | Description | -|-------|-----|-------------| -| `SiLU` | 0 | Sigmoid Linear Unit (x * sigmoid(x)) | -| `GELU` | 1 | Gaussian Error Linear Unit | -| `GELU_Tanh_Approx` | 2 | GELU with tanh approximation | -| `Swish` | 3 | Swish activation (same as SiLU with beta=1) | +| Value | Int | Description | +| ------------------ | --- | ------------------------------------------- | +| `SiLU` | 0 | Sigmoid Linear Unit (x \* sigmoid(x)) | +| `GELU` | 1 | Gaussian Error Linear Unit | +| `GELU_Tanh_Approx` | 2 | GELU with tanh approximation | +| `Swish` | 3 | Swish activation (same as SiLU with beta=1) | **Example:** + ```python from nkilib.core.utils.common_types import ActFnType @@ -93,14 +97,15 @@ act_fn = ActFnType.SiLU # Used in LLaMA-style models Specifies the activation type for Mixture-of-Experts (MoE) router TopK kernel. -| Value | Int | Description | -|-------|-----|-------------| -| `SIGMOID` | 0 | Sigmoid activation for routing scores | -| `SOFTMAX` | 1 | Softmax activation for routing scores | +| Value | Int | Description | +| --------- | --- | ------------------------------------- | +| `SIGMOID` | 0 | Sigmoid activation for routing scores | +| `SOFTMAX` | 1 | Softmax activation for routing scores | Implements `__str__` returning the lowercase name (e.g., `"sigmoid"`, `"softmax"`). **Example:** + ```python from nkilib.core.utils.common_types import RouterActFnType @@ -112,14 +117,15 @@ print(router_act) # prints: "softmax" Controls when and how expert affinity scores are scaled in MoE routing. -| Value | Int | Description | -|-------|-----|-------------| -| `NO_SCALE` | 0 | No scaling applied to affinity scores | -| `POST_SCALE` | 1 | Scale applied after expert selection | -| `PRE_SCALE` | 2 | Scale applied before expert selection | -| `PRE_SCALE_DELAYED` | 3 | Pre-scaling with delayed application | +| Value | Int | Description | +| ------------------- | --- | ------------------------------------- | +| `NO_SCALE` | 0 | No scaling applied to affinity scores | +| `POST_SCALE` | 1 | Scale applied after expert selection | +| `PRE_SCALE` | 2 | Scale applied before expert selection | +| `PRE_SCALE_DELAYED` | 3 | Pre-scaling with delayed application | **Example:** + ```python from nkilib.core.utils.common_types import ExpertAffinityScaleMode @@ -130,14 +136,15 @@ scale_mode = ExpertAffinityScaleMode.POST_SCALE Specifies the quantization strategy for weight or activation tensors. -| Value | Int | Description | -|-------|-----|-------------| -| `NONE` | 0 | No quantization (full precision) | -| `STATIC` | 1 | Static quantization with fixed scale factors | -| `ROW` | 2 | Per-row quantization with individual scale factors | -| `MX` | 3 | Microscaling (MX) quantization format | +| Value | Int | Description | +| -------- | --- | -------------------------------------------------- | +| `NONE` | 0 | No quantization (full precision) | +| `STATIC` | 1 | Static quantization with fixed scale factors | +| `ROW` | 2 | Per-row quantization with individual scale factors | +| `MX` | 3 | Microscaling (MX) quantization format | **Example:** + ```python from nkilib.core.utils.common_types import QuantizationType @@ -151,12 +158,13 @@ if quant != QuantizationType.NONE: Index selector for the gate and up projections in gated MLP architectures (e.g., SwiGLU). -| Value | Int | Description | -|-------|-----|-------------| -| `GATE` | 0 | Index for the gate projection | -| `UP` | 1 | Index for the up projection | +| Value | Int | Description | +| ------ | --- | ----------------------------- | +| `GATE` | 0 | Index for the gate projection | +| `UP` | 1 | Index for the up projection | **Example:** + ```python from nkilib.core.utils.common_types import GateUpDim @@ -168,6 +176,7 @@ up_weight = combined_weights[GateUpDim.UP.value] ## Usage Examples ### Pattern 1: Configuring a fused QKV + normalization kernel + ```python from nkilib.core.utils.common_types import QKVOutputLayout, NormType @@ -178,6 +187,7 @@ def launch_qkv_kernel(input_tensor, weights, config): ``` ### Pattern 2: Selecting MoE router parameters + ```python from nkilib.core.utils.common_types import RouterActFnType, ExpertAffinityScaleMode @@ -189,6 +199,7 @@ router_config = { ``` ### Pattern 3: Quantization-aware kernel dispatch + ```python from nkilib.core.utils.common_types import QuantizationType diff --git a/skills/neuron-nki-writing/references/nkilib/types/logging.md b/skills/neuron-nki-writing/references/nkilib/types/logging.md index 780c23f..eab965c 100644 --- a/skills/neuron-nki-writing/references/nkilib/types/logging.md +++ b/skills/neuron-nki-writing/references/nkilib/types/logging.md @@ -6,14 +6,14 @@ Lightweight logging system for NKI kernels with environment-based configuration ## Quick Reference -| Name | Type | Description | -|------|------|-------------| -| `LogLevel` | Enum | Log severity levels: DEBUG, INFO, WARN, ERROR, OFF | -| `Logger` | Class | Core logger with level-filtered `debug`/`info`/`warn`/`error` methods | -| `get_logger(name, level)` | Function | Factory function that creates a logger respecting env var overrides | -| `logger` | Instance | Pre-configured global logger instance for quick use | -| `LogEntry` | Dataclass | Buffered log entry for tree-style printing | -| `TreeLogger` | Class | Buffers log entries and prints them in tree format | +| Name | Type | Description | +| ------------------------- | --------- | --------------------------------------------------------------------- | +| `LogLevel` | Enum | Log severity levels: DEBUG, INFO, WARN, ERROR, OFF | +| `Logger` | Class | Core logger with level-filtered `debug`/`info`/`warn`/`error` methods | +| `get_logger(name, level)` | Function | Factory function that creates a logger respecting env var overrides | +| `logger` | Instance | Pre-configured global logger instance for quick use | +| `LogEntry` | Dataclass | Buffered log entry for tree-style printing | +| `TreeLogger` | Class | Buffers log entries and prints them in tree format | ## Import Options @@ -21,6 +21,7 @@ Lightweight logging system for NKI kernels with environment-based configuration Source: `references/nkilib/core/utils/logging.py` **If nkilib is installed** in the user's environment: + ```python from nkilib.core.utils.logging import get_logger, Logger, LogLevel, logger from nkilib.core.utils.tree_logger import TreeLogger, LogEntry @@ -32,17 +33,18 @@ from nkilib.core.utils.tree_logger import TreeLogger, LogEntry Log severity levels controlling which messages are emitted. -| Value | Int | Description | -|-------|-----|-------------| -| `DEBUG` | 0 | Detailed diagnostic information | -| `INFO` | 1 | General operational information (default level) | -| `WARN` | 2 | Warning conditions | -| `ERROR` | 3 | Error conditions | -| `OFF` | 999 | Suppress all log output | +| Value | Int | Description | +| ------- | --- | ----------------------------------------------- | +| `DEBUG` | 0 | Detailed diagnostic information | +| `INFO` | 1 | General operational information (default level) | +| `WARN` | 2 | Warning conditions | +| `ERROR` | 3 | Error conditions | +| `OFF` | 999 | Suppress all log output | **Static Method:** #### `LogLevel.from_string(level: str) -> LogLevel` + Converts a string name (e.g., `"DEBUG"`, `"INFO"`) to the corresponding `LogLevel` enum value. Raises `KeyError` for invalid strings. --- @@ -55,24 +57,29 @@ Core logging class. Extends `nl.NKIObject` for NKI compatibility. **Args:** -| Parameter | Type | Default | Description | -|-----------|------|---------|-------------| -| `name` | `str` | (required) | Logger name, displayed in `[name]` prefix | -| `level` | `LogLevel` | `LogLevel.INFO` | Minimum severity level to emit | +| Parameter | Type | Default | Description | +| --------- | ---------- | --------------- | ----------------------------------------- | +| `name` | `str` | (required) | Logger name, displayed in `[name]` prefix | +| `level` | `LogLevel` | `LogLevel.INFO` | Minimum severity level to emit | #### `Logger.debug(msg: str)` + Log a message at DEBUG level. Output format: `[DEBUG] [name] msg` #### `Logger.info(msg: str)` + Log a message at INFO level. Output format: `[INFO] [name] msg` #### `Logger.warn(msg: str)` + Log a message at WARN level. Output format: `[WARN] [name] msg` #### `Logger.error(msg: str)` + Log a message at ERROR level. Output format: `[ERROR] [name] msg` #### `Logger.is_enabled_for(level: LogLevel) -> bool` + Check if a given level would be logged. Useful to guard expensive message construction. ```python @@ -87,6 +94,7 @@ if my_logger.is_enabled_for(LogLevel.DEBUG): Factory function that creates a `Logger` with environment-variable-aware level resolution. **Priority Order (highest to lowest):** + 1. `NKILIB_LOG_LEVEL_` -- Per-logger env var override 2. `NKILIB_LOG_LEVEL` -- Global env var override 3. `level` parameter -- Code-specified default @@ -94,18 +102,20 @@ Factory function that creates a `Logger` with environment-variable-aware level r **Args:** -| Parameter | Type | Default | Description | -|-----------|------|---------|-------------| -| `name` | `str` | (required) | Logger name for output prefix and env var matching | -| `level` | `LogLevel` | `LogLevel.INFO` | Code-level default if no env overrides are found | +| Parameter | Type | Default | Description | +| --------- | ---------- | --------------- | -------------------------------------------------- | +| `name` | `str` | (required) | Logger name for output prefix and env var matching | +| `level` | `LogLevel` | `LogLevel.INFO` | Code-level default if no env overrides are found | **Returns:** Configured `Logger` instance. **Environment Variables:** + - `NKILIB_LOG_LEVEL=` -- Set default level for all loggers (e.g., `NKILIB_LOG_LEVEL=DEBUG`) - `NKILIB_LOG_LEVEL_=` -- Override a specific logger (e.g., `NKILIB_LOG_LEVEL_SBM=DEBUG`) **Example:** + ```python from nkilib.core.utils.logging import get_logger, LogLevel @@ -138,12 +148,12 @@ Buffered log entry used by `TreeLogger` for tree-style output. Extends `nl.NKIOb **Fields:** -| Field | Type | Default | Description | -|-------|------|---------|-------------| -| `msg` | `str` | (required) | Log message text | -| `depth` | `int` | (required) | Nesting depth in the tree (0 = root) | -| `is_stack` | `bool` | (required) | `True` for stack allocations, `False` for heap | -| `is_scope_boundary` | `bool` | `False` | `True` for scope open/close entries | +| Field | Type | Default | Description | +| ------------------- | ------ | ---------- | ---------------------------------------------- | +| `msg` | `str` | (required) | Log message text | +| `depth` | `int` | (required) | Nesting depth in the tree (0 = root) | +| `is_stack` | `bool` | (required) | `True` for stack allocations, `False` for heap | +| `is_scope_boundary` | `bool` | `False` | `True` for scope open/close entries | --- @@ -155,21 +165,23 @@ Buffers log entries and prints them in a tree-formatted structure with box-drawi **Args:** -| Parameter | Type | Description | -|-----------|------|-------------| -| `name` | `str` | Name displayed in the tree header | -| `logger` | `Logger` | Parent logger instance used for the header line | +| Parameter | Type | Description | +| --------- | -------- | ----------------------------------------------- | +| `name` | `str` | Name displayed in the tree header | +| `logger` | `Logger` | Parent logger instance used for the header line | #### `TreeLogger.log(msg: str, depth: int, is_scope_boundary: bool = False)` + Add a log entry to the buffer at the specified depth. -| Parameter | Type | Default | Description | -|-----------|------|---------|-------------| -| `msg` | `str` | (required) | Message text | -| `depth` | `int` | (required) | Tree nesting depth (0 = root level) | -| `is_scope_boundary` | `bool` | `False` | Mark as scope boundary entry | +| Parameter | Type | Default | Description | +| ------------------- | ------ | ---------- | ----------------------------------- | +| `msg` | `str` | (required) | Message text | +| `depth` | `int` | (required) | Tree nesting depth (0 = root level) | +| `is_scope_boundary` | `bool` | `False` | Mark as scope boundary entry | #### `TreeLogger.flush()` + Print all buffered entries in tree format, then clear the buffer. Output format: ``` @@ -182,6 +194,7 @@ Print all buffered entries in tree format, then clear the buffer. Output format: ``` **Example:** + ```python from nkilib.core.utils.logging import get_logger from nkilib.core.utils.tree_logger import TreeLogger @@ -205,6 +218,7 @@ tree.flush() ## Usage Examples ### Pattern 1: Named logger with environment override + ```python from nkilib.core.utils.logging import get_logger, LogLevel @@ -218,6 +232,7 @@ log.debug("This won't show at INFO level") ``` ### Pattern 2: Guard expensive debug messages + ```python from nkilib.core.utils.logging import get_logger, LogLevel @@ -228,6 +243,7 @@ if log.is_enabled_for(LogLevel.DEBUG): ``` ### Pattern 3: Tree logger for allocation visualization + ```python from nkilib.core.utils.logging import get_logger from nkilib.core.utils.tree_logger import TreeLogger diff --git a/skills/neuron-nki-writing/references/performance-basics.md b/skills/neuron-nki-writing/references/performance-basics.md index 7219138..6aa6a9b 100644 --- a/skills/neuron-nki-writing/references/performance-basics.md +++ b/skills/neuron-nki-writing/references/performance-basics.md @@ -22,6 +22,7 @@ for i in range(p_size): ``` **Guidelines:** + - Transfer entire tiles at once, not element by element - Align tile sizes to hardware boundaries (P=128, common F sizes: 512, 2048) - Use reshape to enable contiguous access patterns @@ -84,12 +85,12 @@ for f_idx in nl.affine_range(num_f_tiles): Choose the right loop type for your access pattern. -| Loop Type | When to Use | Unrolling | -|-----------|-------------|-----------| -| `nl.affine_range(N)` | Independent iterations, no dependencies between iterations | Full unroll | -| `nl.sequential_range(N)` | Loop-carried dependencies (e.g., cumsum, running max) | No unroll | -| `nl.static_range(N)` | Small constant N, want partial unroll control | Configurable | -| `TiledRange(total, tile)` | Partition dimension tiling with edge handling | Full unroll | +| Loop Type | When to Use | Unrolling | +| ------------------------- | ---------------------------------------------------------- | ------------ | +| `nl.affine_range(N)` | Independent iterations, no dependencies between iterations | Full unroll | +| `nl.sequential_range(N)` | Loop-carried dependencies (e.g., cumsum, running max) | No unroll | +| `nl.static_range(N)` | Small constant N, want partial unroll control | Configurable | +| `TiledRange(total, tile)` | Partition dimension tiling with edge handling | Full unroll | See [tiled-range.md](nkilib/core/tiled-range.md) for full TiledRange API documentation. @@ -111,11 +112,11 @@ then apply these optimizations based on profiling data. ## Memory Hierarchy Performance -| Memory | Bandwidth | Latency | Use For | -|--------|-----------|---------|---------| -| SBUF | Highest | Lowest | Active compute | -| PSUM | High | Low | MatMul accumulation | -| HBM | Lower | Higher | Input/output storage | +| Memory | Bandwidth | Latency | Use For | +| ------ | --------- | ------- | -------------------- | +| SBUF | Highest | Lowest | Active compute | +| PSUM | High | Low | MatMul accumulation | +| HBM | Lower | Higher | Input/output storage | **Guideline:** Minimize HBM accesses. Load once, compute multiple operations, store once. diff --git a/skills/neuron-nki-writing/references/transpose-and-layout.md b/skills/neuron-nki-writing/references/transpose-and-layout.md index 0ac960b..9fc8ecb 100644 --- a/skills/neuron-nki-writing/references/transpose-and-layout.md +++ b/skills/neuron-nki-writing/references/transpose-and-layout.md @@ -4,22 +4,22 @@ Production-proven transpose and layout transformation patterns. LLMs often strug ## Technique Summary (Quick Reference) -| Technique | When to Use | Hardware Gen | Production Files | -|-----------|-------------|--------------|------------------| -| `nisa.nc_transpose()` | P↔F transpose, after MatMul | All | Examples below, MLP CTE transpose pattern | -| `TensorView` | Zero-copy layout manipulation, broadcast, permute | All | [tensor-view.md](nkilib/core/tensor-view.md), examples below | -| `.ap()` patterns | Complex layouts, custom strides | All | Examples below | -| DMA strided access | Interleaved↔contiguous during DMA | gen3+ (optimized) | [layout-conversion.md](nkilib/patterns/layout-conversion.md), examples below | +| Technique | When to Use | Hardware Gen | Production Files | +| --------------------- | ------------------------------------------------- | ----------------- | ---------------------------------------------------------------------------- | +| `nisa.nc_transpose()` | P↔F transpose, after MatMul | All | Examples below, MLP CTE transpose pattern | +| `TensorView` | Zero-copy layout manipulation, broadcast, permute | All | [tensor-view.md](nkilib/core/tensor-view.md), examples below | +| `.ap()` patterns | Complex layouts, custom strides | All | Examples below | +| DMA strided access | Interleaved↔contiguous during DMA | gen3+ (optimized) | [layout-conversion.md](nkilib/patterns/layout-conversion.md), examples below | ### Constraint Quick Reference -| Constraint | Limit | Notes | -|------------|-------|-------| -| **Partition Dimension (P)** | ≤ 128 | First dimension of SBUF/PSUM, cannot reshape or stride | -| **SBUF Free Dimension (F)** | ≤ 32,767 | Second+ dimensions | -| **PSUM Free Dimension (F)** | 512 (gen2/3); gen4: 4096 fp32 / 8192 bf16 | For nc_transpose destination | -| **nc_transpose step size** | 2 for fp8/int8, 1 otherwise | Generation-specific | -| **TensorView partition rule** | Cannot permute dim 0 in SBUF | Hardware constraint | +| Constraint | Limit | Notes | +| ----------------------------- | ----------------------------------------- | ------------------------------------------------------ | +| **Partition Dimension (P)** | ≤ 128 | First dimension of SBUF/PSUM, cannot reshape or stride | +| **SBUF Free Dimension (F)** | ≤ 32,767 | Second+ dimensions | +| **PSUM Free Dimension (F)** | 512 (gen2/3); gen4: 4096 fp32 / 8192 bf16 | For nc_transpose destination | +| **nc_transpose step size** | 2 for fp8/int8, 1 otherwise | Generation-specific | +| **TensorView partition rule** | Cannot permute dim 0 in SBUF | Hardware constraint | --- @@ -128,14 +128,15 @@ else: # gen4 ### Constraints and Gotchas -| Constraint | Limit | Solution | -|------------|-------|----------| -| PSUM free dimension | 512 (gen2/3); gen4: 4096 fp32 / 8192 bf16 | Tile the transpose operation | -| Step size | 2 for fp8/int8 | Check dtype, adjust PSUM allocation | -| Partition dimension | ≤ 128 | Standard SBUF limit, tile if needed | -| PSUM → HBM | Not direct | Copy PSUM → SBUF → HBM | +| Constraint | Limit | Solution | +| ------------------- | ----------------------------------------- | ----------------------------------- | +| PSUM free dimension | 512 (gen2/3); gen4: 4096 fp32 / 8192 bf16 | Tile the transpose operation | +| Step size | 2 for fp8/int8 | Check dtype, adjust PSUM allocation | +| Partition dimension | ≤ 128 | Standard SBUF limit, tile if needed | +| PSUM → HBM | Not direct | Copy PSUM → SBUF → HBM | **Common error**: "PSUM dimension exceeds limit" + - **Cause**: Free dimension exceeds the PSUM limit — 512 (gen2/3); gen4: 4096 fp32 / 8192 bf16 - **Fix**: Tile the transpose into smaller chunks @@ -156,15 +157,15 @@ High-level abstraction for changing tensor layout without data movement. Uses st ### TensorView Method Reference -| Method | Purpose | Example | Partition Constraint | -|--------|---------|---------|---------------------| -| `.slice(dim, start, end, step)` | Strided slicing | Every 2nd element | Can slice partition (contiguous only) | -| `.permute(dims)` | Reorder dimensions | (P,H,W) → (P,W,H) | Partition must stay dim 0 | -| `.broadcast(dim, size)` | Expand size-1 dim | (P,1,F) → (P,N,F) | Original dim must be size 1 | -| `.reshape_dim(dim, shape)` | Split/merge dimension | (P,24) → (P,2,3,4) | Cannot reshape partition | -| `.flatten_dims(start, end)` | Merge contiguous dims | (P,2,3,4) → (P,24) | Cannot flatten partition | -| `.expand_dim(dim)` | Add size-1 dimension | (P,F) → (P,1,F) | Any position | -| `.rearrange(src, dst)` | Complex reshape+permute | Einops-style | Complex rules | +| Method | Purpose | Example | Partition Constraint | +| ------------------------------- | ----------------------- | ------------------ | ------------------------------------- | +| `.slice(dim, start, end, step)` | Strided slicing | Every 2nd element | Can slice partition (contiguous only) | +| `.permute(dims)` | Reorder dimensions | (P,H,W) → (P,W,H) | Partition must stay dim 0 | +| `.broadcast(dim, size)` | Expand size-1 dim | (P,1,F) → (P,N,F) | Original dim must be size 1 | +| `.reshape_dim(dim, shape)` | Split/merge dimension | (P,24) → (P,2,3,4) | Cannot reshape partition | +| `.flatten_dims(start, end)` | Merge contiguous dims | (P,2,3,4) → (P,24) | Cannot flatten partition | +| `.expand_dim(dim)` | Add size-1 dimension | (P,F) → (P,1,F) | Any position | +| `.rearrange(src, dst)` | Complex reshape+permute | Einops-style | Complex rules | ### Example 1: Strided DMA for Interleaved Layout @@ -276,6 +277,7 @@ TensorView(tensor).slice(dim=0, start=0, end=128, step=2) # ERROR: strided part **Do not use `nl.mgrid` for transpose operations.** This pattern appears in tutorials but has 0 occurrences in production code. **Use instead:** + - `TensorView` for zero-copy layout manipulation (Section 3) - `nisa.nc_transpose()` for P↔F transpose (Section 2) @@ -421,12 +423,12 @@ ap_pattern, ap_offset = tv._get_pattern_and_offset() ### Constraints and Pitfalls -| Issue | Cause | Solution | -|-------|-------|----------| -| Invalid pattern error | Pattern doesn't respect memory layout | Check stride calculations, verify access bounds | -| Partition dimension error | Cannot stride or reshape partition | Keep partition dim with stride=size or simple multiples | -| Compile-time only | Cannot compute pattern at runtime | All strides/sizes must be compile-time constants | -| Hard to debug | Low-level, easy to create invalid patterns | Use TensorView first, only drop to .ap() if needed | +| Issue | Cause | Solution | +| ------------------------- | ------------------------------------------ | ------------------------------------------------------- | +| Invalid pattern error | Pattern doesn't respect memory layout | Check stride calculations, verify access bounds | +| Partition dimension error | Cannot stride or reshape partition | Keep partition dim with stride=size or simple multiples | +| Compile-time only | Cannot compute pattern at runtime | All strides/sizes must be compile-time constants | +| Hard to debug | Low-level, easy to create invalid patterns | Use TensorView first, only drop to .ap() if needed | --- @@ -444,10 +446,10 @@ Combine data movement with layout transformation by using TensorView with DMA op ### Strided DMA vs SBUF Relayout -| Approach | When to Use | Constraint | Performance | -|----------|-------------|------------|-------------| -| **Strided DMA** | Any tensor size | gen3+ optimized | Medium overhead, but better than 2x DMA | -| **SBUF relayout** | Small tensors | B*n_heads*S ≤ gemm_moving_fmax | Fast for small, won't fit for large | +| Approach | When to Use | Constraint | Performance | +| ----------------- | --------------- | ------------------------------ | --------------------------------------- | +| **Strided DMA** | Any tensor size | gen3+ optimized | Medium overhead, but better than 2x DMA | +| **SBUF relayout** | Small tensors | B*n_heads*S ≤ gemm_moving_fmax | Fast for small, won't fit for large | ### Example 1: Interleaved to Contiguous (Strided Load) @@ -639,29 +641,30 @@ Need to change tensor layout? ### Hardware Generation Comparison -| Feature | gen2 (Trn1/Inf2) | gen3 (Trn2) | gen4 (Trn3) | -|---------|------------------|-------------|-------------| -| **PSUM Free Dim** | ≤ 512 | ≤ 512 | ≤ 4,096 | -| **nc_transpose** | Full support | Full support | Full support | -| **TensorView** | Full support | Full support | Full support | -| **Strided DMA** | Basic | Optimized | Optimized | -| **FP8 Support** | No | Yes | Yes (+ MXFP8/4) | -| **nc_transpose step=2** | For int8 | For fp8/int8 | For fp8/int8/mx | +| Feature | gen2 (Trn1/Inf2) | gen3 (Trn2) | gen4 (Trn3) | +| ----------------------- | ---------------- | ------------ | --------------- | +| **PSUM Free Dim** | ≤ 512 | ≤ 512 | ≤ 4,096 | +| **nc_transpose** | Full support | Full support | Full support | +| **TensorView** | Full support | Full support | Full support | +| **Strided DMA** | Basic | Optimized | Optimized | +| **FP8 Support** | No | Yes | Yes (+ MXFP8/4) | +| **nc_transpose step=2** | For int8 | For fp8/int8 | For fp8/int8/mx | ### Constraint Reference #### Dimension Limits -| Constraint | Limit | Memory Type | Notes | -|------------|-------|-------------|-------| -| **Partition (P)** | ≤ 128 | SBUF, PSUM | First dimension, fixed hardware limit | -| **Free (F)** | ≤ 32,767 | SBUF | Second+ dimensions | -| **PSUM Free (F)** | 512 (gen2/3); gen4: 4096 fp32 / 8192 bf16 | PSUM | nc_transpose destination, matmul result | -| **MatMul K** | ≤ 2,048 | Any | Contraction dimension | +| Constraint | Limit | Memory Type | Notes | +| ----------------- | ----------------------------------------- | ----------- | --------------------------------------- | +| **Partition (P)** | ≤ 128 | SBUF, PSUM | First dimension, fixed hardware limit | +| **Free (F)** | ≤ 32,767 | SBUF | Second+ dimensions | +| **PSUM Free (F)** | 512 (gen2/3); gen4: 4096 fp32 / 8192 bf16 | PSUM | nc_transpose destination, matmul result | +| **MatMul K** | ≤ 2,048 | Any | Contraction dimension | #### Partition Dimension Rules (SBUF) **Cannot do:** + - ❌ Reshape: `(128, 512) → (64, 1024)` - changes partition count - ❌ Stride: `tensor[::2, :]` - non-contiguous partition access - ❌ Flatten with free: `tensor.flatten()` - merges partition into free dims @@ -669,6 +672,7 @@ Need to change tensor layout? - ❌ Negative stride: `tensor[::-1, :]` - reverse partition order **Can do:** + - ✅ Contiguous slice: `tensor[0:64, :]` or `tensor[32:96, :]` - ✅ Full range: `tensor[0:128, :]` - ✅ Dynamic offset: `tensor[nl.ds(offset, size), :]` (contiguous only) @@ -676,14 +680,14 @@ Need to change tensor layout? #### Operation-Specific Constraints -| Operation | Constraint | Limit | Workaround | -|-----------|------------|-------|------------| -| `nisa.nc_transpose()` | PSUM free dim | ≤ 512 / ≤ 4,096 (gen) | Tile transpose into chunks | -| `nisa.nc_transpose()` | Step size | 2 for fp8/int8, 1 else | Check dtype, allocate PSUM accordingly | -| `TensorView.permute()` | Partition | Must stay dim 0 | Use nc_transpose for P↔F swap | -| `TensorView.reshape_dim()` | Partition | Cannot reshape | Only reshape free dimensions | -| `TensorView.flatten_dims()` | Partition | Cannot flatten with free | Keep partition separate | -| DMA strided access | Performance | gen3+ optimized | Works on gen2, but slower | +| Operation | Constraint | Limit | Workaround | +| --------------------------- | ------------- | ------------------------ | -------------------------------------- | +| `nisa.nc_transpose()` | PSUM free dim | ≤ 512 / ≤ 4,096 (gen) | Tile transpose into chunks | +| `nisa.nc_transpose()` | Step size | 2 for fp8/int8, 1 else | Check dtype, allocate PSUM accordingly | +| `TensorView.permute()` | Partition | Must stay dim 0 | Use nc_transpose for P↔F swap | +| `TensorView.reshape_dim()` | Partition | Cannot reshape | Only reshape free dimensions | +| `TensorView.flatten_dims()` | Partition | Cannot flatten with free | Keep partition separate | +| DMA strided access | Performance | gen3+ optimized | Works on gen2, but slower | --- @@ -691,14 +695,14 @@ Need to change tensor layout? ### Error Symptoms and Solutions -| Error Message / Symptom | Likely Cause | Solution | Section | -|-------------------------|--------------|----------|---------| -| "Partition dimension exceeds 128" | P > 128 | Tile outer loop to keep P ≤ 128 | 7 | -| "PSUM dimension exceeds limit" | F > 512 (gen2/3) / 4096 fp32 (gen4) / 8192 bf16 (gen4) | Tile nc_transpose operation | 2 | -| "Cannot reshape partition" | Reshape changes dim 0 | Only reshape free dimensions (dim≥1) | 3, 7 | -| "Partition must stay outermost" | TensorView.permute moved dim 0 | Keep partition at dim 0, or use nc_transpose | 3, 7 | -| "Stride not supported on partition" | Used `tensor[::2, :]` | Use contiguous slice + loop instead | 7 | -| Strided access slower than expected | Using gen2 hardware | Expected on gen2, gen3+ has optimization | 4 | +| Error Message / Symptom | Likely Cause | Solution | Section | +| ----------------------------------- | ------------------------------------------------------ | -------------------------------------------- | ------- | +| "Partition dimension exceeds 128" | P > 128 | Tile outer loop to keep P ≤ 128 | 7 | +| "PSUM dimension exceeds limit" | F > 512 (gen2/3) / 4096 fp32 (gen4) / 8192 bf16 (gen4) | Tile nc_transpose operation | 2 | +| "Cannot reshape partition" | Reshape changes dim 0 | Only reshape free dimensions (dim≥1) | 3, 7 | +| "Partition must stay outermost" | TensorView.permute moved dim 0 | Keep partition at dim 0, or use nc_transpose | 3, 7 | +| "Stride not supported on partition" | Used `tensor[::2, :]` | Use contiguous slice + loop instead | 7 | +| Strided access slower than expected | Using gen2 hardware | Expected on gen2, gen3+ has optimization | 4 | ### Common Anti-Patterns @@ -783,6 +787,7 @@ nisa.nc_transpose(dst=psum_result, data=input_large) ### Debugging Strategies 1. **Print tensor shapes at each step** + ```python print(f"Input shape: {input_sb.shape}") # Compile-time print transformed = TensorView(input_sb).permute([0, 2, 1]).get_view() @@ -790,6 +795,7 @@ nisa.nc_transpose(dst=psum_result, data=input_large) ``` 2. **Inspect TensorView patterns** + ```python tv = TensorView(tensor).slice(dim=1, start=0, end=100, step=2) pattern, offset = tv._get_pattern_and_offset() @@ -798,6 +804,7 @@ nisa.nc_transpose(dst=psum_result, data=input_large) ``` 3. **Validate with small test cases first** + ```python # Test with minimal sizes to verify logic test_input = nl.ndarray((8, 16), dtype=nl.float16, buffer=nl.sbuf) # Small @@ -806,6 +813,7 @@ nisa.nc_transpose(dst=psum_result, data=input_large) ``` 4. **Check generated .ap() patterns** + ```python # Manual .ap() pattern for debugging manual_pattern = [[stride0, size0], [stride1, size1]] @@ -818,6 +826,7 @@ nisa.nc_transpose(dst=psum_result, data=input_large) ``` 5. **Verify constraints before kernel launch** + ```python from nkilib.core.utils.kernel_assert import kernel_assert # or inline from references/nkilib/core/utils/kernel_assert.py @@ -836,12 +845,12 @@ nisa.nc_transpose(dst=psum_result, data=input_large) ### Self-Contained Reference Guide -| Topic | Primary Technique | Key Pattern | Reference | -|-------|-------------------|-------------|-----------| -| **TensorView** | Zero-copy views | slice, permute, broadcast, reshape, .ap() generation | [tensor-view.md](nkilib/core/tensor-view.md) | -| **Layout conversion** | TensorView + DMA | Interleaved↔contiguous via strided DMA or permutation matrix | [layout-conversion.md](nkilib/patterns/layout-conversion.md) | -| **Stream shuffle** | nc_stream_shuffle | Partition dimension broadcasting | [stream-shuffle-broadcast.md](nkilib/ops/stream-shuffle-broadcast.md) | -| **Tile tracking** | TiledDimInfo | Subtile indexing for nc_transpose destinations | [tile-info.md](nkilib/core/tile-info.md) | +| Topic | Primary Technique | Key Pattern | Reference | +| --------------------- | ----------------- | ------------------------------------------------------------ | --------------------------------------------------------------------- | +| **TensorView** | Zero-copy views | slice, permute, broadcast, reshape, .ap() generation | [tensor-view.md](nkilib/core/tensor-view.md) | +| **Layout conversion** | TensorView + DMA | Interleaved↔contiguous via strided DMA or permutation matrix | [layout-conversion.md](nkilib/patterns/layout-conversion.md) | +| **Stream shuffle** | nc_stream_shuffle | Partition dimension broadcasting | [stream-shuffle-broadcast.md](nkilib/ops/stream-shuffle-broadcast.md) | +| **Tile tracking** | TiledDimInfo | Subtile indexing for nc_transpose destinations | [tile-info.md](nkilib/core/tile-info.md) | ### Minimal Working Examples