diff --git a/contrib/models/Qwen3.5-35B-A3B/README.md b/contrib/models/Qwen3.5-35B-A3B/README.md new file mode 100644 index 00000000..c480f5f9 --- /dev/null +++ b/contrib/models/Qwen3.5-35B-A3B/README.md @@ -0,0 +1,288 @@ +# Qwen3.5-35B-A3B on NeuronX Distributed Inference (Trn2) + +`Qwen/Qwen3.5-35B-A3B` is the **MoE** flagship of the Qwen3.5 family — 35 B +total parameters, ~3 B activated per token ("A3B") through top-8 routing over +256 experts plus one sigmoid-gated shared expert. It uses the same hybrid +attention stack as the dense siblings — +**[3 gated DeltaNet + 1 full GQA] × 10 = 40 layers** — combined with a sparse +MoE feed-forward on every layer. + +This is the **first** DeltaNet + MoE integration on Neuron. It reuses the +DeltaNet + attention path from PR #173 (originally targeted at Qwen3.6-27B +dense) and plugs NxDI's `initialize_moe_module` (from `moe_v2`) into a new +`Qwen35MoEBlock`. Runs on the stock +`/opt/aws_neuronx_venv_pytorch_2_9_nxd_inference/` DLAMI venv +(`neuronx-cc` 2.26.6360 / `nki` 0.5.0). Only the config / decoder-layer / weight +converter in `modeling_qwen35.py` gained a MoE branch; the rest of the file +is byte-identical to the dense contribs. + +**Status:** both text-only **and vision-language** inference are validated +end-to-end on `trn2.48xlarge` (TP=8, bf16). VL uses the Neuron-compiled +vision encoder (buckets 1024 / 4096) plus a text-model recompile with +`use_text_only_cte_inputs=False` and CTE bucketing enabled, sharing the +same DeltaNet + MoE decoder as the text-only path. + +## Architecture diff vs dense Qwen3.5-27B + +| field | 27B (dense) | **35B-A3B (MoE)** | +|---|---:|---:| +| `hidden_size` | 5120 | **2048** | +| `intermediate_size` (dense MLP) | 17408 | — | +| `moe_intermediate_size` (per routed expert) | — | **512** | +| `shared_expert_intermediate_size` | — | **512** | +| `num_hidden_layers` | 64 | **40** | +| `num_attention_heads` | 24 | **16** | +| `num_key_value_heads` | 4 | **2** | +| `linear_num_value_heads` | 48 | **32** | +| `linear_num_key_heads` | 16 | 16 | +| `head_dim` | 256 | 256 | +| **`num_experts`** | — | **256** | +| **`num_experts_per_tok`** | — | **8** | +| **shared expert count** | — | **1** with per-token sigmoid gate | +| `tie_word_embeddings` | false | false | +| model_type | `qwen3_5_text` | `qwen3_5_moe_text` | + +Total params: 35 B (55.6 GB of bf16 safetensors weight); activated per token +~3 B (top-8 of 256 routed + 1 shared). + +## Contents + +``` +Qwen3.5-35B-A3B/ +├── README.md +├── src/ +│ ├── modeling_qwen35.py — DeltaNet + GQA text stack + NEW `Qwen35MoEBlock` +│ ├── modeling_qwen35_vl.py — VL orchestrator (vision + text) +│ ├── modeling_qwen35_vision.py +│ ├── hybrid_apc.py +│ ├── nki_kernels/ — DeltaNet NKI kernels (unchanged) +│ └── __init__.py +└── test/integration/ — same runner/bench scripts as dense contribs +``` + +## What changed in `modeling_qwen35.py` for MoE + +Deltas vs the dense contrib source (a couple of hundred lines total): + +1. **Config**: `Qwen35InferenceConfig.from_pretrained` preserves + `model_type` (was hardcoded to `qwen3_5_text`) so we can detect + `qwen3_5_moe_text` variants. When a MoE config is detected the + `__init__` sets `num_local_experts = num_experts`, `n_shared_experts = 1`, + maps `moe_intermediate_size → intermediate_size` (used by + `initialize_moe_module` to size the routed experts), and populates + `shared_expert_intermediate_size`. + +2. **`Qwen35MoEBlock`**: new nn.Module inserted between `Qwen35MLP` and + `NeuronQwen35DecoderLayer`. Wraps NxDI's `initialize_moe_module` for the + routed experts and re-implements a shared expert with a **per-token + sigmoid gate** (Qwen3.5-MoE specific — NxDI's built-in `SharedExperts` + only sums into the routed output without a per-token gate). + +3. **Decoder layer**: `NeuronQwen35DecoderLayer.__init__` picks + `Qwen35MoEBlock(config)` when `config._is_moe` is set, else the dense + `Qwen35MLP` / `NeuronLlamaMLP` path. + +4. **Weight converter**: `convert_qwen35_hf_to_neuron_state_dict` gains an + MoE branch. HF stores stacked 3D expert weights `(num_experts, 2*I, H)` + and `(num_experts, H, I)`; NxDI's `ExpertMLPsV2` expects the transposed + layout `(num_experts, H, 2*I)` and `(num_experts, I, H)`. Router key + `mlp.gate.weight` → `mlp.moe.router.linear_router.weight`; shared expert + keys `mlp.shared_expert.{gate,up,down}_proj.weight` → + `mlp.shared_{gate,up,down}_proj.weight`; expert stacked tensors renamed + `mlp.experts.{gate_up_proj,down_proj}` → + `mlp.moe.expert_mlps.mlp_op.{gate_up_proj,down_proj}.weight`; the + `mlp.shared_expert_gate.weight` scalar-output linear is unchanged. + +## Compatibility + +| Component | Version | +|---|---| +| Instance | `trn2.48xlarge` (validated at TP=8) | +| `neuronx-cc` | 2.26.6360.0 | +| `nki` | 0.5.0 | +| `neuronx-distributed` | 0.19.28492 | +| `neuronx-distributed-inference` | 0.10.18399 | +| `torch-neuronx` | 2.9.0.2 (torch 2.9.1) | +| `libneuronxla` | 2.2.17544 | +| Python | 3.12 | +| `transformers` | 4.57.6 (Neuron runtime). HF CPU reference needs ≥ 5.13. | + +## Checkpoint + +- HuggingFace: [`Qwen/Qwen3.5-35B-A3B`](https://huggingface.co/Qwen/Qwen3.5-35B-A3B) +- Architecture identifier: `qwen3_5_moe` +- Weights: **14 shards, ~67 GB bfloat16** + +Download: + +```bash +python -c "from huggingface_hub import snapshot_download; \ + snapshot_download('Qwen/Qwen3.5-35B-A3B', local_dir='/mnt/nvme/models/Qwen3.5-35B-A3B')" +``` + +## Quick start — text-only + +```bash +source /opt/aws_neuronx_venv_pytorch_2_9_nxd_inference/bin/activate +python contrib/models/Qwen3.5-35B-A3B/test/integration/run_text_smoke.py \ + --model-path /mnt/nvme/models/Qwen3.5-35B-A3B \ + --compiled-path /tmp/qwen35_35b_a3b_traced \ + --tp 8 --seq-len 512 --max-new-tokens 32 \ + --prompt "The capital of France is" +``` + +Sample validated on `trn2.48xlarge`, TP=8, bf16, seq_len=512: + +``` +prompt : 'The capital of France is' +output : 'The capital of France is Paris.\nThe capital of France is Paris.\n...' +TTFT : 561.4 ms +TPOT : 7.4 ms (136.03 tok/s) +``` + +Additional prompts produce coherent, factually correct outputs across the +suite (Jupiter for largest planet, 100 °C water boiling point, autumn haiku, +photosynthesis definition). + +## Measured text-only performance (TP=8, bf16, seq_len=512) + +`run_benchmark.py --prompt-lens 16 64 256 --max-new-tokens 64 --repeats 5` + +Three MoE blockwise-matmul backends were evaluated on the shipped +`neuronx-cc 2.26.6360` DLAMI: + +| MoE backend | 16 tok TTFT | 64 tok TTFT | 256 tok TTFT | TPOT | notes | +|---|---:|---:|---:|---:|---| +| `use_torch_block_wise=True` | 553.8 ms | 554.0 ms | 553.6 ms | ~7.7 ms | unrolled per-block loop in NxDI graph; runs on Neuron but no NKI kernel fusion | +| **`use_shard_on_intermediate_dynamic_while=True`** (default) | **480.0 ms** | **575.5 ms** | **667.1 ms** | **~7.7 ms** | LNC=2 NKI kernel sharding on intermediate dim; 13 % faster at prompt=16, grows with prompt length | +| patched `_call_shard_hidden_kernel` (nkilib fwd) | 567.0 ms | 566.5 ms | 566.7 ms | ~7.8 ms | see below | + +`shard_on_intermediate` is the current default in `run_text_smoke.py` / +`run_benchmark.py` — it wins at short prompts (chat use case). The third +row is a **monkey-patched path**: NxDI's default LNC=2 forward MoE kernel +(`_call_shard_hidden_kernel`) is a `NotImplementedError` stub because the +`neuronxcc.nki._private.blockwise_mm` module is absent from the shipped +`neuronx-cc 2.26.6360` DLAMI. However, the same kernel forward +implementation is available at +`nkilib.experimental.moe.forward.bwmm_shard_on_H.blockwise_mm_baseline_shard_hidden`, +so `modeling_qwen35.py::_patch_nxd_shard_hidden_kernel()` wires it in at +import time. That gives the default LNC=2 forward path a runnable +implementation (flat 567 ms TTFT regardless of prompt length) — slower +than shard-on-intermediate for short prompts but more stable at long +prompts. Kept as a safety net; not chosen as the default. + +None of these three paths matches what a genuine LNC=2 shard-hidden NKI +kernel could deliver in a future SDK drop. + +**Why is TTFT / TPOT ratio so different from the dense siblings?** +Dense models spend ~40-60 % of prefill in the FFN block; MoE only dispatches +top-8 experts per token, so decode is extremely cheap (~10× less MLP FLOPs +than 27B despite having more total params) but prefill still has to route +256 experts × 512 tokens = many small block-matmuls. The observed +ratios are consistent with this: + +| model | Params (activated per token) | TTFT (16) | TPOT | TTFT/TPOT | +|---|---:|---:|---:|---:| +| 2B | 2 B | 17.6 ms | 4.00 ms | 4.4× | +| 4B | 4 B | 34.6 ms | 5.68 ms | 6.1× | +| 9B | 9 B | 42.9 ms | 6.89 ms | 6.2× | +| 27B | 27 B | 118.6 ms | 21.58 ms | 5.5× | +| **35B-A3B** | 35 B (**~3 B active**) | **480 ms** | **7.72 ms** | **62×** | + +TPOT of ~7.7 ms sits between 9B (6.9 ms) and 4B (5.7 ms), matching the +"~3 B activated" figure — confirming the MoE sparse-decode benefit. The +inflated TTFT/TPOT ratio is a consequence of the shipped kernel gap, not the +architecture. + +## Measured VL performance (TP=8, bf16, CTE buckets 512/1024/2048/4096/8192) + +Vision encoder compiled to Neuron for two buckets (1024, 4096) via +`compile_vision_encoder.py`; text model compiled with +`use_text_only_cte_inputs=False` and CTE bucketing enabled so the same +artifact handles all three image sizes. Prompt: *"What is in this image? +Describe it briefly."* — three repeats after warmup. + +| Image | Vision path | TTFT (ms) | TPOT (ms) | tok/s | +|---|---|---:|---:|---:| +| 512×512 | Neuron ViT (bucket 1024) | **1036.8** | 8.2 | 121.9 | +| 1024×1024 | Neuron ViT (bucket 4096) | **2004.0** | 8.1 | 123.6 | +| 2048×2048 | 2×2 tile → 4× Neuron ViT (bucket 4096) | **5490.4** | 8.1 | 123.2 | + +TPOT is flat (~8 ms) across sizes because decode is unchanged; TTFT scales +with image → text-token count (image tokens: ~278 / 1046 / 4712, CTE bucket +picked: 512 / 1024 / 8192). Sample output on `test_image_1024.jpg` — the +model correctly identifies the animal as a **Pallas's cat (manul)**; +`test_image_2048.jpg` (a leopard-like feline) resolves to *"close-up of an +animal...rosettes/spots"* using the 2×2 tile fallback. Both descriptions +show the vision+MoE stack is functioning end-to-end. + +To reproduce: + +```bash +# 1. Compile vision encoder buckets (once, ~10 min for 1024, ~12 min for 4096) +python contrib/models/Qwen3.5-35B-A3B/test/integration/compile_vision_encoder.py \ + --model-path /mnt/nvme/models/Qwen3.5-35B-A3B \ + --out-dir /tmp/qwen35_35b_a3b_vl_bench/vision \ + --buckets 1024 4096 + +# 2. Compile text model + benchmark (once, ~15 min compile then benchmark) +python contrib/models/Qwen3.5-35B-A3B/test/integration/run_vl_benchmark.py \ + --model-path /mnt/nvme/models/Qwen3.5-35B-A3B \ + --compiled-path /tmp/qwen35_35b_a3b_vl_bench \ + --tp 8 --images 512 1024 2048 --buckets 512 1024 2048 4096 8192 \ + --max-new-tokens 48 --repeats 3 \ + --vision-compiled-dir /tmp/qwen35_35b_a3b_vl_bench/vision \ + --out-json /tmp/qwen35_35b_a3b_vl_bench.json +``` + +## Notable configuration choices + +- **`MoENeuronConfig`** (not `NeuronConfig`) — required by + `initialize_moe_module` so it can find `router_config`, + `blockwise_matmul_config`, `moe_tp_degree`, etc. +- `moe_tp_degree = 8`, `moe_ep_degree = 1` — no expert parallelism yet, + every rank sees every expert (sharded on the intermediate dim). +- **`blockwise_matmul_config={"use_shard_on_intermediate_dynamic_while": True}`** + (default). Two alternate MoE prefill kernels are evaluated in the + benchmark table above; the shard-on-intermediate path is the fastest one + that runs on the shipped `neuronx-cc 2.26.6360` DLAMI. The truly-preferred + `_call_shard_hidden_kernel` LNC=2 path is not available in the shipped + DLAMI and would require a future SDK drop. +- **`moe_ep_degree=1`** — expert parallelism (`moe_ep_degree > 1`) is + functional during prefill but **not supported during decode** by NxDI's + `ExpertMLPsV2`: with `top_k=8 / num_experts=256`, the per-token expert + fraction is 3 %, below `DEFAULT_SELECTIVE_LOADING_THRESHOLD=1.0`, which + routes decode through `forward_selective_loading` — and that path + raises `NotImplementedError: Selective Loading with Expert parallelism is + not supported in token generation.` +- `router_config.dtype = float32`, `router_config.act_fn = "softmax"` — + Qwen3.5-MoE uses softmax over router logits with fp32 accumulation. +- `normalize_top_k_affinities = True` — Qwen3.5-MoE normalizes the top-k + weights so they sum to 1 per token. +- `QWEN36_DELTANET_CTE_IMPL=legacy_direct`, `QWEN36_DELTANET_MULTIHEAD_CTE=0` + — same DeltaNet numerical stability defaults as the dense siblings; the + fused-multihead NKI kernel is not needed for text-only decode. + +## Known limitations / follow-ups + +- **HF greedy match not run**. 35B-A3B on CPU bf16 is ~67 GB and greedy + generation takes many minutes per prompt; deferred until GPU or larger CPU + is available. All 5 prompts in the accuracy suite produce coherent, + factually correct Neuron output. +- **Shipped SDK MoE kernel gap.** The truly-preferred LNC=2 shard-hidden + NKI kernel (`_call_shard_hidden_kernel`) is not present in the shipped + `neuronx-cc 2.26.6360` DLAMI. Current default + (`use_shard_on_intermediate_dynamic_while`) gets + ~480 ms at prompt=16 (13 % better than the torch fallback) but scales + linearly with prompt length. A future SDK drop should close this. +- **Expert parallelism (EP=1).** `moe_ep_degree > 1` is supported for + prefill but NxDI's `ExpertMLPsV2` raises `NotImplementedError` for + selective-loading decode with EP — and our (top_k=8, num_experts=256) + configuration always goes through selective loading at decode time. + +## Maintainer + +Contributed alongside the 2B/4B/9B/27B dense siblings. This is the first +DeltaNet + MoE integration on Neuron; the MoE plumbing follows NxDI's +`qwen3_moe` model as a reference and adapts it for Qwen3.5-MoE's +sigmoid-gated shared expert. diff --git a/contrib/models/Qwen3.5-35B-A3B/src/__init__.py b/contrib/models/Qwen3.5-35B-A3B/src/__init__.py new file mode 100644 index 00000000..7e79aa03 --- /dev/null +++ b/contrib/models/Qwen3.5-35B-A3B/src/__init__.py @@ -0,0 +1,41 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +from src.modeling_qwen35 import ( + NeuronGatedDeltaNet, + NeuronQwen35Attention, + NeuronQwen35DecoderLayer, + NeuronQwen35ForCausalLM, + NeuronQwen35Model, + Qwen35DecoderModelInstance, + Qwen35InferenceConfig, + Qwen35MLP, + Qwen35ModelWrapper, +) +from src.modeling_qwen35_vision import ( + NeuronQwen35VisionForImageEncoding, + NeuronQwen35VisionModel, +) +from src.modeling_qwen35_vl import ( + NeuronQwen35VLForCausalLM, + Qwen35VLInferenceConfig, +) + +__all__ = [ + # Text decoder + "NeuronGatedDeltaNet", + "NeuronQwen35Attention", + "NeuronQwen35DecoderLayer", + "NeuronQwen35ForCausalLM", + "NeuronQwen35Model", + "Qwen35DecoderModelInstance", + "Qwen35InferenceConfig", + "Qwen35MLP", + "Qwen35ModelWrapper", + # Vision encoder + "NeuronQwen35VisionForImageEncoding", + "NeuronQwen35VisionModel", + # Vision-language + "NeuronQwen35VLForCausalLM", + "Qwen35VLInferenceConfig", +] diff --git a/contrib/models/Qwen3.5-35B-A3B/src/hybrid_apc.py b/contrib/models/Qwen3.5-35B-A3B/src/hybrid_apc.py new file mode 100644 index 00000000..f1304a73 --- /dev/null +++ b/contrib/models/Qwen3.5-35B-A3B/src/hybrid_apc.py @@ -0,0 +1,1798 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Qwen hybrid APC metadata lifecycle. + +This module intentionally stores only control-plane metadata. GDN recurrent and +conv checkpoint tensors live in the model-side checkpoint bank; the metadata +store owns prefix identity, validity, refcounts, LRU state, and memory +accounting. +""" + +from __future__ import annotations + +import hashlib +import os +import struct +from collections import OrderedDict, deque +from dataclasses import dataclass +from typing import Callable, Hashable, Iterable, NamedTuple + +import torch + + +class HybridPrefixKey(NamedTuple): + cumulative_prefix_hash: Hashable + prefix_len: int + block_size: int + cache_salt: Hashable | None + model_revision: str + layout_version: int + tp_rank: int + recurrent_dtype: str + conv_dtype: str + + +class HybridAPCHitPlan(NamedTuple): + attention_hit_len: int + recurrent_hit_len: int + conv_hit_len: int + usable_hit_len: int + restore_checkpoint_prefix_len: int + residual_replay_len: int + suffix_len: int + checkpoint_slot: int | None + checkpoint_key: HybridPrefixKey | None + + +class HybridAPCPreparedRequest(NamedTuple): + request_id: Hashable + input_dict: dict[str, torch.Tensor] + plan: HybridAPCHitPlan + commit_prefix_len: int + commit_key: HybridPrefixKey | None + commit_slot: int | None + attention_block_refs: tuple[int, ...] + + +@dataclass +class HybridAPCStats: + checkpoints: int = 0 + bytes_used: int = 0 + evictions: int = 0 + hits: int = 0 + misses: int = 0 + + +@dataclass +class HybridAPCRequestRecord: + request_id: Hashable + state: str + restored_key: HybridPrefixKey | None = None + committed_keys: list[HybridPrefixKey] | None = None + reserved_slots: list[int] | None = None + + def __post_init__(self): + if self.committed_keys is None: + self.committed_keys = [] + if self.reserved_slots is None: + self.reserved_slots = [] + + +@dataclass +class HybridPrefixCheckpoint: + key: HybridPrefixKey + prefix_len: int + attention_block_refs: tuple[int, ...] + gdn_checkpoint_slot: int + valid_recurrent_layers: torch.Tensor + valid_conv_layers: torch.Tensor + refcount: int = 0 + last_access_step: int = 0 + bytes_used: int = 0 + evictable: bool = True + attention_valid: bool = True + + def has_valid_recurrent(self, required_layers: tuple[int, ...]) -> bool: + return _mask_has_layers(self.valid_recurrent_layers, required_layers) + + def has_valid_conv(self, required_layers: tuple[int, ...]) -> bool: + return _mask_has_layers(self.valid_conv_layers, required_layers) + + def has_valid_gdn(self, required_layers: tuple[int, ...]) -> bool: + return self.has_valid_recurrent(required_layers) and self.has_valid_conv( + required_layers + ) + + def has_valid_hybrid_state(self, required_layers: tuple[int, ...]) -> bool: + return self.attention_valid and self.has_valid_gdn(required_layers) + + +def _normalize_dtype(dtype: str | torch.dtype) -> str: + if dtype == torch.float32: + return "float32" + if dtype == torch.bfloat16: + return "bfloat16" + normalized = str(dtype).lower() + aliases = { + "fp32": "float32", + "float32": "float32", + "torch.float32": "float32", + "bf16": "bfloat16", + "bfloat16": "bfloat16", + "torch.bfloat16": "bfloat16", + } + if normalized not in aliases: + raise ValueError(f"unsupported hybrid APC dtype: {dtype}") + return aliases[normalized] + + +def _mask_has_layers(mask: torch.Tensor, required_layers: tuple[int, ...]) -> bool: + if mask.numel() == 0: + return False + for layer in required_layers: + if layer >= mask.numel() or not bool(mask[layer].item()): + return False + return True + + +def _env_flag(name: str) -> bool: + value = os.environ.get(name) + return value is not None and value.strip().lower() not in { + "", + "0", + "false", + "no", + "off", + } + + +def _publish_scheduler_gdn_checkpoint(key): + try: + from qwen36_hybrid_apc_scheduler_patch import ( # noqa: WPS433 + register_hybrid_apc_gdn_checkpoint, + ) + except Exception: + return + try: + register_hybrid_apc_gdn_checkpoint(key) + except Exception: + return + + +def _unpublish_scheduler_gdn_checkpoint(key): + try: + from qwen36_hybrid_apc_scheduler_patch import ( # noqa: WPS433 + unregister_hybrid_apc_gdn_checkpoint, + ) + except Exception: + return + try: + unregister_hybrid_apc_gdn_checkpoint(key) + except Exception: + return + + +def estimate_qwen_gdn_checkpoint_bytes_per_rank( + *, + num_gdn_layers: int = 48, + local_value_heads: int = 12, + local_key_heads: int = 4, + key_dim: int = 128, + value_dim: int = 128, + conv_kernel_size: int = 4, + recurrent_dtype: str | torch.dtype = "float32", + conv_dtype: str | torch.dtype = "bfloat16", +) -> int: + recurrent_dtype = _normalize_dtype(recurrent_dtype) + conv_dtype = _normalize_dtype(conv_dtype) + recurrent_bytes = 4 if recurrent_dtype == "float32" else 2 + conv_bytes = 4 if conv_dtype == "float32" else 2 + recurrent_numel = num_gdn_layers * local_value_heads * key_dim * value_dim + conv_dim = 2 * local_key_heads * key_dim + local_value_heads * value_dim + conv_numel = num_gdn_layers * conv_dim * (conv_kernel_size - 1) + return recurrent_numel * recurrent_bytes + conv_numel * conv_bytes + + +def estimate_qwen_hybrid_cache_bytes_per_rank( + *, + max_context_len: int, + checkpoint_interval: int, + num_attention_layers: int = 16, + local_kv_heads: int = 1, + attention_head_dim: int = 256, + attention_kv_dtype: str | torch.dtype = "bfloat16", + **gdn_kwargs, +) -> dict[str, int]: + attention_dtype = _normalize_dtype(attention_kv_dtype) + attention_bytes = 4 if attention_dtype == "float32" else 2 + attention_kv = ( + int(max_context_len) + * num_attention_layers + * 2 + * local_kv_heads + * attention_head_dim + * attention_bytes + ) + checkpoints = max(0, int(max_context_len)) // int(checkpoint_interval) + gdn_per_checkpoint = estimate_qwen_gdn_checkpoint_bytes_per_rank(**gdn_kwargs) + gdn_total = checkpoints * gdn_per_checkpoint + return { + "attention_kv_bytes": attention_kv, + "gdn_checkpoint_bytes": gdn_total, + "gdn_bytes_per_checkpoint": gdn_per_checkpoint, + "num_gdn_checkpoints": checkpoints, + "total_bytes": attention_kv + gdn_total, + } + + +def _flatten_single_request_tokens(token_ids: torch.Tensor | Iterable[int]) -> torch.Tensor: + if isinstance(token_ids, torch.Tensor): + tokens = token_ids.detach().cpu() + else: + tokens = torch.tensor(list(token_ids), dtype=torch.int64) + if tokens.ndim == 2 and tokens.shape[0] == 1: + tokens = tokens.reshape(-1) + elif tokens.ndim != 1: + raise ValueError( + "token_ids must be a single request tensor with shape [seq] or [1, seq], " + f"got {tuple(tokens.shape)}" + ) + return tokens.to(torch.int64).contiguous() + + +def build_cumulative_prefix_hashes( + token_ids: torch.Tensor | Iterable[int], + *, + block_size: int, + prefix_lens: Iterable[int] | None = None, +) -> dict[int, str]: + """Build deterministic cumulative prefix hashes at block boundaries. + + This is a local scheduler bridge helper, not a replacement for vLLM's + production block hash. It deliberately hashes the parent digest plus the + next block's token bytes so a reused final block with a different parent + prefix produces a different cumulative hash. + """ + + block_size = int(block_size) + if block_size <= 0: + raise ValueError(f"block_size must be positive, got {block_size}") + + tokens = _flatten_single_request_tokens(token_ids) + seq_len = int(tokens.numel()) + if prefix_lens is None: + requested_lens = set(range(block_size, seq_len + 1, block_size)) + else: + requested_lens = {int(prefix_len) for prefix_len in prefix_lens} + requested_lens = {prefix_len for prefix_len in requested_lens if prefix_len > 0} + for prefix_len in requested_lens: + if prefix_len > seq_len: + raise ValueError(f"prefix_len {prefix_len} exceeds token length {seq_len}") + if prefix_len % block_size != 0: + raise ValueError( + f"prefix_len {prefix_len} must be a multiple of block_size {block_size}" + ) + + if not requested_lens: + return {} + + max_prefix_len = max(requested_lens) + parent_digest = b"" + hashes: dict[int, str] = {} + for block_start in range(0, max_prefix_len, block_size): + block_end = block_start + block_size + block = tokens[block_start:block_end] + digest = hashlib.blake2b(digest_size=16) + digest.update(parent_digest) + digest.update(struct.pack(" int: + checkpoint_interval = int(checkpoint_interval) + if checkpoint_interval <= 0: + raise ValueError( + f"checkpoint_interval must be positive, got {checkpoint_interval}" + ) + return max(0, int(prefix_len)) // checkpoint_interval * checkpoint_interval + + +def apply_hybrid_apc_prefill_plan( + input_dict: dict[str, torch.Tensor], + *, + plan: HybridAPCHitPlan, + commit_slot: int | None = None, + request_prefix_len: int | None = None, + gdn_active_carry: bool = False, + block_size: int | None = None, +) -> dict[str, torch.Tensor]: + """Materialize model inputs for a scheduler-selected hybrid APC hit plan. + + The serving scheduler owns prefix hashing, attention block-table selection, + checkpoint lookup, and checkpoint-slot reservation. This helper only applies + the chosen restore boundary to the token tensors and emits explicit + restore/commit control tensors. GDN state is restored only when the plan has + a checkpoint slot; slot ID presence alone is never treated as a cache hit. + """ + + if "input_ids" not in input_dict: + raise KeyError("input_ids is required to apply a hybrid APC prefill plan") + + input_ids = input_dict["input_ids"] + if input_ids.ndim != 2: + raise ValueError(f"input_ids must be [batch, seq], got {tuple(input_ids.shape)}") + + batch_size, available_len = input_ids.shape + prompt_len = available_len if request_prefix_len is None else int(request_prefix_len) + restore_len = int(plan.restore_checkpoint_prefix_len) + if prompt_len < 0: + raise ValueError(f"request_prefix_len must be non-negative, got {prompt_len}") + if restore_len < 0 or restore_len > prompt_len: + raise ValueError( + "restore_checkpoint_prefix_len must be in [0, request_prefix_len], " + f"got {restore_len} and {prompt_len}" + ) + if prompt_len > available_len: + raise ValueError( + f"request_prefix_len {prompt_len} exceeds input_ids length {available_len}" + ) + if plan.checkpoint_slot is None and restore_len != 0: + raise ValueError("restore checkpoint prefix length requires a checkpoint slot") + if plan.checkpoint_slot is not None and restore_len == 0: + raise ValueError("checkpoint slot restore requires a positive prefix length") + + output = dict(input_dict) + suffix_len = prompt_len - restore_len + device = input_ids.device + + output["input_ids"] = input_ids[:, restore_len:prompt_len] + + attention_mask = input_dict.get("attention_mask") + if ( + isinstance(attention_mask, torch.Tensor) + and attention_mask.ndim >= 2 + and attention_mask.shape[0] == batch_size + and attention_mask.shape[1] >= prompt_len + ): + output["attention_mask"] = attention_mask[:, restore_len:prompt_len] + + inputs_embeds = input_dict.get("inputs_embeds") + if ( + isinstance(inputs_embeds, torch.Tensor) + and inputs_embeds.ndim >= 3 + and inputs_embeds.shape[0] == batch_size + and inputs_embeds.shape[1] >= prompt_len + ): + output["inputs_embeds"] = inputs_embeds[:, restore_len:prompt_len] + + def _slot_mapping_covers_suffix(value: torch.Tensor) -> bool: + if value.ndim == 1: + if batch_size == 1: + return int(value.numel()) >= suffix_len + return int(value.numel()) >= batch_size * suffix_len + if value.ndim >= 2: + return value.shape[0] >= batch_size and value.shape[1] >= suffix_len + return False + + def _slot_mapping_needs_repair(value) -> bool: + if not isinstance(value, torch.Tensor) or value.numel() == 0: + return True + if not _slot_mapping_covers_suffix(value): + return True + return bool((value.to(torch.int64) < 0).any().item()) + + unbacked_attention_hit = ( + plan.checkpoint_slot is None + and int(plan.attention_hit_len) > 0 + and restore_len == 0 + ) + + def _synthesize_suffix_slot_mapping() -> torch.Tensor | None: + if block_size is None or int(block_size) <= 0 or suffix_len <= 0: + return None + block_table = input_dict.get("block_table") + if not isinstance(block_table, torch.Tensor) or block_table.numel() == 0: + return None + table = block_table + if table.ndim == 1: + table = table.unsqueeze(0) + if table.ndim != 2 or table.shape[0] < batch_size: + return None + block_size_int = int(block_size) + positions = torch.arange( + restore_len, + prompt_len, + dtype=torch.int64, + device=table.device, + ) + logical_blocks = torch.div(positions, block_size_int, rounding_mode="floor") + if logical_blocks.numel() == 0 or int(logical_blocks.max().item()) >= table.shape[1]: + return None + offsets = positions.remainder(block_size_int) + rows = [] + table_i64 = table.to(torch.int64) + for batch_idx in range(batch_size): + physical_blocks = torch.index_select( + table_i64[batch_idx], + 0, + logical_blocks, + ) + rows.append(physical_blocks * block_size_int + offsets) + return torch.stack(rows, dim=0) + + slot_mapping = input_dict.get("slot_mapping") + if ( + isinstance(slot_mapping, torch.Tensor) + and slot_mapping.ndim >= 2 + and slot_mapping.shape[0] == batch_size + and slot_mapping.shape[1] >= prompt_len + ): + output["slot_mapping"] = slot_mapping[:, restore_len:prompt_len] + elif isinstance(slot_mapping, torch.Tensor) and slot_mapping.ndim == 1: + if batch_size == 1 and slot_mapping.numel() >= prompt_len: + output["slot_mapping"] = slot_mapping[restore_len:prompt_len] + elif slot_mapping.numel() >= batch_size * prompt_len: + flattened = slot_mapping.reshape(batch_size, -1) + output["slot_mapping"] = flattened[:, restore_len:prompt_len] + if unbacked_attention_hit: + synthesized_slot_mapping = _synthesize_suffix_slot_mapping() + if synthesized_slot_mapping is not None: + dtype = ( + slot_mapping.dtype + if isinstance(slot_mapping, torch.Tensor) + else torch.int32 + ) + output["slot_mapping"] = synthesized_slot_mapping.to(dtype=dtype) + elif _slot_mapping_needs_repair(output.get("slot_mapping")): + synthesized_slot_mapping = _synthesize_suffix_slot_mapping() + if synthesized_slot_mapping is not None: + dtype = ( + slot_mapping.dtype + if isinstance(slot_mapping, torch.Tensor) + else torch.int32 + ) + repaired_slot_mapping = synthesized_slot_mapping.to(dtype=dtype) + current_slot_mapping = output.get("slot_mapping") + if ( + isinstance(current_slot_mapping, torch.Tensor) + and current_slot_mapping.numel() == repaired_slot_mapping.numel() + ): + repaired_slot_mapping = torch.where( + current_slot_mapping.to(torch.int64) < 0, + repaired_slot_mapping.reshape(current_slot_mapping.shape), + current_slot_mapping, + ) + output["slot_mapping"] = repaired_slot_mapping + + position_template = input_dict.get("position_ids") + position_dtype = ( + position_template.dtype + if isinstance(position_template, torch.Tensor) + else torch.int64 + ) + position_ids = torch.arange( + restore_len, + prompt_len, + dtype=position_dtype, + device=device, + ).unsqueeze(0) + output["position_ids"] = position_ids.expand(batch_size, suffix_len).contiguous() + default_rotary_positions = torch.arange( + restore_len, + prompt_len, + dtype=torch.int32, + device=device, + ) + output["rotary_position_ids"] = default_rotary_positions.view( + 1, + 1, + suffix_len, + ).expand(3, batch_size, suffix_len).contiguous() + + for key in ("rotary_position_id", "rotary_position_ids"): + value = input_dict.get(key) + if not isinstance(value, torch.Tensor): + continue + if ( + value.ndim == 2 + and value.shape[0] == batch_size + and value.shape[1] >= prompt_len + ): + output[key] = value[:, restore_len:prompt_len] + elif ( + value.ndim == 3 + and value.shape[1] == batch_size + and value.shape[2] >= prompt_len + ): + output[key] = value[:, :, restore_len:prompt_len] + + def _batch_i32(value: int) -> torch.Tensor: + return torch.full((batch_size,), int(value), dtype=torch.int32, device=device) + + def _batch_i32_col(value: int) -> torch.Tensor: + return torch.full((batch_size, 1), int(value), dtype=torch.int32, device=device) + + disable_restore = _env_flag("QWEN36_DISABLE_HYBRID_GDN_RESTORE") + disable_commit = _env_flag("QWEN36_DISABLE_HYBRID_GDN_COMMIT") + restore_available = plan.checkpoint_slot is not None and not disable_restore + restore_enabled = restore_available and not gdn_active_carry + commit_enabled = commit_slot is not None and not disable_commit + output["computed_context_lens"] = _batch_i32_col(restore_len) + output["full_context_lens"] = _batch_i32_col(prompt_len) + output["num_queries"] = _batch_i32_col(suffix_len) + output["hybrid_restore_slot_ids"] = _batch_i32( + 0 if not restore_available else int(plan.checkpoint_slot) + ) + output["hybrid_restore_mask"] = _batch_i32(1 if restore_enabled else 0) + output["hybrid_restore_prefix_lens"] = _batch_i32( + restore_len if restore_available else 0 + ) + output["hybrid_commit_slot_ids"] = _batch_i32( + 0 if not commit_enabled else commit_slot + ) + output["hybrid_commit_mask"] = _batch_i32(1 if commit_enabled else 0) + + if os.environ.get("QWEN36_HYBRID_APC_DEBUG") == "1": + print( + "[hybrid_apc_debug] apply " + f"prompt_len={prompt_len} restore_len={restore_len} " + f"suffix_len={suffix_len} restore_slot={plan.checkpoint_slot} " + f"commit_slot={commit_slot} gdn_active_carry={gdn_active_carry} " + f"input_shape={tuple(input_ids.shape)} " + f"output_shape={tuple(output['input_ids'].shape)}", + flush=True, + ) + + return output + + +def apply_hybrid_apc_suffix_prefill_plan( + input_dict: dict[str, torch.Tensor], + *, + plan: HybridAPCHitPlan, + request_prefix_len: int, + commit_slot: int | None = None, + attention_block_refs: Iterable[int] | None = None, + gdn_active_carry: bool = False, +) -> dict[str, torch.Tensor]: + """Materialize Hybrid APC controls when vLLM already sliced to suffix. + + This diagnostic path is used only when the caller explicitly allows an + unhashed single-checkpoint restore. The input tokens are already the active + suffix, so this helper must not slice token tensors by ``restore_len``. + """ + + if "input_ids" not in input_dict: + raise KeyError("input_ids is required to apply a hybrid APC suffix plan") + + input_ids = input_dict["input_ids"] + if input_ids.ndim != 2: + raise ValueError(f"input_ids must be [batch, seq], got {tuple(input_ids.shape)}") + + batch_size, suffix_len = input_ids.shape + prompt_len = int(request_prefix_len) + restore_len = int(plan.restore_checkpoint_prefix_len) + expected_suffix_len = prompt_len - restore_len + if plan.checkpoint_slot is None or restore_len <= 0: + raise ValueError("suffix-only Hybrid APC restore requires a checkpoint slot") + if expected_suffix_len != suffix_len: + raise ValueError( + "suffix-only Hybrid APC input length mismatch: " + f"expected {expected_suffix_len}, got {suffix_len}" + ) + + output = dict(input_dict) + device = input_ids.device + output["input_ids"] = input_ids + refs: tuple[int, ...] = () + if attention_block_refs is not None: + refs = tuple(int(ref) for ref in attention_block_refs) + if refs: + block_table_template = input_dict.get("block_table") + refs_table = torch.tensor( + [refs] * batch_size, + dtype=torch.int32, + device=device, + ) + has_block_table = ( + isinstance(block_table_template, torch.Tensor) + and block_table_template.numel() > 0 + and block_table_template.ndim >= 2 + and block_table_template.shape[0] >= batch_size + ) + if has_block_table: + active_table = block_table_template[:batch_size].to( + dtype=torch.int32, + device=device, + ) + if active_table.shape[1] > len(refs): + suffix_table = active_table[:, len(refs) :] + else: + suffix_table = active_table + output["block_table"] = torch.cat([refs_table, suffix_table], dim=1) + else: + output["block_table"] = refs_table + + position_template = input_dict.get("position_ids") + position_dtype = ( + position_template.dtype + if isinstance(position_template, torch.Tensor) + else torch.int64 + ) + position_ids = torch.arange( + restore_len, + prompt_len, + dtype=position_dtype, + device=device, + ).unsqueeze(0) + output["position_ids"] = position_ids.expand(batch_size, suffix_len).contiguous() + default_rotary_positions = torch.arange( + restore_len, + prompt_len, + dtype=torch.int32, + device=device, + ) + output["rotary_position_ids"] = default_rotary_positions.view( + 1, + 1, + suffix_len, + ).expand(3, batch_size, suffix_len).contiguous() + + for key in ("rotary_position_id", "rotary_position_ids"): + value = input_dict.get(key) + if not isinstance(value, torch.Tensor): + continue + rotary_positions = torch.arange( + restore_len, + prompt_len, + dtype=value.dtype, + device=device, + ) + if value.ndim == 2: + output[key] = rotary_positions.unsqueeze(0).expand( + batch_size, + suffix_len, + ).contiguous() + elif value.ndim == 3: + output[key] = rotary_positions.view(1, 1, suffix_len).expand( + value.shape[0], + batch_size, + suffix_len, + ).contiguous() + + def _batch_i32(value: int) -> torch.Tensor: + return torch.full((batch_size,), int(value), dtype=torch.int32, device=device) + + def _batch_i32_col(value: int) -> torch.Tensor: + return torch.full((batch_size, 1), int(value), dtype=torch.int32, device=device) + + output["computed_context_lens"] = _batch_i32_col(restore_len) + output["full_context_lens"] = _batch_i32_col(prompt_len) + output["num_queries"] = _batch_i32_col(suffix_len) + output["hybrid_restore_slot_ids"] = _batch_i32(int(plan.checkpoint_slot)) + output["hybrid_restore_mask"] = _batch_i32(0 if gdn_active_carry else 1) + output["hybrid_restore_prefix_lens"] = _batch_i32(restore_len) + output["hybrid_commit_slot_ids"] = _batch_i32(0 if commit_slot is None else commit_slot) + output["hybrid_commit_mask"] = _batch_i32(1 if commit_slot is not None else 0) + + if os.environ.get("QWEN36_HYBRID_APC_DEBUG") == "1": + print( + "[hybrid_apc_debug] apply-suffix " + f"prompt_len={prompt_len} restore_len={restore_len} " + f"suffix_len={suffix_len} restore_slot={plan.checkpoint_slot} " + f"commit_slot={commit_slot} " + f"gdn_active_carry={gdn_active_carry} " + f"attention_block_refs={refs} " + f"input_shape={tuple(input_ids.shape)}", + flush=True, + ) + + return output + + +class HybridAPCSlotAllocator: + """Small checkpoint-slot allocator for local scheduler integration tests.""" + + def __init__(self, num_slots: int): + num_slots = int(num_slots) + if num_slots <= 0: + raise ValueError(f"num_slots must be positive, got {num_slots}") + self.num_slots = num_slots + self._free = deque(range(num_slots)) + self._reserved: set[int] = set() + self._committed: set[int] = set() + + @property + def free_slots(self) -> tuple[int, ...]: + return tuple(self._free) + + @property + def reserved_slots(self) -> tuple[int, ...]: + return tuple(sorted(self._reserved)) + + @property + def committed_slots(self) -> tuple[int, ...]: + return tuple(sorted(self._committed)) + + def reserve(self) -> int: + if not self._free: + raise RuntimeError("no hybrid APC checkpoint slots available") + slot = int(self._free.popleft()) + self._reserved.add(slot) + return slot + + def mark_committed(self, slot: int): + slot = int(slot) + self.validate_slot_range(slot) + if slot not in self._reserved and slot not in self._committed: + raise ValueError(f"hybrid APC checkpoint slot {slot} is not reserved") + self._reserved.discard(slot) + self._committed.add(slot) + + def release(self, slot: int): + slot = int(slot) + self.validate_slot_range(slot) + was_known = slot in self._reserved or slot in self._committed + self._reserved.discard(slot) + self._committed.discard(slot) + if was_known and slot not in self._free: + self._free.append(slot) + + def release_committed(self, slot: int) -> bool: + slot = int(slot) + self.validate_slot_range(slot) + if slot in self._reserved: + return False + was_committed = slot in self._committed + self._committed.discard(slot) + if was_committed and slot not in self._free: + self._free.append(slot) + return was_committed + + def validate_slot_range(self, slot: int): + slot = int(slot) + if slot < 0 or slot >= self.num_slots: + raise ValueError( + f"hybrid APC checkpoint slot {slot} is outside " + f"[0, {self.num_slots})" + ) + + +class HybridAPCSchedulerBridge: + """Local request-prep bridge for production hybrid APC scheduler wiring. + + The real vLLM/NxDI scheduler must supply the attention APC hit length, + active attention block refs, and tenant/cache metadata. This bridge performs + the Qwen hybrid-specific part: intersect attention hits with GDN checkpoint + metadata, materialize suffix model inputs, reserve a GDN checkpoint slot, + and commit checkpoint metadata after a successful prefill. + """ + + def __init__( + self, + *, + store: "HybridAPCMetadataStore", + slot_allocator: HybridAPCSlotAllocator, + cache_salt: Hashable | None = None, + model_revision: str | None = None, + layout_version: int | None = None, + tp_rank: int | None = None, + recurrent_dtype: str | torch.dtype | None = None, + conv_dtype: str | torch.dtype | None = None, + allow_local_hash_fallback: bool = True, + require_attention_block_refs: bool = False, + reject_unbacked_attention_hits: bool = True, + ): + self.store = store + self.slot_allocator = slot_allocator + self.cache_salt = cache_salt + self.model_revision = model_revision + self.layout_version = layout_version + self.tp_rank = tp_rank + self.recurrent_dtype = recurrent_dtype + self.conv_dtype = conv_dtype + self.allow_local_hash_fallback = bool(allow_local_hash_fallback) + self.require_attention_block_refs = bool(require_attention_block_refs) + self.reject_unbacked_attention_hits = bool(reject_unbacked_attention_hits) + self._same_request_committed_keys: dict[Hashable, set[HybridPrefixKey]] = {} + self.store.set_checkpoint_slot_releaser( + self.slot_allocator.release_committed + ) + + @property + def requires_external_metadata(self) -> bool: + return ( + not self.allow_local_hash_fallback + or self.require_attention_block_refs + ) + + def prepare_request( + self, + *, + request_id: Hashable, + input_dict: dict[str, torch.Tensor], + attention_hit_len: int, + request_prefix_len: int | None = None, + cumulative_hashes_by_prefix_len: dict[int, Hashable] | None = None, + attention_block_refs_by_prefix_len: dict[int, Iterable[int]] | None = None, + ) -> HybridAPCPreparedRequest: + if "input_ids" not in input_dict: + raise KeyError("input_ids is required for hybrid APC request prep") + input_ids = input_dict["input_ids"] + prompt_len = ( + int(input_ids.shape[1]) + if request_prefix_len is None + else int(request_prefix_len) + ) + commit_prefix_len = floor_to_checkpoint_boundary( + prompt_len, + self.store.checkpoint_interval, + ) + + if cumulative_hashes_by_prefix_len is None: + if not self.allow_local_hash_fallback: + raise ValueError( + "hybrid APC production mode requires vLLM cumulative prefix " + "hashes; set hybrid_apc_allow_local_hash_fallback=True only " + "for controlled local validation" + ) + cumulative_hashes_by_prefix_len = build_cumulative_prefix_hashes( + input_ids, + block_size=self.store.block_size, + ) + + plan = self.store.compute_hit_plan( + cumulative_hashes_by_prefix_len=cumulative_hashes_by_prefix_len, + attention_hit_len=attention_hit_len, + request_prefix_len=prompt_len, + cache_salt=self.cache_salt, + model_revision=self.model_revision, + layout_version=self.layout_version, + tp_rank=self.tp_rank, + recurrent_dtype=self.recurrent_dtype, + conv_dtype=self.conv_dtype, + ) + disable_restore = _env_flag("QWEN36_DISABLE_HYBRID_GDN_RESTORE") + disable_commit = _env_flag("QWEN36_DISABLE_HYBRID_GDN_COMMIT") + if disable_restore and plan.checkpoint_slot is not None: + plan = HybridAPCHitPlan( + attention_hit_len=0, + recurrent_hit_len=0, + conv_hit_len=0, + usable_hit_len=0, + restore_checkpoint_prefix_len=0, + residual_replay_len=0, + suffix_len=prompt_len, + checkpoint_slot=None, + checkpoint_key=None, + ) + if ( + self.reject_unbacked_attention_hits + and not _env_flag("QWEN36_ALLOW_UNBACKED_HYBRID_APC_FALLBACK") + and not disable_restore + and int(attention_hit_len) > 0 + and plan.checkpoint_slot is None + ): + raise ValueError( + "hybrid APC received an attention prefix hit without a matching " + "GDN checkpoint; scheduler must intersect attention KV hits with " + "GDN checkpoint hits or disable prefix reuse for this request" + ) + if plan.checkpoint_slot is not None: + self.slot_allocator.validate_slot_range(plan.checkpoint_slot) + + commit_key = None + commit_slot = None + attention_block_refs: tuple[int, ...] = () + # The Neuron checkpoint bank can only commit the active GDN state at + # the end of this traced prefill call. Do not label that state as an + # earlier checkpoint boundary unless the current prefill ends exactly + # at that boundary; scheduler-level chunking must create those boundary + # calls. + can_commit_boundary = commit_prefix_len > 0 and commit_prefix_len == prompt_len + if can_commit_boundary and not disable_commit: + if commit_prefix_len not in cumulative_hashes_by_prefix_len: + raise ValueError( + f"missing cumulative prefix hash for commit boundary {commit_prefix_len}" + ) + commit_key = self.store.make_key( + cumulative_prefix_hash=cumulative_hashes_by_prefix_len[commit_prefix_len], + prefix_len=commit_prefix_len, + cache_salt=self.cache_salt, + model_revision=self.model_revision, + layout_version=self.layout_version, + tp_rank=self.tp_rank, + recurrent_dtype=self.recurrent_dtype, + conv_dtype=self.conv_dtype, + ) + if attention_block_refs_by_prefix_len is not None: + attention_block_refs = tuple( + int(ref) + for ref in attention_block_refs_by_prefix_len.get( + commit_prefix_len, + (), + ) + ) + if not attention_block_refs and plan.checkpoint_key is not None: + suffix_refs = tuple( + int(ref) + for ref in attention_block_refs_by_prefix_len.get( + plan.suffix_len, + (), + ) + ) + checkpoint = self.store.lookup(plan.checkpoint_key) + if ( + checkpoint is not None + and suffix_refs + and commit_prefix_len + == checkpoint.prefix_len + plan.suffix_len + ): + attention_block_refs = ( + tuple(int(ref) for ref in checkpoint.attention_block_refs) + + suffix_refs + ) + if not attention_block_refs and not self.require_attention_block_refs: + attention_block_refs = tuple( + range(commit_prefix_len // self.store.block_size) + ) + if self.store.lookup(commit_key) is None: + commit_slot = self._reserve_commit_slot() + + same_request_keys = self._same_request_committed_keys.get(request_id, set()) + existing_record = self.store._requests.get(request_id) + if existing_record is not None: + same_request_keys = same_request_keys | set(existing_record.committed_keys) + gdn_active_carry = ( + plan.checkpoint_key is not None + and plan.checkpoint_key in same_request_keys + ) + if ( + os.environ.get("QWEN36_HYBRID_APC_DEBUG") == "1" + and gdn_active_carry + ): + print( + "[hybrid_apc_debug] prefill-active-carry " + f"request_id={request_id!r} prefix_len={plan.restore_checkpoint_prefix_len} " + f"slot={plan.checkpoint_slot}", + flush=True, + ) + + model_inputs = apply_hybrid_apc_prefill_plan( + input_dict, + plan=plan, + commit_slot=commit_slot, + request_prefix_len=prompt_len, + block_size=self.store.block_size, + gdn_active_carry=gdn_active_carry, + ) + record = self.store.on_request_restore( + request_id=request_id, + checkpoint_key=plan.checkpoint_key, + ) + if commit_slot is not None: + record.reserved_slots.append(commit_slot) + self.store.on_prefill_running(request_id) + + return HybridAPCPreparedRequest( + request_id=request_id, + input_dict=model_inputs, + plan=plan, + commit_prefix_len=commit_prefix_len, + commit_key=commit_key, + commit_slot=commit_slot, + attention_block_refs=attention_block_refs, + ) + + def prepare_suffix_only_request( + self, + *, + request_id: Hashable, + input_dict: dict[str, torch.Tensor], + attention_hit_len: int, + request_prefix_len: int, + cumulative_hashes_by_prefix_len: dict[int, Hashable] | None = None, + attention_block_refs_by_prefix_len: dict[int, Iterable[int]] | None = None, + ) -> HybridAPCPreparedRequest | None: + """Prepare a suffix-only request using scheduler-approved restore metadata.""" + + if "input_ids" not in input_dict: + raise KeyError("input_ids is required for hybrid APC request prep") + input_ids = input_dict["input_ids"] + if input_ids.ndim != 2: + raise ValueError( + f"input_ids must be [batch, seq], got {tuple(input_ids.shape)}" + ) + request_prefix_len = int(request_prefix_len) + attention_hit_len = max(0, int(attention_hit_len)) + suffix_len = int(input_ids.shape[1]) + restore_len = min(attention_hit_len, request_prefix_len) + restore_len = floor_to_checkpoint_boundary( + restore_len, + self.store.checkpoint_interval, + ) + if restore_len <= 0 or request_prefix_len - restore_len != suffix_len: + return None + + checkpoint = None + checkpoint_key = None + try: + from qwen36_hybrid_apc_scheduler_patch import ( # noqa: WPS433 + pop_hybrid_apc_authorized_prefix_key, + ) + except Exception: + pop_hybrid_apc_authorized_prefix_key = None + + if pop_hybrid_apc_authorized_prefix_key is not None: + checkpoint_key = pop_hybrid_apc_authorized_prefix_key( + prefix_len=restore_len, + request_id=request_id, + cache_salt=self.cache_salt, + model_revision=self.model_revision or self.store.model_revision, + layout_version=( + self.layout_version + if self.layout_version is not None + else self.store.layout_version + ), + tp_rank=self.tp_rank if self.tp_rank is not None else self.store.tp_rank, + recurrent_dtype=( + self.recurrent_dtype + if self.recurrent_dtype is not None + else self.store.recurrent_dtype + ), + conv_dtype=( + self.conv_dtype + if self.conv_dtype is not None + else self.store.conv_dtype + ), + ) + if checkpoint_key is not None: + checkpoint = self.store.lookup(checkpoint_key) + + if checkpoint is None and checkpoint_key is not None: + raise ValueError( + "suffix-only hybrid APC received a scheduler-authorized " + "prefix key that is missing from the GDN checkpoint store" + ) + + if checkpoint is None and _env_flag( + "QWEN36_HYBRID_APC_ALLOW_UNHASHED_SINGLE_PREFIX_RESTORE" + ): + checkpoint = self.store.lookup_unique_prefix_len( + prefix_len=restore_len, + cache_salt=self.cache_salt, + model_revision=self.model_revision, + layout_version=self.layout_version, + tp_rank=self.tp_rank, + recurrent_dtype=self.recurrent_dtype, + conv_dtype=self.conv_dtype, + ) + if checkpoint is None: + if self.reject_unbacked_attention_hits: + raise ValueError( + "suffix-only hybrid APC received an attention prefix hit " + "without scheduler-authorized GDN checkpoint metadata" + ) + return None + + plan = HybridAPCHitPlan( + attention_hit_len=attention_hit_len, + recurrent_hit_len=checkpoint.prefix_len, + conv_hit_len=checkpoint.prefix_len, + usable_hit_len=checkpoint.prefix_len, + restore_checkpoint_prefix_len=checkpoint.prefix_len, + residual_replay_len=0, + suffix_len=suffix_len, + checkpoint_slot=checkpoint.gdn_checkpoint_slot, + checkpoint_key=checkpoint.key, + ) + disable_commit = _env_flag("QWEN36_DISABLE_HYBRID_GDN_COMMIT") + commit_prefix_len = floor_to_checkpoint_boundary( + request_prefix_len, + self.store.checkpoint_interval, + ) + commit_key = None + commit_slot = None + attention_block_refs: tuple[int, ...] = () + can_commit_boundary = ( + commit_prefix_len > 0 + and commit_prefix_len == request_prefix_len + and not disable_commit + ) + if can_commit_boundary: + if cumulative_hashes_by_prefix_len is None: + if not self.allow_local_hash_fallback: + raise ValueError( + "hybrid APC production mode requires vLLM cumulative prefix " + f"hashes to commit suffix-only boundary {commit_prefix_len}" + ) + elif commit_prefix_len not in cumulative_hashes_by_prefix_len: + raise ValueError( + f"missing cumulative prefix hash for commit boundary {commit_prefix_len}" + ) + if cumulative_hashes_by_prefix_len is None: + can_commit_boundary = False + if can_commit_boundary: + commit_key = self.store.make_key( + cumulative_prefix_hash=cumulative_hashes_by_prefix_len[commit_prefix_len], + prefix_len=commit_prefix_len, + cache_salt=self.cache_salt, + model_revision=self.model_revision, + layout_version=self.layout_version, + tp_rank=self.tp_rank, + recurrent_dtype=self.recurrent_dtype, + conv_dtype=self.conv_dtype, + ) + if attention_block_refs_by_prefix_len is not None: + attention_block_refs = tuple( + int(ref) + for ref in attention_block_refs_by_prefix_len.get( + commit_prefix_len, + (), + ) + ) + if not attention_block_refs: + suffix_refs = tuple( + int(ref) + for ref in attention_block_refs_by_prefix_len.get( + suffix_len, + (), + ) + ) + if ( + suffix_refs + and commit_prefix_len == checkpoint.prefix_len + suffix_len + ): + attention_block_refs = ( + tuple(int(ref) for ref in checkpoint.attention_block_refs) + + suffix_refs + ) + if not attention_block_refs and not self.require_attention_block_refs: + attention_block_refs = tuple( + range(commit_prefix_len // self.store.block_size) + ) + if self.store.lookup(commit_key) is None: + commit_slot = self._reserve_commit_slot() + + same_request_keys = self._same_request_committed_keys.get(request_id, set()) + existing_record = self.store._requests.get(request_id) + if existing_record is not None: + same_request_keys = same_request_keys | set(existing_record.committed_keys) + gdn_active_carry = checkpoint.key in same_request_keys + if ( + os.environ.get("QWEN36_HYBRID_APC_DEBUG") == "1" + and gdn_active_carry + ): + print( + "[hybrid_apc_debug] suffix-active-carry " + f"request_id={request_id!r} prefix_len={checkpoint.prefix_len} " + f"slot={checkpoint.gdn_checkpoint_slot}", + flush=True, + ) + + model_inputs = apply_hybrid_apc_suffix_prefill_plan( + input_dict, + plan=plan, + request_prefix_len=request_prefix_len, + commit_slot=commit_slot, + attention_block_refs=checkpoint.attention_block_refs, + gdn_active_carry=gdn_active_carry, + ) + record = self.store.on_request_restore( + request_id=request_id, + checkpoint_key=plan.checkpoint_key, + ) + if commit_slot is not None: + record.reserved_slots.append(commit_slot) + self.store.on_prefill_running(request_id) + + return HybridAPCPreparedRequest( + request_id=request_id, + input_dict=model_inputs, + plan=plan, + commit_prefix_len=commit_prefix_len, + commit_key=commit_key, + commit_slot=commit_slot, + attention_block_refs=attention_block_refs or checkpoint.attention_block_refs, + ) + + def commit_prefill( + self, + prepared: HybridAPCPreparedRequest, + *, + attention_block_refs: Iterable[int] | None = None, + bytes_used: int = 0, + ) -> HybridPrefixCheckpoint | None: + if prepared.commit_key is None or prepared.commit_slot is None: + return None + refs = ( + tuple(int(ref) for ref in attention_block_refs) + if attention_block_refs is not None + else prepared.attention_block_refs + ) + if self.require_attention_block_refs and not refs: + raise ValueError( + "hybrid APC checkpoint commit requires real attention block refs " + "from the vLLM/NxDI APC allocator" + ) + checkpoint = self.store.insert( + key=prepared.commit_key, + attention_block_refs=refs, + gdn_checkpoint_slot=prepared.commit_slot, + bytes_used=bytes_used, + ) + _publish_scheduler_gdn_checkpoint(prepared.commit_key) + self.slot_allocator.mark_committed(prepared.commit_slot) + record = self.store.on_checkpoint_committed( + request_id=prepared.request_id, + checkpoint_key=prepared.commit_key, + ) + self._same_request_committed_keys.setdefault( + prepared.request_id, + set(), + ).add(prepared.commit_key) + if len(self._same_request_committed_keys) > 4096: + self._same_request_committed_keys.clear() + if prepared.commit_slot in record.reserved_slots: + record.reserved_slots.remove(prepared.commit_slot) + return checkpoint + + def _reserve_commit_slot(self) -> int: + try: + return self.slot_allocator.reserve() + except RuntimeError: + target_checkpoints = self.slot_allocator.num_slots - 1 + if self.store.max_checkpoints is not None: + target_checkpoints = min( + target_checkpoints, + int(self.store.max_checkpoints) - 1, + ) + evicted = self.store.evict_lru( + target_checkpoints=max(0, target_checkpoints) + ) + if evicted: + return self.slot_allocator.reserve() + raise + + def finish_request(self, request_id: Hashable) -> HybridAPCRequestRecord | None: + record = self.store.on_request_finish(request_id) + if record is not None: + for slot in record.reserved_slots: + self.slot_allocator.release(slot) + return record + + def cancel_request( + self, + prepared: HybridAPCPreparedRequest, + ) -> HybridAPCRequestRecord | None: + record = self.store._requests.get(prepared.request_id) + if ( + prepared.commit_slot is not None + and record is not None + and prepared.commit_slot in record.reserved_slots + ): + self.slot_allocator.release(prepared.commit_slot) + return self.store.on_request_cancel(prepared.request_id) + + +class HybridAPCMetadataStore: + """CPU-side lifecycle store for hybrid prefix-boundary checkpoints.""" + + def __init__( + self, + *, + required_gdn_layers: Iterable[int], + block_size: int, + checkpoint_interval: int | None = None, + max_checkpoints: int | None = None, + max_bytes: int | None = None, + layout_version: int = 1, + model_revision: str = "unknown", + tp_rank: int = 0, + recurrent_dtype: str | torch.dtype = "float32", + conv_dtype: str | torch.dtype = "bfloat16", + allow_residual_replay: bool = False, + checkpoint_slot_releaser: Callable[[int], object] | None = None, + ): + self.required_gdn_layers = tuple(sorted({int(x) for x in required_gdn_layers})) + if not self.required_gdn_layers: + raise ValueError("required_gdn_layers must not be empty") + self.num_layer_mask_bits = max(self.required_gdn_layers) + 1 + self.block_size = int(block_size) + if self.block_size <= 0: + raise ValueError(f"block_size must be positive, got {block_size}") + self.checkpoint_interval = ( + self.block_size + if checkpoint_interval is None + else int(checkpoint_interval) + ) + if self.checkpoint_interval <= 0: + raise ValueError( + f"checkpoint_interval must be positive, got {checkpoint_interval}" + ) + if self.checkpoint_interval % self.block_size != 0: + raise ValueError( + "checkpoint_interval must be a multiple of block_size for v0 " + f"hybrid APC, got {self.checkpoint_interval} and {self.block_size}" + ) + self.max_checkpoints = max_checkpoints + if self.max_checkpoints is not None and self.max_checkpoints <= 0: + raise ValueError(f"max_checkpoints must be positive, got {max_checkpoints}") + self.max_bytes = max_bytes + if self.max_bytes is not None and self.max_bytes <= 0: + raise ValueError(f"max_bytes must be positive, got {max_bytes}") + self.layout_version = int(layout_version) + self.model_revision = str(model_revision) + self.tp_rank = int(tp_rank) + self.recurrent_dtype = _normalize_dtype(recurrent_dtype) + self.conv_dtype = _normalize_dtype(conv_dtype) + self.allow_residual_replay = bool(allow_residual_replay) + self._checkpoint_slot_releaser = checkpoint_slot_releaser + + self._by_key: OrderedDict[HybridPrefixKey, HybridPrefixCheckpoint] = ( + OrderedDict() + ) + self._slot_to_key: dict[int, HybridPrefixKey] = {} + self._requests: dict[Hashable, HybridAPCRequestRecord] = {} + self._step = 0 + self.stats = HybridAPCStats() + + def set_checkpoint_slot_releaser( + self, + releaser: Callable[[int], object] | None, + ): + self._checkpoint_slot_releaser = releaser + + def __len__(self) -> int: + return len(self._by_key) + + @property + def bytes_used(self) -> int: + return sum(checkpoint.bytes_used for checkpoint in self._by_key.values()) + + def _next_step(self) -> int: + self._step += 1 + return self._step + + def make_key( + self, + *, + cumulative_prefix_hash: Hashable, + prefix_len: int, + cache_salt: Hashable | None = None, + model_revision: str | None = None, + layout_version: int | None = None, + tp_rank: int | None = None, + recurrent_dtype: str | torch.dtype | None = None, + conv_dtype: str | torch.dtype | None = None, + ) -> HybridPrefixKey: + prefix_len = int(prefix_len) + if prefix_len < 0: + raise ValueError(f"prefix_len must be non-negative, got {prefix_len}") + if prefix_len % self.checkpoint_interval != 0: + raise ValueError( + "prefix_len must align to checkpoint_interval " + f"{self.checkpoint_interval}, got {prefix_len}" + ) + return HybridPrefixKey( + cumulative_prefix_hash=cumulative_prefix_hash, + prefix_len=prefix_len, + block_size=self.block_size, + cache_salt=cache_salt, + model_revision=self.model_revision + if model_revision is None + else str(model_revision), + layout_version=self.layout_version + if layout_version is None + else int(layout_version), + tp_rank=self.tp_rank if tp_rank is None else int(tp_rank), + recurrent_dtype=self.recurrent_dtype + if recurrent_dtype is None + else _normalize_dtype(recurrent_dtype), + conv_dtype=self.conv_dtype if conv_dtype is None else _normalize_dtype(conv_dtype), + ) + + def _make_mask(self, valid_layers: torch.Tensor | int | Iterable[int] | None): + if valid_layers is None: + layers = self.required_gdn_layers + mask = torch.zeros(self.num_layer_mask_bits, dtype=torch.bool) + mask[list(layers)] = True + return mask + if isinstance(valid_layers, torch.Tensor): + mask = valid_layers.detach().cpu().to(torch.bool).flatten().clone() + if mask.numel() < self.num_layer_mask_bits: + padded = torch.zeros(self.num_layer_mask_bits, dtype=torch.bool) + padded[: mask.numel()] = mask + mask = padded + return mask + mask = torch.zeros(self.num_layer_mask_bits, dtype=torch.bool) + if isinstance(valid_layers, int): + bitmask = int(valid_layers) + for layer in range(self.num_layer_mask_bits): + mask[layer] = bool(bitmask & (1 << layer)) + return mask + for layer in valid_layers: + layer = int(layer) + if layer >= mask.numel(): + padded = torch.zeros(layer + 1, dtype=torch.bool) + padded[: mask.numel()] = mask + mask = padded + mask[layer] = True + return mask + + def insert( + self, + *, + key: HybridPrefixKey, + attention_block_refs: Iterable[int], + gdn_checkpoint_slot: int, + valid_recurrent_layers: torch.Tensor | int | Iterable[int] | None = None, + valid_conv_layers: torch.Tensor | int | Iterable[int] | None = None, + bytes_used: int = 0, + evictable: bool = True, + ) -> HybridPrefixCheckpoint: + if key.block_size != self.block_size: + raise ValueError( + f"key block_size {key.block_size} does not match store block_size {self.block_size}" + ) + if key.layout_version != self.layout_version: + raise ValueError( + f"key layout_version {key.layout_version} does not match store layout_version {self.layout_version}" + ) + recurrent_mask = self._make_mask(valid_recurrent_layers) + conv_mask = self._make_mask(valid_conv_layers) + checkpoint = HybridPrefixCheckpoint( + key=key, + prefix_len=key.prefix_len, + attention_block_refs=tuple(int(ref) for ref in attention_block_refs), + gdn_checkpoint_slot=int(gdn_checkpoint_slot), + valid_recurrent_layers=recurrent_mask, + valid_conv_layers=conv_mask, + last_access_step=self._next_step(), + bytes_used=int(bytes_used), + evictable=bool(evictable), + ) + if not checkpoint.has_valid_gdn(self.required_gdn_layers): + raise ValueError("checkpoint is missing recurrent or conv state") + + old = self._slot_to_key.get(checkpoint.gdn_checkpoint_slot) + if old is not None and old != key: + self.mark_invalid(old) + if key in self._by_key: + old_checkpoint = self._by_key[key] + self._slot_to_key.pop(old_checkpoint.gdn_checkpoint_slot, None) + if ( + old_checkpoint.gdn_checkpoint_slot != checkpoint.gdn_checkpoint_slot + and self._checkpoint_slot_releaser is not None + ): + self._checkpoint_slot_releaser(old_checkpoint.gdn_checkpoint_slot) + self._by_key[key] = checkpoint + self._by_key.move_to_end(key) + self._slot_to_key[checkpoint.gdn_checkpoint_slot] = key + self._evict_over_budget() + self._refresh_stats() + return checkpoint + + def lookup( + self, + key: HybridPrefixKey, + *, + require_attention: bool = True, + require_gdn: bool = True, + ) -> HybridPrefixCheckpoint | None: + checkpoint = self._by_key.get(key) + if checkpoint is None: + self.stats.misses += 1 + return None + if require_attention and not checkpoint.attention_valid: + self.stats.misses += 1 + return None + if require_gdn and not checkpoint.has_valid_gdn(self.required_gdn_layers): + self.stats.misses += 1 + return None + checkpoint.last_access_step = self._next_step() + self._by_key.move_to_end(key) + self.stats.hits += 1 + return checkpoint + + def lookup_unique_prefix_len( + self, + *, + prefix_len: int, + cache_salt: Hashable | None = None, + model_revision: str | None = None, + layout_version: int | None = None, + tp_rank: int | None = None, + recurrent_dtype: str | torch.dtype | None = None, + conv_dtype: str | torch.dtype | None = None, + ) -> HybridPrefixCheckpoint | None: + """Return the only valid checkpoint at a prefix length, if unambiguous.""" + + prefix_len = int(prefix_len) + model_revision = self.model_revision if model_revision is None else str(model_revision) + layout_version = self.layout_version if layout_version is None else int(layout_version) + tp_rank = self.tp_rank if tp_rank is None else int(tp_rank) + recurrent_dtype = ( + self.recurrent_dtype + if recurrent_dtype is None + else _normalize_dtype(recurrent_dtype) + ) + conv_dtype = self.conv_dtype if conv_dtype is None else _normalize_dtype(conv_dtype) + + candidates: list[HybridPrefixKey] = [] + for key, checkpoint in self._by_key.items(): + if key.prefix_len != prefix_len: + continue + if key.cache_salt != cache_salt: + continue + if key.model_revision != model_revision: + continue + if key.layout_version != layout_version: + continue + if key.tp_rank != tp_rank: + continue + if key.recurrent_dtype != recurrent_dtype or key.conv_dtype != conv_dtype: + continue + if not checkpoint.attention_valid: + continue + if not checkpoint.has_valid_gdn(self.required_gdn_layers): + continue + candidates.append(key) + + if not candidates: + self.stats.misses += 1 + return None + if len(candidates) > 1: + raise ValueError( + "ambiguous unhashed Hybrid APC restore: " + f"{len(candidates)} checkpoints match prefix_len={prefix_len}" + ) + return self.lookup(candidates[0]) + + def mark_invalid( + self, + key: HybridPrefixKey | None = None, + *, + checkpoint_slot: int | None = None, + state_kind: str | None = None, + layer_id: int | None = None, + ) -> bool: + if key is None: + if checkpoint_slot is None: + raise ValueError("key or checkpoint_slot is required") + key = self._slot_to_key.get(int(checkpoint_slot)) + if key is None: + return False + checkpoint = self._by_key.get(key) + if checkpoint is None: + return False + + if state_kind is None: + self._delete_checkpoint(key) + self._refresh_stats() + return True + if state_kind == "attention": + checkpoint.attention_valid = False + elif state_kind == "recurrent": + if layer_id is None: + checkpoint.valid_recurrent_layers.zero_() + elif int(layer_id) < checkpoint.valid_recurrent_layers.numel(): + checkpoint.valid_recurrent_layers[int(layer_id)] = False + elif state_kind == "conv": + if layer_id is None: + checkpoint.valid_conv_layers.zero_() + elif int(layer_id) < checkpoint.valid_conv_layers.numel(): + checkpoint.valid_conv_layers[int(layer_id)] = False + else: + raise ValueError(f"unknown state_kind: {state_kind}") + _unpublish_scheduler_gdn_checkpoint(key) + return True + + def inc_ref(self, key: HybridPrefixKey) -> int: + checkpoint = self.lookup(key, require_attention=False, require_gdn=False) + if checkpoint is None: + raise KeyError(key) + checkpoint.refcount += 1 + return checkpoint.refcount + + def dec_ref(self, key: HybridPrefixKey) -> int: + checkpoint = self.lookup(key, require_attention=False, require_gdn=False) + if checkpoint is None: + raise KeyError(key) + checkpoint.refcount = max(0, checkpoint.refcount - 1) + return checkpoint.refcount + + def on_request_restore( + self, + *, + request_id: Hashable, + checkpoint_key: HybridPrefixKey | None, + ) -> HybridAPCRequestRecord: + record = HybridAPCRequestRecord( + request_id=request_id, + state="NEW", + restored_key=checkpoint_key, + ) + if checkpoint_key is not None: + self.inc_ref(checkpoint_key) + record.state = "RESTORED_FROM_HYBRID_APC" + self._requests[request_id] = record + return record + + def on_prefill_running(self, request_id: Hashable) -> HybridAPCRequestRecord: + record = self._requests[request_id] + record.state = "PREFILL_RUNNING" + return record + + def on_checkpoint_committed( + self, + *, + request_id: Hashable, + checkpoint_key: HybridPrefixKey, + ) -> HybridAPCRequestRecord: + record = self._requests.setdefault( + request_id, + HybridAPCRequestRecord(request_id=request_id, state="PREFILL_RUNNING"), + ) + record.state = "PREFILL_COMMIT_PENDING" + record.committed_keys.append(checkpoint_key) + return record + + def on_decode_running(self, request_id: Hashable) -> HybridAPCRequestRecord: + record = self._requests[request_id] + record.state = "DECODE_RUNNING" + return record + + def on_request_finish(self, request_id: Hashable) -> HybridAPCRequestRecord | None: + record = self._requests.pop(request_id, None) + if record is None: + return None + if record.restored_key is not None and record.restored_key in self._by_key: + self.dec_ref(record.restored_key) + record.state = "FINISHED" + return record + + def on_request_cancel(self, request_id: Hashable) -> HybridAPCRequestRecord | None: + record = self._requests.pop(request_id, None) + if record is None: + return None + if record.restored_key is not None and record.restored_key in self._by_key: + self.dec_ref(record.restored_key) + for key in record.committed_keys: + if key in self._by_key: + self.mark_invalid(key) + record.state = "CANCELLED" + return record + + def evict_lru(self, *, target_checkpoints: int | None = None) -> list[HybridPrefixKey]: + target = self.max_checkpoints if target_checkpoints is None else target_checkpoints + if target is None: + return [] + evicted: list[HybridPrefixKey] = [] + for key, checkpoint in list(self._by_key.items()): + if len(self._by_key) <= target: + break + if checkpoint.refcount > 0 or not checkpoint.evictable: + continue + self._delete_checkpoint(key) + evicted.append(key) + self.stats.evictions += len(evicted) + self._refresh_stats() + return evicted + + def on_attention_block_evicted(self, block_ref: int) -> list[HybridPrefixKey]: + invalidated: list[HybridPrefixKey] = [] + for key, checkpoint in self._by_key.items(): + if int(block_ref) in checkpoint.attention_block_refs: + checkpoint.attention_valid = False + _unpublish_scheduler_gdn_checkpoint(key) + invalidated.append(key) + return invalidated + + def on_gdn_checkpoint_evicted(self, checkpoint_slot: int) -> bool: + return self.mark_invalid(checkpoint_slot=int(checkpoint_slot)) + + def compute_hit_plan( + self, + *, + cumulative_hashes_by_prefix_len: dict[int, Hashable], + attention_hit_len: int, + request_prefix_len: int, + cache_salt: Hashable | None = None, + model_revision: str | None = None, + layout_version: int | None = None, + tp_rank: int | None = None, + recurrent_dtype: str | torch.dtype | None = None, + conv_dtype: str | torch.dtype | None = None, + ) -> HybridAPCHitPlan: + attention_hit_len = max(0, int(attention_hit_len)) + request_prefix_len = max(0, int(request_prefix_len)) + target_hit_len = min(attention_hit_len, request_prefix_len) + candidate_lens = sorted( + ( + int(prefix_len) + for prefix_len in cumulative_hashes_by_prefix_len + if int(prefix_len) <= target_hit_len + and int(prefix_len) % self.checkpoint_interval == 0 + ), + reverse=True, + ) + + for prefix_len in candidate_lens: + key = self.make_key( + cumulative_prefix_hash=cumulative_hashes_by_prefix_len[prefix_len], + prefix_len=prefix_len, + cache_salt=cache_salt, + model_revision=model_revision, + layout_version=layout_version, + tp_rank=tp_rank, + recurrent_dtype=recurrent_dtype, + conv_dtype=conv_dtype, + ) + checkpoint = self.lookup(key) + if checkpoint is None: + continue + + if self.allow_residual_replay: + usable_hit_len = target_hit_len + residual_replay_len = target_hit_len - prefix_len + suffix_len = request_prefix_len - target_hit_len + else: + usable_hit_len = prefix_len + residual_replay_len = 0 + suffix_len = request_prefix_len - prefix_len + return HybridAPCHitPlan( + attention_hit_len=attention_hit_len, + recurrent_hit_len=prefix_len, + conv_hit_len=prefix_len, + usable_hit_len=usable_hit_len, + restore_checkpoint_prefix_len=prefix_len, + residual_replay_len=residual_replay_len, + suffix_len=suffix_len, + checkpoint_slot=checkpoint.gdn_checkpoint_slot, + checkpoint_key=key, + ) + + return HybridAPCHitPlan( + attention_hit_len=attention_hit_len, + recurrent_hit_len=0, + conv_hit_len=0, + usable_hit_len=0, + restore_checkpoint_prefix_len=0, + residual_replay_len=0, + suffix_len=request_prefix_len, + checkpoint_slot=None, + checkpoint_key=None, + ) + + def _evict_over_budget(self): + if self.max_checkpoints is not None: + self.evict_lru(target_checkpoints=self.max_checkpoints) + if self.max_bytes is None: + return + evicted = 0 + for key, checkpoint in list(self._by_key.items()): + if self.bytes_used <= self.max_bytes: + break + if checkpoint.refcount > 0 or not checkpoint.evictable: + continue + self._delete_checkpoint(key) + evicted += 1 + self.stats.evictions += evicted + + def _delete_checkpoint(self, key: HybridPrefixKey): + checkpoint = self._by_key.pop(key, None) + if checkpoint is not None: + self._slot_to_key.pop(checkpoint.gdn_checkpoint_slot, None) + _unpublish_scheduler_gdn_checkpoint(key) + if self._checkpoint_slot_releaser is not None: + self._checkpoint_slot_releaser(checkpoint.gdn_checkpoint_slot) + + def _refresh_stats(self): + self.stats.checkpoints = len(self._by_key) + self.stats.bytes_used = self.bytes_used diff --git a/contrib/models/Qwen3.5-35B-A3B/src/modeling_qwen35.py b/contrib/models/Qwen3.5-35B-A3B/src/modeling_qwen35.py new file mode 100644 index 00000000..cc76b07b --- /dev/null +++ b/contrib/models/Qwen3.5-35B-A3B/src/modeling_qwen35.py @@ -0,0 +1,8378 @@ +""" +NxDI contrib: Qwen3.5-27B / Qwen3.6-27B (qwen3_5 -- dense model) + +Supports both Qwen3.5-27B and Qwen3.6-27B. These models share identical +architecture (qwen3_5 model_type). Qwen3.6-27B is a post-training update +with improved agentic coding and thinking preservation -- no architecture +changes, only weight differences. + +Hybrid DeltaNet + Standard Attention + Dense MLP architecture. +Adapted from Qwen3.5-35B-A3B (MoE) -- MoE removed, dense MLP added. + +48 of 64 layers use Gated DeltaNet (linear recurrent attention) +16 of 64 layers use standard GQA with KV cache + output gate +All 64 layers use a dense SwiGLU MLP (intermediate_size=17408) + +Architecture details: +- DeltaNet layers: separate in_proj_{qkv, z, a, b}, causal conv1d on QKV, gated delta rule +- Attention layers: q_proj doubled (Q + gate), partial RoPE (25% of head_dim), sigmoid output gate +- Dense MLP: standard SwiGLU (gate_proj, up_proj, down_proj) -- no MoE, no router, no experts +- KV cache: NxDI KVCacheManager for attention layers; DeltaNet layers store recurrent+conv + state as nn.Parameter buffers and return dummy KV tuples + +Config compatibility notes: +- Qwen3.6-27B adds output_gate_type="swish" to text_config. This field is + unused by both HF transformers and this NxDI code (gate uses sigmoid, as + confirmed across transformers v4.57.6, v5.6.0, and GitHub main). Safe to ignore. +""" + +import copy +import gc +import json +import math +import logging +import os +import re +import sys +import time +from typing import Any, Hashable, List, NamedTuple, Optional, Tuple + +import torch +import torch.nn.functional as F +from torch import nn + +from neuronx_distributed_inference.models.model_base import ( + NeuronBaseForCausalLM, + NeuronBaseModel, +) +try: + from neuronx_distributed_inference.modules.async_execution import ( + cancel_hybrid_apc_request, + finish_hybrid_apc_request, + prepare_hybrid_apc_model_inputs, + prepare_hybrid_apc_request_for_execution, + ) +except ImportError: + def cancel_hybrid_apc_request(*a, **kw): + return None + + def finish_hybrid_apc_request(*a, **kw): + return None + + def prepare_hybrid_apc_model_inputs(*a, **kw): + return None + + def prepare_hybrid_apc_request_for_execution(*a, **kw): + return None +from neuronx_distributed_inference.modules.custom_calls import CustomRMSNorm + +try: + from neuronxcc.nki._private_kernels.attention import attention_isa_kernel +except ImportError: + from neuronxcc.nki.kernels.attention import attention_isa_kernel + +from neuronx_distributed.parallel_layers import parallel_state +from neuronx_distributed.parallel_layers.layers import ( + ColumnParallelLinear, + ParallelEmbedding, + RowParallelLinear, +) +from neuronx_distributed.parallel_layers.mappings import _gather_along_dim +from neuronx_distributed.utils import cpu_mode + +try: + from nki import jit as nki_jit # standalone nki package (>=0.3.0) +except ImportError: + from torch_neuronx.xla_impl.ops import nki_jit # legacy embedded path +from transformers.models.qwen3_moe.modeling_qwen3_moe import Qwen3MoeRMSNorm + +from src.nki_kernels.nki_deltanet import deltanet_recurrent_fwd as _deltanet_nki_kernel +from src.nki_kernels.nki_deltanet import ( + deltanet_recurrent_fwd_state as _deltanet_nki_kernel_state, +) +from src.nki_kernels.nki_deltanet import ( + deltanet_recurrent_step_batched as _deltanet_nki_step_batched, +) +from src.nki_kernels.nki_deltanet_chunked import ( + deltanet_chunk_step as _deltanet_nki_chunk_step, +) +from src.nki_kernels.nki_deltanet_fused import ( + deltanet_autocp_affine_sequence as _deltanet_autocp_affine_sequence, + deltanet_autocp_apply_output as _deltanet_autocp_apply_output, + deltanet_autocp_prefix_apply_output as _deltanet_autocp_prefix_apply_output, + deltanet_autocp_state_summary_sequence as _deltanet_autocp_state_summary_sequence, + deltanet_autocp_state_prefix as _deltanet_autocp_state_prefix, + deltanet_fused_chunked_fwd as _deltanet_fused_kernel, + deltanet_fused_chunked_fwd_multihead as _deltanet_fused_multihead_kernel, +) +from src.nki_kernels.nki_deltanet_fused_legacy import ( + deltanet_fused_chunked_fwd as _deltanet_fused_legacy_direct_kernel, +) +from src.nki_kernels.nki_deltanet_fused import ( + _make_lower_mask, + _make_lower_mask_diag, + _make_identity, +) +try: + import nki as _nkilib_nki + from nkilib.core.qkv.qkv import qkv as _nkilib_qkv + from nkilib.core.utils.common_types import ( + NormType as _NkilibNormType, + QKVOutputLayout as _NkilibQKVOutputLayout, + QuantizationType as _NkilibQuantizationType, + ) + + _qwen_gate_projection_kernel = _nkilib_nki.jit(_nkilib_qkv) +except Exception: + _NkilibNormType = None + _NkilibQKVOutputLayout = None + _NkilibQuantizationType = None + _qwen_gate_projection_kernel = None + +try: + from src.nki_kernels.qwen_qk_norm_rope import ( + qwen_qk_norm_partial_rope_kernel as _qwen_qk_norm_partial_rope_kernel, + ) +except Exception: + _qwen_qk_norm_partial_rope_kernel = None +from src.hybrid_apc import ( + HybridAPCMetadataStore, + HybridAPCSchedulerBridge, + HybridAPCSlotAllocator, +) + + +def _infer_neuron_lnc(default: int = 1) -> int: + flags = os.environ.get("NEURON_CC_FLAGS", "") + match = re.search(r"(?:^|\s)--lnc(?:=|\s+)(\d+)", flags) + if match is None: + return default + return max(1, int(match.group(1))) + + +def _resolve_deltanet_multihead_group_size(total_heads: int) -> int: + lnc = _infer_neuron_lnc() + raw_group_size = os.environ.get("QWEN36_DELTANET_MULTIHEAD_GROUP_SIZE") + if raw_group_size is None: + requested_group_size = 2 if lnc >= 2 else 1 + else: + requested_group_size = max(1, int(raw_group_size)) + if requested_group_size > lnc: + raise ValueError( + f"QWEN36_DELTANET_MULTIHEAD_GROUP_SIZE={requested_group_size} " + f"requires NEURON_CC_FLAGS --lnc >= {requested_group_size}; " + f"inferred lnc={lnc}" + ) + return max(1, min(total_heads, requested_group_size)) + + +def _deltanet_multihead_launch_spec(num_heads: int): + """Return the launch spec for a grouped multihead DeltaNet CTE kernel. + + The legacy ``kernel[2]`` launch only covers two programs. For larger + grouped launches we need an SPMD axis distributed over the available NCs, + while each program still handles exactly one flattened (batch, head) row. + """ + lnc = _infer_neuron_lnc() + if num_heads <= lnc: + return num_heads + if os.environ.get("QWEN36_DELTANET_MULTIHEAD_SPMD", "1") == "0": + raise ValueError( + "QWEN36_DELTANET_MULTIHEAD_GROUP_SIZE exceeds inferred LNC but " + "QWEN36_DELTANET_MULTIHEAD_SPMD=0; " + f"group_size={num_heads}, inferred_lnc={lnc}" + ) + + import nki.language as _nl # Imported lazily so CPU-only unit stubs still load. + + if not hasattr(_nl, "spmd_dim") or not hasattr(_nl, "nc"): + if os.environ.get("QWEN36_DELTANET_MULTIHEAD_GRID_FALLBACK", "0") == "1": + return (num_heads, 1) + raise ValueError( + "QWEN36_DELTANET_MULTIHEAD_GROUP_SIZE exceeds inferred LNC, but " + "this NKI runtime does not expose spmd_dim/nc; " + f"group_size={num_heads}, inferred_lnc={lnc}" + ) + return (_nl.spmd_dim(num_heads, _nl.nc(lnc)),) + + +def _qwen35_grouped_prefix_attention( + Q, + K_cache, + V_cache, + query_positions, + cache_positions, + key_valid_mask=None, +): + """GQA-native prefix attention without materializing repeated KV heads.""" + B, q_heads, q_len, head_dim = Q.shape + kv_heads = K_cache.shape[1] + if q_heads % kv_heads != 0: + raise ValueError( + "Qwen grouped prefix attention requires q_heads to be divisible " + f"by kv_heads, got q_heads={q_heads}, kv_heads={kv_heads}." + ) + + q_per_kv = q_heads // kv_heads + if cache_positions.ndim == 4: + cache_positions = cache_positions.reshape(B, -1) + elif cache_positions.ndim != 2: + raise ValueError( + "cache_positions must have shape (B, K) or (B, 1, 1, K), " + f"got {tuple(cache_positions.shape)}." + ) + + if key_valid_mask is not None: + if key_valid_mask.ndim == 4: + key_valid_mask = key_valid_mask.reshape(B, -1) + elif key_valid_mask.ndim != 2: + raise ValueError( + "key_valid_mask must have shape (B, K) or (B, 1, 1, K), " + f"got {tuple(key_valid_mask.shape)}." + ) + + q_grouped = Q.reshape(B, kv_heads, q_per_kv, q_len, head_dim) + k_grouped = K_cache.transpose(-1, -2).unsqueeze(2) + attn_weights = torch.matmul(q_grouped, k_grouped) / math.sqrt(head_dim) + + causal_mask = cache_positions[:, None, None, None, :] <= query_positions[ + :, None, None, :, None + ] + if key_valid_mask is not None: + causal_mask = causal_mask & key_valid_mask[:, None, None, None, :] + attn_weights = attn_weights.masked_fill(~causal_mask, -65504.0) + attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(Q.dtype) + + attn_output = torch.matmul(attn_weights, V_cache.unsqueeze(2)) + return attn_output.reshape(B, q_heads, q_len, head_dim) + + +def _qwen35_expanded_prefix_attention( + Q, + K_cache, + V_cache, + query_positions, + cache_positions, + key_valid_mask=None, +): + B, q_heads, q_len, head_dim = Q.shape + kv_heads = K_cache.shape[1] + cache_len = K_cache.shape[2] + + if q_heads != kv_heads: + kv_rep = q_heads // kv_heads + K_full = ( + K_cache.unsqueeze(2) + .expand(-1, -1, kv_rep, -1, -1) + .reshape(B, q_heads, cache_len, head_dim) + ) + V_full = ( + V_cache.unsqueeze(2) + .expand(-1, -1, kv_rep, -1, -1) + .reshape(B, q_heads, cache_len, head_dim) + ) + else: + K_full = K_cache + V_full = V_cache + + attn_weights = torch.matmul(Q, K_full.transpose(-1, -2)) / math.sqrt(head_dim) + causal_mask = cache_positions <= query_positions[:, None, :, None] + if key_valid_mask is not None: + causal_mask = causal_mask & key_valid_mask + attn_weights = attn_weights.masked_fill(~causal_mask, -65504.0) + attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(Q.dtype) + return torch.matmul(attn_weights, V_full) + + +def _qwen36_prefix_attention_impl() -> str: + raw = os.environ.get("QWEN36_PREFIX_ATTENTION_IMPL", "grouped").strip().lower() + aliases = { + "grouped": "grouped", + "current": "grouped", + "expanded": "expanded", + "legacy": "expanded", + "legacy_expanded": "expanded", + } + if raw not in aliases: + raise ValueError( + "QWEN36_PREFIX_ATTENTION_IMPL must be grouped/current or " + f"expanded/legacy, got {raw!r}" + ) + return aliases[raw] + + +def _resolve_deltanet_autocp_lnc(num_chunks: int) -> int: + lnc = _infer_neuron_lnc() + raw_lnc = os.environ.get("QWEN36_DELTANET_AUTOCP_LNC") + if raw_lnc is None: + launch_lnc = 2 if lnc >= 2 and num_chunks % 2 == 0 else 1 + else: + launch_lnc = max(1, int(raw_lnc)) + if launch_lnc > lnc: + raise ValueError( + f"QWEN36_DELTANET_AUTOCP_LNC={launch_lnc} requires " + f"NEURON_CC_FLAGS --lnc >= {launch_lnc}; inferred lnc={lnc}" + ) + if launch_lnc not in (1, 2): + raise ValueError( + f"QWEN36_DELTANET_AUTOCP_LNC must be 1 or 2, got {launch_lnc}" + ) + if num_chunks % launch_lnc != 0: + raise ValueError( + "QWEN36_DELTANET_AUTOCP_CTE requires the number of 128-token " + f"chunks to be divisible by launch LNC; chunks={num_chunks}, " + f"launch_lnc={launch_lnc}" + ) + return launch_lnc + + +def _deltanet_autocp_affine_launch_spec(num_chunks: int, launch_lnc: int): + """Return a SPMD affine launch grid, falling back to legacy LNC split. + + Bare ``kernel[2]`` launches only two logical cores. For AutoCP affine + generation we need one independent program per 128-token chunk, sharded + across those logical cores. NKI represents that as a SPMD grid dimension + with an attached NC distribution. + """ + if os.environ.get("QWEN36_DELTANET_AUTOCP_SPMD_AFFINE", "1") == "0": + return launch_lnc + import nki.language as _nl # Imported lazily so CPU-only unit stubs still load. + + if not hasattr(_nl, "spmd_dim") or not hasattr(_nl, "nc"): + return launch_lnc + + if launch_lnc == 2: + return (_nl.spmd_dim(num_chunks, _nl.nc(2)), 1) + return (num_chunks, 1) + + +def _resolve_deltanet_autocp_cp_chunks(num_chunks: int) -> int: + cp_chunks = max(1, int(os.environ.get("QWEN36_DELTANET_AUTOCP_CP_CHUNKS", "4"))) + if num_chunks % cp_chunks != 0: + raise ValueError( + "QWEN36_DELTANET_COMPACT_AUTOCP_CTE requires the number of " + "128-token chunks to be divisible by QWEN36_DELTANET_AUTOCP_CP_CHUNKS; " + f"chunks={num_chunks}, cp_chunks={cp_chunks}" + ) + return cp_chunks + +from neuronx_distributed_inference.models.config import ( + InferenceConfig, + MoENeuronConfig, + NeuronConfig, +) +try: + from neuronx_distributed_inference.modules.moe_v2 import initialize_moe_module + HAS_MOE_V2 = True +except ImportError: + HAS_MOE_V2 = False + + +def _patch_nxd_shard_hidden_kernel(): + """Monkey-patch NxD's `_call_shard_hidden_kernel` to use the nkilib + blockwise-matmul-shard-on-H NKI kernel present in the shipped `nkilib` + package (neuronx-cc 2.26.6360). Upstream NxD gates the fastest MoE + blockwise path on a + `neuronxcc.nki._private.blockwise_mm` module which is absent from the + DLAMI; falling through, it raises NotImplementedError. But the same + kernel exists at `nkilib.experimental.moe.forward.bwmm_shard_on_H. + blockwise_mm_baseline_shard_hidden` — we wire that up here so the + default LNC=2 forward path becomes available. + """ + try: + from nkilib.experimental.moe.forward.bwmm_shard_on_H import ( + blockwise_mm_baseline_shard_hidden as _nki_shard_H, + ) + except ImportError: + return # nothing to do; NxD's fallback handling stays in effect + + from neuronx_distributed.modules.moe import blockwise as _nxd_blockwise + + if getattr(_nxd_blockwise, "_qwen35_shard_hidden_patched", False): + return + + def _call_shard_hidden_kernel_patched(args): + """Drop-in replacement matching the NxD `_call_shard_hidden_kernel` + signature but using the nkilib forward-only kernel. + + The nkilib kernel returns (output, gate_up_activations_T, down_activations) + — no `output=` kwarg — so we match that layout. + """ + result = _nki_shard_H[2]( + hidden_states=args.hidden_states, + expert_affinities_masked=args.expert_affinities_masked, + gate_up_proj_weight=args.gate_up_proj_weight, + down_proj_weight=args.down_proj_weight, + block_size=args.block_size, + token_position_to_id=args.token_position_to_id.to(dtype=torch.int32), + block_to_expert=args.block_to_expert.to(dtype=torch.int32), + gate_up_activations_T=args.gate_up_activations_T, + down_activations=args.down_activations, + skip_dma=args.skip_dma, + is_tensor_update_accumulating=args.is_tensor_update_accumulating, + expert_affinities_scaling_mode=args.expert_affinities_scaling_mode, + ) + if isinstance(result, tuple) and len(result) == 3: + return result + # Some kernel builds return only the output tensor; adapt. + return result, args.gate_up_activations_T, args.down_activations + + _nxd_blockwise._call_shard_hidden_kernel = _call_shard_hidden_kernel_patched + _nxd_blockwise._qwen35_shard_hidden_patched = True + + +_patch_nxd_shard_hidden_kernel() +from neuronx_distributed_inference.models.llama.modeling_llama import NeuronLlamaMLP +from neuronx_distributed_inference.models.model_wrapper import ( + CONTEXT_ENCODING_MODEL_TAG, + TOKEN_GENERATION_MODEL_TAG, + DecoderModelInstance, + ModelWrapper, +) +from neuronx_distributed_inference.modules.attention.attention_base import NeuronAttentionBase +from neuronx_distributed_inference.modules.attention.utils import ( + RotaryEmbedding, + move_heads_front, + transpose_parallel_linear_layer, +) +try: + from neuronx_distributed_inference.modules.attention.utils import ( + preprocess_quantized_linear_layer, + ) +except (ImportError, AttributeError): + def preprocess_quantized_linear_layer(layer): + return layer + +from neuronx_distributed_inference.modules.kvcache.block_kv_cache_manager import ( + BlockKVCacheManager, +) +from neuronx_distributed_inference.modules.kvcache.kv_cache_manager import KVCacheManager +from neuronx_distributed_inference.models.layer_boundary_marker import ( + ModuleMarkerEndWrapper, + ModuleMarkerStartWrapper, +) + +logger = logging.getLogger(__name__) + +try: + from neuronxcc.nki._pre_prod_kernels import ( + NormType as _QKVNormType, + QKVOutputLayout as _QKVOutputLayout, + QuantizationType as _QKVQuantizationType, + ) + from neuronxcc.nki._pre_prod_kernels.qkv_tkg_impl import ( + nki_qkv_projection_tkg_impl as _qkv_tkg_nki_kernel, + ) +except ImportError: + _QKVNormType = None + _QKVOutputLayout = None + _QKVQuantizationType = None + _qkv_tkg_nki_kernel = None + +try: + _flash_fwd_call = nki_jit()(attention_isa_kernel) +except TypeError: + from torch_neuronx.xla_impl.ops import nki_jit as _torch_xla_nki_jit + + _flash_fwd_call = _torch_xla_nki_jit()(attention_isa_kernel) + +# Option B: Direct nkilib flash attention for head_dim > 128 +USE_NKILIB_KERNEL = os.environ.get("USE_NKILIB_KERNEL", "0") == "1" + +_nkilib_flash_attn = None +if USE_NKILIB_KERNEL: + try: + import neuronxcc.nki as _nki + from neuronx_distributed_inference.modules.attention.attention_base import ( + peel_decorations as _peel_decorations, + get_platform_target as _get_platform_target, + ) + from neuronxcc.nki.compiler import ( + skip_middle_end_transformations as _skip_middle_end, + enable_stack_allocator as _enable_stack_allocator, + ) + + import importlib + + _fork_path = "/home/ubuntu/nki-library-fork/nkilib_src" + if os.path.isdir(_fork_path) and _fork_path not in sys.path: + sys.path.insert(0, _fork_path) + _to_remove = [k for k in sys.modules if k.startswith("nkilib")] + for k in _to_remove: + del sys.modules[k] + import nki.language as _stub_nl + import neuronxcc.nki.language as _real_nl + + for _attr in [ + "NKIObject", + "float8_e4m3fn", + "float8_e4m3fn_x4", + "float8_e5m2_x4", + "float4_e2m1fn_x4", + ]: + if not hasattr(_real_nl, _attr) and hasattr(_stub_nl, _attr): + setattr(_real_nl, _attr, getattr(_stub_nl, _attr)) + from nkilib.core.attention.attention_cte import ( + attention_cte as _attention_cte_raw, + _MAX_HEAD_DIM, + ) + + assert _MAX_HEAD_DIM == 256, ( + f"nkilib fork has _MAX_HEAD_DIM={_MAX_HEAD_DIM}, expected 256. " + f"System nkilib may have been loaded instead of fork." + ) + logger.info( + f"Loaded nkilib attention_cte from fork (_MAX_HEAD_DIM={_MAX_HEAD_DIM})" + ) + + _raw_fn = _peel_decorations(_attention_cte_raw) + os.environ.setdefault("NEURON_PLATFORM_TARGET_OVERRIDE", _get_platform_target()) + _nkilib_flash_attn = _nki.jit( + _raw_fn, + show_compiler_tb=True, + debug_kernel=True, + ) + _nkilib_flash_attn = _skip_middle_end(_nkilib_flash_attn) + _nkilib_flash_attn = _enable_stack_allocator( + _nkilib_flash_attn, log_level=logging.INFO + ) + logger.info("Option B: nkilib flash attention loaded for head_dim > 128") + except Exception as e: + logger.warning(f"Option B: Failed to load nkilib flash attention: {e}") + import traceback as _tb + + _tb.print_exc() + _nkilib_flash_attn = None + +# Option A: Detect if patch_attn_kernel was imported +NKILIB_PATCH_ACTIVE = False +try: + from importlib import import_module as _import_module + + _attn_mod = _import_module("neuronxcc.nki._pre_prod_kernels.attn_fwd") + if hasattr(_attn_mod, "_original_attention_nki_kernel_adapter"): + NKILIB_PATCH_ACTIVE = True + logger.info("Option A detected: _pre_prod_kernels patched with nkilib kernel") +except Exception: + pass + + +# ============================================================ +# Newton-Raphson Refined RMSNorm +# ============================================================ +USE_NEWTON_RMSNORM = os.environ.get("USE_NEWTON_RMSNORM") == "1" +USE_PYTHON_RMSNORM = os.environ.get("USE_PYTHON_RMSNORM") == "1" + + +class NewtonRMSNorm(nn.Module): + """RMSNorm with Newton-Raphson refined rsqrt for improved numerical accuracy.""" + + def __init__(self, hidden_size=None, eps=1e-6): + super().__init__() + self.weight = None + if hidden_size is not None: + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.hidden_size = hidden_size + self.variance_epsilon = eps + + def forward(self, hidden_states): + original_dtype = hidden_states.dtype + x = hidden_states.to(torch.float32) + variance = x.pow(2).mean(-1, keepdim=True) + y = torch.rsqrt(variance + self.variance_epsilon) + y = y * (3.0 - (variance + self.variance_epsilon) * y * y) * 0.5 + result = x * y + if self.weight is not None: + result = result * self.weight.float() + return result.to(original_dtype) + + +def get_rmsnorm_cls(): + if cpu_mode() or USE_PYTHON_RMSNORM: + return Qwen3MoeRMSNorm + return NewtonRMSNorm if USE_NEWTON_RMSNORM else CustomRMSNorm + + +def l2norm(x, dim=-1, eps=1e-6): + return F.normalize(x, p=2, dim=dim, eps=eps) + + +class GDNAPCReusePlan(NamedTuple): + """Exact hybrid-APC reuse plan for attention KV plus GDN checkpoints.""" + + attention_hit_len: int + recurrent_hit_len: int + conv_hit_len: int + reusable_prefix_len: int + restore_checkpoint_prefix_len: int + residual_replay_len: int + suffix_len: int + + +def _non_negative_len(name: str, value: int) -> int: + value = int(value) + if value < 0: + raise ValueError(f"{name} must be non-negative, got {value}") + return value + + +def _normalize_hybrid_cache_dtype(name: str, value, default: str) -> str: + if value is None: + value = default + if isinstance(value, torch.dtype): + if value == torch.float32: + return "float32" + if value == torch.bfloat16: + return "bfloat16" + normalized = str(value).lower() + aliases = { + "fp32": "float32", + "float32": "float32", + "torch.float32": "float32", + "bf16": "bfloat16", + "bfloat16": "bfloat16", + "torch.bfloat16": "bfloat16", + } + if normalized not in aliases: + raise ValueError( + f"{name} must be one of fp32/float32 or bf16/bfloat16, got {value}" + ) + return aliases[normalized] + + +def _torch_dtype_from_hybrid_cache_dtype(value: str) -> torch.dtype: + value = _normalize_hybrid_cache_dtype("hybrid cache dtype", value, "bfloat16") + if value == "float32": + return torch.float32 + if value == "bfloat16": + return torch.bfloat16 + raise AssertionError(f"unexpected hybrid cache dtype {value}") + + +def plan_gdn_apc_reuse( + *, + attention_hit_len: int, + recurrent_hit_len: int, + conv_hit_len: int, + request_prefix_len: int, + gdn_checkpoint_interval: int, +) -> GDNAPCReusePlan: + """Plan exact prefix reuse for Qwen hybrid APC. + + Attention KV can be reused up to the vLLM APC hit, but DeltaNet can only + resume exactly from a boundary with both recurrent and conv checkpoint + state. When the attention hit is inside a GDN interval, restore the nearest + earlier checkpoint and replay the residual tokens before running the suffix. + """ + attention_hit_len = _non_negative_len("attention_hit_len", attention_hit_len) + recurrent_hit_len = _non_negative_len("recurrent_hit_len", recurrent_hit_len) + conv_hit_len = _non_negative_len("conv_hit_len", conv_hit_len) + request_prefix_len = _non_negative_len("request_prefix_len", request_prefix_len) + gdn_checkpoint_interval = int(gdn_checkpoint_interval) + if gdn_checkpoint_interval <= 0: + raise ValueError( + f"gdn_checkpoint_interval must be positive, got {gdn_checkpoint_interval}" + ) + + reusable_prefix_len = min( + attention_hit_len, + recurrent_hit_len, + conv_hit_len, + request_prefix_len, + ) + restore_checkpoint_prefix_len = ( + reusable_prefix_len // gdn_checkpoint_interval + ) * gdn_checkpoint_interval + residual_replay_len = reusable_prefix_len - restore_checkpoint_prefix_len + suffix_len = request_prefix_len - reusable_prefix_len + + return GDNAPCReusePlan( + attention_hit_len=attention_hit_len, + recurrent_hit_len=recurrent_hit_len, + conv_hit_len=conv_hit_len, + reusable_prefix_len=reusable_prefix_len, + restore_checkpoint_prefix_len=restore_checkpoint_prefix_len, + residual_replay_len=residual_replay_len, + suffix_len=suffix_len, + ) + + +# ============================================================ +# Gated DeltaNet Module (Linear Recurrent Attention) +# ============================================================ + + +class NeuronGatedDeltaNet(nn.Module): + """ + Gated DeltaNet linear attention for Neuron. + + Replaces standard attention for 48 of 64 layers in Qwen3.5/3.6-27B. + Uses a chunk-based linear recurrence instead of KV cache. + + HF weight layout (27B dense -- scaled dimensions): + - in_proj_qkv.weight: (key_dim*2 + value_dim, hidden_size) = (10240, 5120) + - in_proj_z.weight: (value_dim, hidden_size) = (6144, 5120) + - in_proj_a.weight: (num_v_heads, hidden_size) = (48, 5120) + - in_proj_b.weight: (num_v_heads, hidden_size) = (48, 5120) + - conv1d.weight: (conv_dim, 1, conv_kernel_size) = (10240, 1, 4) + - A_log: (num_v_heads,) = (48,) + - dt_bias: (num_v_heads,) = (48,) + - norm.weight: (head_v_dim,) = (128,) + - out_proj.weight: (hidden_size, value_dim) = (5120, 6144) + """ + + def __init__(self, config, layer_idx: int): + super().__init__() + tc = config + + self.hidden_size = tc.hidden_size # 5120 + self.tp_degree = tc.neuron_config.tp_degree + self.global_num_v_heads = tc.linear_num_value_heads # 48 + self.global_num_k_heads = tc.linear_num_key_heads # 16 + self.head_k_dim = tc.linear_key_head_dim # 128 + self.head_v_dim = tc.linear_value_head_dim # 128 + if self.global_num_v_heads % self.tp_degree != 0: + raise ValueError( + f"linear_num_value_heads={self.global_num_v_heads} must be divisible " + f"by tp_degree={self.tp_degree}" + ) + if self.global_num_k_heads % self.tp_degree != 0: + raise ValueError( + f"linear_num_key_heads={self.global_num_k_heads} must be divisible " + f"by tp_degree={self.tp_degree}" + ) + self.num_v_heads = self.global_num_v_heads // self.tp_degree + self.num_k_heads = self.global_num_k_heads // self.tp_degree + self.global_key_dim = self.head_k_dim * self.global_num_k_heads # 2048 + self.global_value_dim = self.head_v_dim * self.global_num_v_heads # 6144 + self.key_dim = self.head_k_dim * self.num_k_heads # 512 at TP=4 + self.value_dim = self.head_v_dim * self.num_v_heads # 1536 at TP=4 + self.conv_kernel_size = tc.linear_conv_kernel_dim # 4 + self.layer_idx = layer_idx + self.rms_norm_eps = tc.rms_norm_eps + self.use_hybrid_cache_manager = getattr(tc, "use_hybrid_cache_manager", False) + self.use_hybrid_apc_manager = getattr(tc, "use_hybrid_apc_manager", False) + self.use_qwen_hybrid_chunked_prefill = getattr( + tc, "use_qwen_hybrid_chunked_prefill", False + ) + self.use_qwen_hybrid_chunked_prefill_nki = getattr( + tc, "use_qwen_hybrid_chunked_prefill_nki", False + ) + self.use_qwen_deltanet_decode_nki = getattr( + tc, "use_qwen_deltanet_decode_nki", False + ) + self.use_cold_zero_conv_fast_path = getattr( + tc, "use_cold_zero_conv_fast_path", False + ) + + # KV cache dummy shape info + self.head_dim = tc.head_dim # 256 + tp_degree = tc.neuron_config.tp_degree + raw_kv_heads = tc.num_key_value_heads + if raw_kv_heads < tp_degree: + replicated_kv_heads = tp_degree + else: + replicated_kv_heads = raw_kv_heads + self.kv_heads_per_rank = replicated_kv_heads // tp_degree + + # Conv1d on concatenated QKV (NOT Z). Store the depthwise kernel in a + # ColumnParallelLinear parameter container so NxD's checkpoint sharder + # can split it by output channel. Forward still uses it as Conv1d + # weight after unsqueezing the singleton input-channel dimension. + self.global_conv_dim = self.global_key_dim * 2 + self.global_value_dim # 10240 + self.conv_dim = self.key_dim * 2 + self.value_dim # 2560 at TP=4 + self.conv1d_weight = ColumnParallelLinear( + self.conv_kernel_size, + self.global_conv_dim, + bias=False, + gather_output=False, + ) + + # Input/output projections are the large DeltaNet tensors. Shard them + # with tensor parallelism; convert_qwen35_hf_to_neuron_state_dict() + # reorders in_proj_qkv into per-rank [Q_local | K_local | V_local] + # blocks before NxD slices the output dimension. + self.in_proj_qkv = ColumnParallelLinear( + self.hidden_size, + self.global_key_dim * 2 + self.global_value_dim, + bias=False, + gather_output=False, + ) + self.in_proj_z = ColumnParallelLinear( + self.hidden_size, + self.global_value_dim, + bias=False, + gather_output=False, + ) + self.in_proj_b = ColumnParallelLinear( + self.hidden_size, + self.global_num_v_heads, + bias=False, + gather_output=False, + ) + self.in_proj_a = ColumnParallelLinear( + self.hidden_size, + self.global_num_v_heads, + bias=False, + gather_output=False, + ) + + # Same parameter-container pattern for per-value-head decay vectors. + # These are used as vectors in forward but sharded by output dim during + # checkpoint conversion/loading. + self.dt_bias_weight = ColumnParallelLinear( + 1, + self.global_num_v_heads, + bias=False, + gather_output=False, + ) + self.A_log_weight = ColumnParallelLinear( + 1, + self.global_num_v_heads, + bias=False, + gather_output=False, + ) + + # Output norm and projection + self.norm = Qwen3MoeRMSNorm(self.head_v_dim, eps=self.rms_norm_eps) + self.out_proj = RowParallelLinear( + self.global_value_dim, + self.hidden_size, + bias=False, + input_is_parallel=True, + ) + + # State buffers for CTE -> TKG carry-over + alloc_batch_size = getattr(config.neuron_config, "max_batch_size", 1) + self._phase_batch_size = getattr(config.neuron_config, "batch_size", 1) + recurrent_buffer_dtype = ( + _torch_dtype_from_hybrid_cache_dtype(config.hybrid_recurrent_cache_dtype) + if self.use_hybrid_apc_manager + else config.neuron_config.torch_dtype + ) + conv_buffer_dtype = ( + _torch_dtype_from_hybrid_cache_dtype(config.hybrid_conv_cache_dtype) + if self.use_hybrid_apc_manager + else config.neuron_config.torch_dtype + ) + self.recurrent_state_buffer = nn.Parameter( + torch.zeros( + alloc_batch_size, + self.num_v_heads, + self.head_k_dim, + self.head_v_dim, + dtype=recurrent_buffer_dtype, + ), + requires_grad=False, + ) + self.conv_state_buffer = nn.Parameter( + torch.zeros( + alloc_batch_size, + self.conv_dim, + self.conv_kernel_size - 1, + dtype=conv_buffer_dtype, + ), + requires_grad=False, + ) + + def _conv1d_weight(self): + return self.conv1d_weight.weight.unsqueeze(1) + + def _dt_bias(self): + return self.dt_bias_weight.weight.squeeze(1) + + def _A_log(self): + return self.A_log_weight.weight.squeeze(1) + + def _recurrent_step(self, query, key, value, g, beta, recurrent_state): + """Single-step recurrent update for token generation.""" + query = l2norm(query, dim=-1) + key = l2norm(key, dim=-1) + scale = 1.0 / (query.shape[-1] ** 0.5) + query = query * scale + + q_t = query[:, :, 0] + k_t = key[:, :, 0] + v_t = value[:, :, 0] + g_t = g[:, :, 0].exp().unsqueeze(-1).unsqueeze(-1) + beta_t = beta[:, :, 0].unsqueeze(-1) + + new_state = recurrent_state * g_t + kv_mem = (new_state * k_t.unsqueeze(-1)).sum(dim=-2) + delta = (v_t - kv_mem) * beta_t + new_state = new_state + k_t.unsqueeze(-1) * delta.unsqueeze(-2) + output = (new_state * q_t.unsqueeze(-1)).sum(dim=-2) + + return output.unsqueeze(2), new_state + + def _nki_recurrent_step(self, query, key, value, g, beta, recurrent_state): + """Single-step recurrent update using the stateful NKI decode kernel.""" + query = l2norm(query, dim=-1) + key = l2norm(key, dim=-1) + B, H, S, k_dim = query.shape + v_dim = value.shape[-1] + scale = 1.0 / (k_dim**0.5) + query = query * scale + + BH = B * H + query_flat = query.reshape(BH, S, k_dim)[:, 0, :].contiguous() + key_flat = key.reshape(BH, S, k_dim)[:, 0, :].contiguous() + value_flat = value.reshape(BH, S, v_dim)[:, 0, :].contiguous() + g_flat = g.reshape(BH, S)[:, 0:1].contiguous() + beta_flat = beta.reshape(BH, S)[:, 0:1].contiguous() + state_flat = recurrent_state.reshape(BH * k_dim, v_dim).contiguous() + + output_flat, state_flat_out = _deltanet_nki_step_batched( + query_flat, + key_flat, + value_flat, + g_flat, + beta_flat, + state_flat, + ) + + output = output_flat.reshape(B, H, S, v_dim) + new_state = state_flat_out.reshape(B, H, k_dim, v_dim) + + return output, new_state + + def _nki_recurrent_forward(self, query, key, value, g, beta): + """Full-sequence recurrent forward using NKI kernel for context encoding.""" + query = l2norm(query, dim=-1) + key = l2norm(key, dim=-1) + B, H, S, k_dim = query.shape + v_dim = value.shape[-1] + scale = 1.0 / (k_dim**0.5) + query = query * scale + + BH = B * H + query_flat = query.reshape(BH, S, k_dim).contiguous() + key_flat = key.reshape(BH, S, k_dim).contiguous() + value_flat = value.reshape(BH, S, v_dim).contiguous() + + g_flat = g.reshape(BH, S).unsqueeze(-1).expand(-1, -1, v_dim).contiguous() + beta_flat = beta.reshape(BH, S).unsqueeze(-1).expand(-1, -1, v_dim).contiguous() + + outputs = [] + states = [] + for bh in range(BH): + out_bh, state_bh = _deltanet_nki_kernel_state( + query_flat[bh], + key_flat[bh], + value_flat[bh], + g_flat[bh], + beta_flat[bh], + ) + outputs.append(out_bh) + states.append(state_bh) + + output = torch.stack(outputs, dim=0) + output = output.reshape(B, H, S, v_dim) + + final_state = torch.stack(states, dim=0) + final_state = final_state.reshape(B, H, k_dim, v_dim) + + return output, final_state + + def _nki_chunked_forward( + self, query, key, value, g, beta, output_final_state=False, initial_state=None + ): + """Chunked NKI kernel forward for context encoding (prefill).""" + chunk_size = 128 + + query = l2norm(query, dim=-1) + key = l2norm(key, dim=-1) + B, H, S, k_dim = query.shape + v_dim = value.shape[-1] + scale = 1.0 / (k_dim**0.5) + query = query * scale + + pad_size = (chunk_size - S % chunk_size) % chunk_size + if pad_size > 0: + query = F.pad(query, (0, 0, 0, pad_size)) + key = F.pad(key, (0, 0, 0, pad_size)) + value = F.pad(value, (0, 0, 0, pad_size)) + beta = F.pad(beta, (0, pad_size)) + g = F.pad(g, (0, pad_size)) + total_seq_len = S + pad_size + + num_chunks = total_seq_len // chunk_size + g_reshaped = g.reshape(B, H, num_chunks, chunk_size) + g_cs = g_reshaped.cumsum(dim=-1) + g_last_per_chunk = g_cs[:, :, :, -1:] + g_last_expanded = g_last_per_chunk.expand(-1, -1, -1, chunk_size) + + query_chunks = query.reshape(B, H, num_chunks, chunk_size, k_dim) + key_chunks = key.reshape(B, H, num_chunks, chunk_size, k_dim) + value_chunks = value.reshape(B, H, num_chunks, chunk_size, v_dim) + + beta_chunks = ( + beta.reshape(B, H, num_chunks, chunk_size) + .unsqueeze(-1) + .expand(-1, -1, -1, -1, v_dim) + ) + gc_chunks = g_cs.unsqueeze(-1).expand(-1, -1, -1, -1, v_dim) + gl_chunks = g_last_expanded.unsqueeze(-1).expand(-1, -1, -1, -1, v_dim) + + BH = B * H + query_chunks = query_chunks.reshape( + BH, num_chunks, chunk_size, k_dim + ).contiguous() + key_chunks = key_chunks.reshape(BH, num_chunks, chunk_size, k_dim).contiguous() + value_chunks = value_chunks.reshape( + BH, num_chunks, chunk_size, v_dim + ).contiguous() + beta_chunks = beta_chunks.reshape( + BH, num_chunks, chunk_size, v_dim + ).contiguous() + gc_chunks = gc_chunks.reshape(BH, num_chunks, chunk_size, v_dim).contiguous() + gl_chunks = gl_chunks.reshape(BH, num_chunks, chunk_size, v_dim).contiguous() + + device = query.device + lower_mask = torch.tril( + torch.ones(chunk_size, chunk_size, dtype=torch.float32, device=device), + diagonal=-1, + ) + identity_mat = torch.eye(chunk_size, dtype=torch.float32, device=device) + lower_mask_diag = torch.tril( + torch.ones(chunk_size, chunk_size, dtype=torch.float32, device=device), + diagonal=0, + ) + + initial_state_flat = None + if initial_state is not None: + initial_state_flat = initial_state.reshape(BH, k_dim, v_dim).float().contiguous() + + all_outputs = [] + all_states = [] + for bh in range(BH): + if initial_state_flat is None: + state = torch.zeros(k_dim, v_dim, dtype=torch.float32, device=device) + else: + state = initial_state_flat[bh] + + head_chunks = [] + for c_idx in range(num_chunks): + q_chunk = query_chunks[bh, c_idx].contiguous() + k_chunk = key_chunks[bh, c_idx].contiguous() + v_chunk = value_chunks[bh, c_idx].contiguous() + beta_chunk = beta_chunks[bh, c_idx].contiguous() + gc_chunk = gc_chunks[bh, c_idx].contiguous() + gl_chunk = gl_chunks[bh, c_idx].contiguous() + + out_chunk, state = _deltanet_nki_chunk_step( + q_chunk, + k_chunk, + v_chunk, + beta_chunk, + gc_chunk, + gl_chunk, + state, + lower_mask, + identity_mat, + lower_mask_diag, + ) + head_chunks.append(out_chunk) + + head_output = torch.cat(head_chunks, dim=0) + all_outputs.append(head_output) + all_states.append(state) + + output = torch.stack(all_outputs, dim=0) + output = output.reshape(B, H, total_seq_len, v_dim) + output = output[:, :, :S] + + if output_final_state: + final_state = torch.stack(all_states, dim=0) + last_recurrent_state = final_state.reshape(B, H, k_dim, v_dim) + else: + last_recurrent_state = None + + return output, last_recurrent_state + + def _fused_chunked_forward( + self, + query, + key, + value, + g, + beta, + output_final_state=False, + initial_state=None, + _segment_disabled=False, + ): + """Fused single-kernel chunked forward for CTE — SSD-style. + + Processes all chunks in a single NKI kernel call per (B,H) pair. + State persists in SBUF across chunks (no HBM round-trips). + Cumsum of g computed in-kernel via tensor_tensor_scan. + + This is the optimized version of _nki_chunked_forward with: + 1. Single kernel call per (B,H) instead of B*H*num_chunks + 2. State in SBUF across all chunks (biggest perf win) + 3. In-kernel cumsum (avoids PyTorch cumsum overhead) + 4. tensor_scalar for broadcasts (no explicit loops) + + initial_state is the restored GDN recurrent checkpoint for warm or + partial-prefix suffix prefill. Cold prefill passes zeros. + """ + chunk_size = int(os.environ.get("QWEN36_DELTANET_CHUNK_SIZE", "128")) + if chunk_size not in (64, 128): + raise ValueError( + "QWEN36_DELTANET_CHUNK_SIZE must be 64 or 128 for fused CTE; " + f"got {chunk_size}" + ) + + cte_impl = os.environ.get("QWEN36_DELTANET_CTE_IMPL", "current").lower() + if cte_impl in ("legacy", "legacy_direct", "direct"): + use_legacy_direct_cte = True + elif cte_impl in ("current", "optimized"): + use_legacy_direct_cte = False + else: + raise ValueError( + "QWEN36_DELTANET_CTE_IMPL must be current or legacy_direct; " + f"got {cte_impl!r}" + ) + if use_legacy_direct_cte and chunk_size != 128: + raise ValueError( + "QWEN36_DELTANET_CTE_IMPL=legacy_direct requires " + f"QWEN36_DELTANET_CHUNK_SIZE=128; got {chunk_size}" + ) + + B, H, S, k_dim = query.shape + v_dim = value.shape[-1] + + # Pad sequence to multiple of chunk_size + pad_size = (chunk_size - S % chunk_size) % chunk_size + if pad_size > 0: + query = F.pad(query, (0, 0, 0, pad_size)) + key = F.pad(key, (0, 0, 0, pad_size)) + value = F.pad(value, (0, 0, 0, pad_size)) + beta = F.pad(beta, (0, pad_size)) + g = F.pad(g, (0, pad_size)) + total_seq_len = S + pad_size + + segment_tokens = int( + os.environ.get("QWEN36_DELTANET_FUSED_SEGMENT_TOKENS", "0") or "0" + ) + if ( + not _segment_disabled + and segment_tokens > 0 + and total_seq_len > segment_tokens + ): + if segment_tokens < chunk_size or segment_tokens % chunk_size != 0: + raise ValueError( + "QWEN36_DELTANET_FUSED_SEGMENT_TOKENS must be a positive " + "multiple of QWEN36_DELTANET_CHUNK_SIZE; " + f"got segment_tokens={segment_tokens}, chunk_size={chunk_size}" + ) + segment_outputs = [] + state = initial_state + for start in range(0, total_seq_len, segment_tokens): + end = min(start + segment_tokens, total_seq_len) + segment_output, state = self._fused_chunked_forward( + query[:, :, start:end, :], + key[:, :, start:end, :], + value[:, :, start:end, :], + g[:, :, start:end], + beta[:, :, start:end], + output_final_state=True, + initial_state=state, + _segment_disabled=True, + ) + segment_outputs.append(segment_output) + output = torch.cat(segment_outputs, dim=2)[:, :, :S, :] + return output, state if output_final_state else None + + if use_legacy_direct_cte: + query = l2norm(query, dim=-1) + key = l2norm(key, dim=-1) + query = query * (1.0 / (k_dim ** 0.5)) + + BH = B * H + # Flatten to (BH, S, dim). Grouped multihead launches are opt-in + # because isolated validation must pass before using them in artifacts. + query_flat = query.reshape(BH, total_seq_len, k_dim).contiguous() + key_flat = key.reshape(BH, total_seq_len, k_dim).contiguous() + value_flat = value.reshape(BH, total_seq_len, v_dim).contiguous() + + # g and beta: (BH, S) -> (BH, S, 1) for the kernel's (S, 1) input layout + g_flat = g.reshape(BH, total_seq_len).unsqueeze(-1).contiguous() + beta_flat = beta.reshape(BH, total_seq_len).unsqueeze(-1).contiguous() + if initial_state is None: + initial_state_flat = torch.zeros( + BH, k_dim, v_dim, dtype=torch.float32, device=query.device + ) + else: + initial_state_flat = initial_state.reshape(BH, k_dim, v_dim).float().contiguous() + + # Create constant mask tensors (shared across all B*H calls) + device = query.device + lower_mask = torch.tensor( + _make_lower_mask(), dtype=torch.float32, device=device + ) + identity_mat = torch.tensor( + _make_identity(), dtype=torch.float32, device=device + ) + lower_mask_diag = torch.tensor( + _make_lower_mask_diag(), dtype=torch.float32, device=device + ) + + use_multihead_cte = ( + not use_legacy_direct_cte + and os.environ.get("QWEN36_DELTANET_MULTIHEAD_CTE", "1") != "0" + ) + if use_multihead_cte: + pair_outputs = [] + pair_states = [] + head_group_size = _resolve_deltanet_multihead_group_size(BH) + for bh_start in range(0, BH, head_group_size): + bh_end = min(bh_start + head_group_size, BH) + launch_heads = bh_end - bh_start + launch_spec = _deltanet_multihead_launch_spec(launch_heads) + out_pair, state_pair = _deltanet_fused_multihead_kernel[launch_spec]( + query_flat[bh_start:bh_end], # (G, S, 128) + key_flat[bh_start:bh_end], # (G, S, 128) + value_flat[bh_start:bh_end], # (G, S, 128) + g_flat[bh_start:bh_end], # (G, S, 1) — RAW g, not cumsum + beta_flat[bh_start:bh_end], # (G, S, 1) — sigmoid(b) + initial_state_flat[bh_start:bh_end], + lower_mask, # (128, 128) + identity_mat, # (128, 128) + lower_mask_diag, # (128, 128) + ) + pair_outputs.append(out_pair) + pair_states.append(state_pair) + + output = torch.cat(pair_outputs, dim=0) + final_state = torch.cat(pair_states, dim=0) + else: + fused_singlehead_kernel = ( + _deltanet_fused_legacy_direct_kernel + if use_legacy_direct_cte + else _deltanet_fused_kernel + ) + all_outputs = [] + all_states = [] + for bh in range(BH): + out_bh, state_bh = fused_singlehead_kernel( + query_flat[bh], # (S, 128) + key_flat[bh], # (S, 128) + value_flat[bh], # (S, 128) + g_flat[bh], # (S, 1) — RAW g, not cumsum + beta_flat[bh], # (S, 1) — sigmoid(b) + initial_state_flat[bh], # (128, 128) recurrent checkpoint + lower_mask, # (128, 128) + identity_mat, # (128, 128) + lower_mask_diag, # (128, 128) + ) + all_outputs.append(out_bh) + all_states.append(state_bh) + + output = torch.stack(all_outputs, dim=0) + final_state = torch.stack(all_states, dim=0) + + output = output.reshape(B, H, total_seq_len, v_dim) + output = output[:, :, :S] + + if output_final_state: + last_recurrent_state = final_state.reshape(B, H, k_dim, v_dim) + else: + last_recurrent_state = None + + return output, last_recurrent_state + + def _compact_autocp_chunked_forward( + self, query, key, value, g, beta, output_final_state=False, initial_state=None + ): + """Compact AutoCP CTE probe: prefix segment state summaries, replay segments. + + Compared with ``_autocp_chunked_forward``, this avoids materializing + per-chunk output-affine tensors. It is intentionally opt-in because the + first version reuses the existing recurrent fused kernel for segment + replay; a later NKI replay kernel can collapse the segment loop. + """ + chunk_size = 128 + + B, H, S, k_dim = query.shape + v_dim = value.shape[-1] + if k_dim != 128 or v_dim != 128: + raise ValueError( + "QWEN36_DELTANET_COMPACT_AUTOCP_CTE requires 128-wide " + f"key/value heads; got k_dim={k_dim}, v_dim={v_dim}" + ) + + pad_size = (chunk_size - S % chunk_size) % chunk_size + if pad_size > 0: + query = F.pad(query, (0, 0, 0, pad_size)) + key = F.pad(key, (0, 0, 0, pad_size)) + value = F.pad(value, (0, 0, 0, pad_size)) + beta = F.pad(beta, (0, pad_size)) + g = F.pad(g, (0, pad_size)) + total_seq_len = S + pad_size + num_chunks = total_seq_len // chunk_size + if num_chunks <= 0: + raise ValueError("QWEN36_DELTANET_COMPACT_AUTOCP_CTE requires chunks") + cp_chunks = _resolve_deltanet_autocp_cp_chunks(num_chunks) + num_segments = num_chunks // cp_chunks + launch_lnc = _resolve_deltanet_autocp_lnc(num_segments) + summary_launch_spec = _deltanet_autocp_affine_launch_spec( + num_segments, + launch_lnc, + ) + + BH = B * H + query_flat = query.reshape(BH, total_seq_len, k_dim).contiguous() + key_flat = key.reshape(BH, total_seq_len, k_dim).contiguous() + value_flat = value.reshape(BH, total_seq_len, v_dim).contiguous() + g_flat = g.reshape(BH, total_seq_len).unsqueeze(-1).contiguous() + beta_flat = beta.reshape(BH, total_seq_len).unsqueeze(-1).contiguous() + if initial_state is None: + initial_state_flat = torch.zeros( + BH, k_dim, v_dim, dtype=torch.float32, device=query.device + ) + else: + initial_state_flat = initial_state.reshape(BH, k_dim, v_dim).float().contiguous() + + device = query.device + lower_mask = torch.tensor( + _make_lower_mask(), dtype=torch.float32, device=device + ) + identity_mat = torch.tensor( + _make_identity(), dtype=torch.float32, device=device + ) + lower_mask_diag = torch.tensor( + _make_lower_mask_diag(), dtype=torch.float32, device=device + ) + + segment_len = cp_chunks * chunk_size + all_outputs = [] + all_states = [] + for bh in range(BH): + segment_matrix, segment_bias = ( + _deltanet_autocp_state_summary_sequence[summary_launch_spec]( + key_flat[bh], + value_flat[bh], + g_flat[bh], + beta_flat[bh], + lower_mask, + identity_mat, + ) + ) + segment_states, final_state = _deltanet_autocp_state_prefix( + segment_matrix, + segment_bias, + initial_state_flat[bh], + ) + + q_segments = query_flat[bh].reshape(num_segments, segment_len, k_dim).contiguous() + k_segments = key_flat[bh].reshape(num_segments, segment_len, k_dim).contiguous() + v_segments = value_flat[bh].reshape(num_segments, segment_len, v_dim).contiguous() + g_segments = g_flat[bh].reshape(num_segments, segment_len, 1).contiguous() + beta_segments = beta_flat[bh].reshape(num_segments, segment_len, 1).contiguous() + + replay_group_size = _resolve_deltanet_multihead_group_size(num_segments) + replay_outputs = [] + for segment_start in range(0, num_segments, replay_group_size): + segment_end = min(segment_start + replay_group_size, num_segments) + launch_segments = segment_end - segment_start + replay_launch_spec = _deltanet_multihead_launch_spec(launch_segments) + out_group, _ = _deltanet_fused_multihead_kernel[replay_launch_spec]( + q_segments[segment_start:segment_end], + k_segments[segment_start:segment_end], + v_segments[segment_start:segment_end], + g_segments[segment_start:segment_end], + beta_segments[segment_start:segment_end], + segment_states[segment_start:segment_end], + lower_mask, + identity_mat, + lower_mask_diag, + ) + replay_outputs.append(out_group) + out_segments = torch.cat(replay_outputs, dim=0) + + all_outputs.append(out_segments.reshape(total_seq_len, v_dim)) + all_states.append(final_state) + + output = torch.stack(all_outputs, dim=0) + output = output.reshape(B, H, total_seq_len, v_dim) + output = output[:, :, :S] + + if output_final_state: + final_state = torch.stack(all_states, dim=0) + last_recurrent_state = final_state.reshape(B, H, k_dim, v_dim) + else: + last_recurrent_state = None + + return output, last_recurrent_state + + def _autocp_chunked_forward( + self, query, key, value, g, beta, output_final_state=False, initial_state=None + ): + """FlashQLA-style AutoCP CTE path for exact GDN prefill probes. + + This path decomposes each 128-token chunk into an affine state transform, + scans chunk states, then applies the per-chunk initial state to outputs. + It is gated by QWEN36_DELTANET_AUTOCP_CTE while we measure whether the + extra custom-call/HBM traffic beats the recurrent fused path. + """ + chunk_size = 128 + + B, H, S, k_dim = query.shape + v_dim = value.shape[-1] + if k_dim != 128 or v_dim != 128: + raise ValueError( + "QWEN36_DELTANET_AUTOCP_CTE requires 128-wide key/value heads; " + f"got k_dim={k_dim}, v_dim={v_dim}" + ) + + pad_size = (chunk_size - S % chunk_size) % chunk_size + if pad_size > 0: + query = F.pad(query, (0, 0, 0, pad_size)) + key = F.pad(key, (0, 0, 0, pad_size)) + value = F.pad(value, (0, 0, 0, pad_size)) + beta = F.pad(beta, (0, pad_size)) + g = F.pad(g, (0, pad_size)) + total_seq_len = S + pad_size + num_chunks = total_seq_len // chunk_size + if num_chunks <= 0: + raise ValueError("QWEN36_DELTANET_AUTOCP_CTE requires at least one chunk") + launch_lnc = _resolve_deltanet_autocp_lnc(num_chunks) + + BH = B * H + query_flat = query.reshape(BH, total_seq_len, k_dim).contiguous() + key_flat = key.reshape(BH, total_seq_len, k_dim).contiguous() + value_flat = value.reshape(BH, total_seq_len, v_dim).contiguous() + g_flat = g.reshape(BH, total_seq_len).unsqueeze(-1).contiguous() + beta_flat = beta.reshape(BH, total_seq_len).unsqueeze(-1).contiguous() + if initial_state is None: + initial_state_flat = torch.zeros( + BH, k_dim, v_dim, dtype=torch.float32, device=query.device + ) + else: + initial_state_flat = initial_state.reshape(BH, k_dim, v_dim).float().contiguous() + + device = query.device + lower_mask = torch.tensor( + _make_lower_mask(), dtype=torch.float32, device=device + ) + identity_mat = torch.tensor( + _make_identity(), dtype=torch.float32, device=device + ) + lower_mask_diag = torch.tensor( + _make_lower_mask_diag(), dtype=torch.float32, device=device + ) + affine_launch_spec = _deltanet_autocp_affine_launch_spec( + num_chunks, + launch_lnc, + ) + + all_outputs = [] + all_states = [] + for bh in range(BH): + output_base, output_state, state_matrix, state_bias = ( + _deltanet_autocp_affine_sequence[affine_launch_spec]( + query_flat[bh], + key_flat[bh], + value_flat[bh], + g_flat[bh], + beta_flat[bh], + lower_mask, + identity_mat, + lower_mask_diag, + ) + ) + if os.environ.get("QWEN36_DELTANET_AUTOCP_SPLIT_APPLY") == "1": + chunk_states, final_state = _deltanet_autocp_state_prefix( + state_matrix, + state_bias, + initial_state_flat[bh], + ) + out_bh = _deltanet_autocp_apply_output( + output_base, + output_state, + chunk_states, + ) + else: + out_bh, final_state = _deltanet_autocp_prefix_apply_output( + output_base, + output_state, + state_matrix, + state_bias, + initial_state_flat[bh], + ) + all_outputs.append(out_bh) + all_states.append(final_state) + + output = torch.stack(all_outputs, dim=0) + output = output.reshape(B, H, total_seq_len, v_dim) + output = output[:, :, :S] + + if output_final_state: + final_state = torch.stack(all_states, dim=0) + last_recurrent_state = final_state.reshape(B, H, k_dim, v_dim) + else: + last_recurrent_state = None + + return output, last_recurrent_state + + def _sequential_forward(self, query, key, value, g, beta, output_final_state=False): + """Sequential full-sequence gated delta rule for CTE. + + Uses the same per-step recurrence as _recurrent_step but loops over the + full sequence. Avoids the slice-assignment loop in _chunk_forward that + may compile incorrectly on Neuron/XLA. + """ + query = l2norm(query, dim=-1) + key = l2norm(key, dim=-1) + + B, H, S, k_dim = query.shape + v_dim = value.shape[-1] + scale = 1.0 / (k_dim**0.5) + query = query * scale + + state = query.new_zeros(B, H, k_dim, v_dim) + all_outputs = [] + for t in range(S): + q_t = query[:, :, t] # (B, H, K) + k_t = key[:, :, t] # (B, H, K) + v_t = value[:, :, t] # (B, H, V) + beta_t = beta[:, :, t].unsqueeze(-1) # (B, H, 1) + g_t = g[:, :, t].exp().unsqueeze(-1).unsqueeze(-1) # (B, H, 1, 1) + + # Gated delta rule + state = state * g_t + kv_mem = (state * k_t.unsqueeze(-1)).sum(dim=-2) # (B, H, V) + delta = (v_t - kv_mem) * beta_t # (B, H, V) + state = state + k_t.unsqueeze(-1) * delta.unsqueeze(-2) # (B, H, K, V) + + o_t = (state * q_t.unsqueeze(-1)).sum(dim=-2) # (B, H, V) + all_outputs.append(o_t.unsqueeze(2)) + + output = torch.cat(all_outputs, dim=2) # (B, H, S, V) + final_state = state if output_final_state else None + return output, final_state + + def _chunk_forward( + self, query, key, value, g, beta, output_final_state=False, initial_state=None + ): + """Chunk-based forward for context encoding (prefill).""" + chunk_size = 64 + + query = l2norm(query, dim=-1) + key = l2norm(key, dim=-1) + + B, H, S, k_dim = query.shape + v_dim = value.shape[-1] + scale = 1.0 / (k_dim**0.5) + query = query * scale + + pad_size = (chunk_size - S % chunk_size) % chunk_size + if pad_size > 0: + query = F.pad(query, (0, 0, 0, pad_size)) + key = F.pad(key, (0, 0, 0, pad_size)) + value = F.pad(value, (0, 0, 0, pad_size)) + beta = F.pad(beta, (0, pad_size)) + g = F.pad(g, (0, pad_size)) + total_seq_len = S + pad_size + + v_beta = value * beta.unsqueeze(-1) + k_beta = key * beta.unsqueeze(-1) + + num_chunks = total_seq_len // chunk_size + query = query.reshape(B, H, num_chunks, chunk_size, k_dim) + key = key.reshape(B, H, num_chunks, chunk_size, k_dim) + value = value.reshape(B, H, num_chunks, chunk_size, v_dim) + k_beta = k_beta.reshape(B, H, num_chunks, chunk_size, k_dim) + v_beta = v_beta.reshape(B, H, num_chunks, chunk_size, v_dim) + g = g.reshape(B, H, num_chunks, chunk_size) + + mask = torch.triu( + torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), + diagonal=0, + ) + + g = g.cumsum(dim=-1) + decay_mask = (g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().tril() + + attn = -((k_beta @ key.transpose(-1, -2)) * decay_mask).masked_fill(mask, 0) + for i in range(1, chunk_size): + row = attn[..., i, :i].clone() + sub = attn[..., :i, :i].clone() + attn[..., i, :i] = row + (row.unsqueeze(-1) * sub).sum(-2) + attn = attn + torch.eye(chunk_size, dtype=attn.dtype, device=attn.device) + + value = attn @ v_beta + k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1)) + + if initial_state is None: + last_recurrent_state = torch.zeros( + B, H, k_dim, v_dim, dtype=query.dtype, device=query.device + ) + else: + last_recurrent_state = initial_state.to(dtype=query.dtype) + core_attn_out = torch.zeros_like(value) + mask2 = torch.triu( + torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), + diagonal=1, + ) + + for i in range(num_chunks): + q_i = query[:, :, i] + k_i = key[:, :, i] + v_i = value[:, :, i] + + attn_i = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_( + mask2, 0 + ) + + v_prime = k_cumdecay[:, :, i] @ last_recurrent_state + v_new = v_i - v_prime + + attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_recurrent_state + core_attn_out[:, :, i] = attn_inter + attn_i @ v_new + + last_recurrent_state = ( + last_recurrent_state * g[:, :, i, -1, None, None].exp() + + ( + k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None] + ).transpose(-1, -2) + @ v_new + ) + + core_attn_out = core_attn_out.reshape(B, H, -1, v_dim) + core_attn_out = core_attn_out[:, :, :S] + + if not output_final_state: + last_recurrent_state = None + + return core_attn_out, last_recurrent_state + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask=None, + position_ids=None, + past_key_value=None, + **kwargs, + ): + """Forward pass compatible with NxDI decoder layer interface.""" + batch_size, seq_len, _ = hidden_states.shape + + seq_ids = kwargs.get("seq_ids", None) + is_for_context_encoding = bool(kwargs.get("is_for_context_encoding", False)) + qwen_chunked_prefill_active = ( + self.use_qwen_hybrid_chunked_prefill + and past_key_value is not None + and seq_len > 1 + ) + is_decode = ( + past_key_value is not None + and not qwen_chunked_prefill_active + and not is_for_context_encoding + ) + + # Padding mask for DeltaNet: [B, S, 1] with 1.0 for real tokens, 0.0 for padding. + # Passed from get_model_output where it's computed from input_ids != pad_token_id. + # Embeddings are already zeroed for padding tokens; this mask additionally + # zeros the decay gate so the recurrent state is preserved unchanged + # through padding positions (no spurious decay). + valid_mask_1d = kwargs.get("deltanet_padding_mask", None) # [B, S, 1] or None + static_hybrid_cache_active = self.use_hybrid_cache_manager + recurrent_state_cache = None + conv_state_cache = None + if static_hybrid_cache_active and past_key_value is not None: + recurrent_state_cache, conv_state_cache = past_key_value + elif ( + self.use_hybrid_apc_manager + and past_key_value is not None + and len(past_key_value) == 2 + and getattr(past_key_value[0], "dim", lambda: 0)() == 4 + and getattr(past_key_value[1], "dim", lambda: 0)() == 3 + and past_key_value[0].shape[1:] == self.recurrent_state_buffer.shape[1:] + and past_key_value[1].shape[1:] == self.conv_state_buffer.shape[1:] + ): + recurrent_state_cache, conv_state_cache = past_key_value + + # Project inputs + deltanet_fp32 = os.environ.get("DELTANET_FP32") == "1" + if deltanet_fp32 and isinstance(self.in_proj_qkv, nn.Linear): + hs_f32 = hidden_states.float() + qkv = F.linear(hs_f32, self.in_proj_qkv.weight.float()).to( + hidden_states.dtype + ) + z = F.linear(hs_f32, self.in_proj_z.weight.float()).to(hidden_states.dtype) + b = F.linear(hs_f32, self.in_proj_b.weight.float()).to(hidden_states.dtype) + a = F.linear(hs_f32, self.in_proj_a.weight.float()).to(hidden_states.dtype) + else: + qkv = self.in_proj_qkv(hidden_states) + z = self.in_proj_z(hidden_states) + b = self.in_proj_b(hidden_states) + a = self.in_proj_a(hidden_states) + + # Split QKV + query = qkv[..., : self.key_dim] + key = qkv[..., self.key_dim : self.key_dim * 2] + value = qkv[..., self.key_dim * 2 :] + + # Causal Conv1d on QKV + mixed = torch.cat([query, key, value], dim=-1) + mixed = mixed.transpose(1, 2) + + if is_decode: + if conv_state_cache is not None: + conv_state = conv_state_cache[:batch_size] + elif seq_ids is not None: + conv_state = torch.index_select(self.conv_state_buffer, 0, seq_ids) + else: + conv_state = self.conv_state_buffer[:batch_size] + conv_input = torch.cat([conv_state, mixed], dim=-1) + + w = self._conv1d_weight().squeeze(1) + if seq_len == 1: + conv_out = ( + conv_input[:, :, : self.conv_kernel_size] * w.unsqueeze(0) + ).sum(dim=-1, keepdim=True) + else: + conv_out = torch.zeros_like(mixed) + for k in range(self.conv_kernel_size): + conv_out = ( + conv_out + + w[:, k].unsqueeze(0).unsqueeze(-1) + * conv_input[:, :, k : k + 1] + ) + mixed_post_conv = F.silu(conv_out) + + new_conv_state = torch.cat([conv_state[:, :, 1:], mixed], dim=-1) + expected_state_len = self.conv_state_buffer.shape[-1] + if new_conv_state.shape[-1] != expected_state_len: + if new_conv_state.shape[-1] > expected_state_len: + new_conv_state = new_conv_state[:, :, -expected_state_len:] + else: + new_conv_state = F.pad( + new_conv_state, + (expected_state_len - new_conv_state.shape[-1], 0), + ) + alloc_bs = self.conv_state_buffer.shape[0] + if static_hybrid_cache_active: + new_conv_state = new_conv_state.to(self.conv_state_buffer.dtype) + elif seq_ids is not None: + new_conv_state = _qwen36_update_state_rows_by_seq_ids( + self.conv_state_buffer, + new_conv_state.to(self.conv_state_buffer.dtype), + seq_ids, + ) + elif batch_size < alloc_bs: + pad_size = alloc_bs - batch_size + new_conv_state = torch.cat( + [ + new_conv_state, + self.conv_state_buffer[batch_size:] * 0, + ], + dim=0, + ) + else: + new_conv_state = new_conv_state + self.conv_state_buffer * 0 + else: + if ( + conv_state_cache is not None + and (qwen_chunked_prefill_active or is_for_context_encoding) + ): + cold_prefill_from_zero = self.use_cold_zero_conv_fast_path + if cold_prefill_from_zero: + mixed_post_conv = F.silu( + F.conv1d( + mixed, + self._conv1d_weight(), + bias=None, + padding=self.conv_kernel_size - 1, + groups=self.conv_dim, + )[:, :, :seq_len] + ) + state_source = mixed + else: + conv_state = conv_state_cache[:batch_size] + if position_ids is not None: + reset_mask = (position_ids[:, :1].long() == 0).to( + dtype=conv_state.dtype, device=conv_state.device + ) + conv_state = conv_state * ( + 1.0 - reset_mask[:, :, None] + ) + conv_input = torch.cat([conv_state, mixed], dim=-1) + w = self._conv1d_weight().squeeze(1) + conv_out = torch.zeros_like(mixed) + for k in range(self.conv_kernel_size): + conv_out = ( + conv_out + + w[:, k].unsqueeze(0).unsqueeze(-1) + * conv_input[:, :, k : k + seq_len] + ) + mixed_post_conv = F.silu(conv_out) + state_source = conv_input + + state_len = self.conv_kernel_size - 1 + if valid_mask_1d is not None: + num_valid = valid_mask_1d.squeeze(-1).sum(dim=-1, keepdim=True).long() + idx_base = (state_source.shape[-1] - seq_len + num_valid - state_len).clamp(min=0) + offsets = torch.arange(state_len, device=mixed.device).unsqueeze(0) + gather_idx = idx_base + offsets + gather_idx = gather_idx.unsqueeze(1).expand(-1, self.conv_dim, -1) + new_conv_state = torch.gather(state_source, 2, gather_idx) + else: + new_conv_state = state_source[:, :, -state_len:].contiguous() + else: + mixed_post_conv = F.silu( + F.conv1d( + mixed, + self._conv1d_weight(), + bias=None, + padding=self.conv_kernel_size - 1, + groups=self.conv_dim, + )[:, :, :seq_len] + ) + + if valid_mask_1d is not None: + # valid_mask_1d is [B, S, 1]; count valid tokens per batch + state_len = self.conv_kernel_size - 1 + num_valid = ( + valid_mask_1d.squeeze(-1).sum(dim=-1, keepdim=True).long() + ) # [B, 1] + idx_base = num_valid - state_len + idx_base = idx_base.clamp(min=0) + offsets = torch.arange(state_len, device=mixed.device).unsqueeze(0) + gather_idx = idx_base + offsets # [B, state_len] + gather_idx = gather_idx.unsqueeze(1).expand(-1, self.conv_dim, -1) + new_conv_state = torch.gather(mixed, 2, gather_idx) + else: + new_conv_state = mixed[:, :, -self.conv_kernel_size + 1 :].contiguous() + + alloc_bs = self.conv_state_buffer.shape[0] + if static_hybrid_cache_active: + new_conv_state = new_conv_state.to(self.conv_state_buffer.dtype) + elif seq_ids is not None: + new_conv_state = _qwen36_update_state_rows_by_seq_ids( + self.conv_state_buffer, + new_conv_state.to(self.conv_state_buffer.dtype), + seq_ids, + ) + elif batch_size < alloc_bs: + pad_size = alloc_bs - batch_size + new_conv_state = torch.cat( + [ + new_conv_state, + torch.zeros( + pad_size, + self.conv_dim, + self.conv_kernel_size - 1, + dtype=new_conv_state.dtype, + device=new_conv_state.device, + ), + ], + dim=0, + ) + new_conv_state = new_conv_state + self.conv_state_buffer * 0 + else: + new_conv_state = new_conv_state + self.conv_state_buffer * 0 + + mixed_post_conv = mixed_post_conv.transpose(1, 2) + + # Zero out conv1d output for padding positions. + # Conv1d with kernel_size=4 leaks real token info into the first + # few padding positions. Zeroing here ensures Q, K, V are exactly + # zero for all padding positions so the recurrence is unaffected. + if valid_mask_1d is not None: + mixed_post_conv = ( + mixed_post_conv * valid_mask_1d + ) # [B, S, conv_dim] * [B, S, 1] + + query = mixed_post_conv[..., : self.key_dim] + key = mixed_post_conv[..., self.key_dim : self.key_dim * 2] + value = mixed_post_conv[..., self.key_dim * 2 :] + + # Reshape to heads + query = query.reshape(batch_size, seq_len, self.num_k_heads, self.head_k_dim) + key = key.reshape(batch_size, seq_len, self.num_k_heads, self.head_k_dim) + value = value.reshape(batch_size, seq_len, self.num_v_heads, self.head_v_dim) + + # Compute gating + beta = b.sigmoid() + g = -self._A_log().float().exp() * F.softplus(a.float() + self._dt_bias()) + + if valid_mask_1d is not None: + # Zero g for padding → alpha=exp(0)=1 → state preserved through padding + # Zero beta for padding → no state update from padding tokens + mask_2d = valid_mask_1d.squeeze(-1).float() # [B, S] + g = g * mask_2d.unsqueeze(-1) + beta = beta * mask_2d.unsqueeze(-1) + + # Expand K heads to match V heads (16 -> 48) using expand+reshape + if self.num_v_heads // self.num_k_heads > 1: + rep = self.num_v_heads // self.num_k_heads # 3 + query = ( + query.unsqueeze(3) + .expand(-1, -1, -1, rep, -1) + .reshape(batch_size, seq_len, self.num_v_heads, self.head_k_dim) + ) + key = ( + key.unsqueeze(3) + .expand(-1, -1, -1, rep, -1) + .reshape(batch_size, seq_len, self.num_v_heads, self.head_k_dim) + ) + + # Transpose to (B, H, S, dim) + query = query.transpose(1, 2).contiguous().float() + key = key.transpose(1, 2).contiguous().float() + value = value.transpose(1, 2).contiguous().float() + g = g.transpose(1, 2).contiguous().float() + beta = beta.transpose(1, 2).contiguous().float() + + if is_decode: + # TKG: single-step recurrent update + if recurrent_state_cache is not None: + recurrent_state = recurrent_state_cache[:batch_size] + elif seq_ids is not None: + recurrent_state = torch.index_select( + self.recurrent_state_buffer, 0, seq_ids + ) + else: + recurrent_state = self.recurrent_state_buffer[:batch_size] + + use_nki_decode = ( + self.use_qwen_deltanet_decode_nki + or os.environ.get("USE_NKI_DECODE") == "1" + ) + if use_nki_decode and seq_len == 1: + output, new_state = self._nki_recurrent_step( + query, key, value, g, beta, recurrent_state + ) + else: + output, new_state = self._recurrent_step( + query, key, value, g, beta, recurrent_state.float() + ) + new_state_bf16 = new_state.to(self.recurrent_state_buffer.dtype) + alloc_bs = self.recurrent_state_buffer.shape[0] + if static_hybrid_cache_active: + new_rec_state = new_state_bf16 + elif seq_ids is not None: + new_rec_state = _qwen36_update_state_rows_by_seq_ids( + self.recurrent_state_buffer, + new_state_bf16, + seq_ids, + ) + elif batch_size < alloc_bs: + new_rec_state = torch.cat( + [ + new_state_bf16, + self.recurrent_state_buffer[batch_size:] * 0, + ], + dim=0, + ) + else: + new_rec_state = new_state_bf16 + self.recurrent_state_buffer * 0 + else: + # CTE: fused NKI kernel by default (PyTorch _chunk_forward can hit + # neuronx-cc codegen ICE NCC_INLA001 with these DeltaNet dimensions). + # Override with env vars for debugging/benchmarking. + use_nki_fused = os.environ.get("USE_NKI_FUSED", "1") != "0" + use_nki_chunked = os.environ.get("USE_NKI_CHUNKED") == "1" + use_nki = os.environ.get("USE_NKI") == "1" + use_sequential = os.environ.get("DELTANET_SEQUENTIAL") == "1" + use_pytorch_chunk = os.environ.get("USE_PYTORCH_CHUNK") == "1" + use_autocp_cte = os.environ.get("QWEN36_DELTANET_AUTOCP_CTE") == "1" + use_compact_autocp_cte = ( + os.environ.get("QWEN36_DELTANET_COMPACT_AUTOCP_CTE") == "1" + ) + + if recurrent_state_cache is not None and ( + qwen_chunked_prefill_active or is_for_context_encoding + ): + initial_state = recurrent_state_cache[:batch_size].float() + if position_ids is not None: + reset_mask = (position_ids[:, :1].long() == 0).to( + dtype=initial_state.dtype, device=initial_state.device + ) + initial_state = initial_state * (1.0 - reset_mask[:, :, None, None]) + if use_autocp_cte and use_compact_autocp_cte: + output, final_state = self._compact_autocp_chunked_forward( + query, + key, + value, + g, + beta, + output_final_state=True, + initial_state=initial_state, + ) + elif use_autocp_cte: + output, final_state = self._autocp_chunked_forward( + query, + key, + value, + g, + beta, + output_final_state=True, + initial_state=initial_state, + ) + elif use_nki_chunked or ( + self.use_qwen_hybrid_chunked_prefill_nki + and os.environ.get("USE_NKI_FUSED", "1") == "0" + ): + output, final_state = self._nki_chunked_forward( + query, + key, + value, + g, + beta, + output_final_state=True, + initial_state=initial_state, + ) + elif use_pytorch_chunk: + output, final_state = self._chunk_forward( + query, + key, + value, + g, + beta, + output_final_state=True, + initial_state=initial_state, + ) + else: + output, final_state = self._fused_chunked_forward( + query, + key, + value, + g, + beta, + output_final_state=True, + initial_state=initial_state, + ) + elif use_pytorch_chunk: + output, final_state = self._chunk_forward( + query, key, value, g, beta, output_final_state=True + ) + elif use_autocp_cte and use_compact_autocp_cte: + output, final_state = self._compact_autocp_chunked_forward( + query, key, value, g, beta, output_final_state=True + ) + elif use_autocp_cte: + output, final_state = self._autocp_chunked_forward( + query, key, value, g, beta, output_final_state=True + ) + elif use_nki_chunked: + output, final_state = self._nki_chunked_forward( + query, key, value, g, beta, output_final_state=True + ) + elif use_nki: + output, final_state = self._nki_recurrent_forward( + query, key, value, g, beta + ) + elif use_sequential: + output, final_state = self._sequential_forward( + query, key, value, g, beta, output_final_state=True + ) + elif use_nki_fused: + output, final_state = self._fused_chunked_forward( + query, key, value, g, beta, output_final_state=True + ) + else: + output, final_state = self._fused_chunked_forward( + query, key, value, g, beta, output_final_state=True + ) + + if final_state is not None: + final_state_bf16 = final_state.to(self.recurrent_state_buffer.dtype) + alloc_bs = self.recurrent_state_buffer.shape[0] + if static_hybrid_cache_active: + new_rec_state = final_state_bf16 + elif seq_ids is not None: + new_rec_state = _qwen36_update_state_rows_by_seq_ids( + self.recurrent_state_buffer, + final_state_bf16, + seq_ids, + ) + elif batch_size < alloc_bs: + new_rec_state = torch.cat( + [ + final_state_bf16, + torch.zeros( + alloc_bs - batch_size, + self.num_v_heads, + self.head_k_dim, + self.head_v_dim, + dtype=final_state_bf16.dtype, + device=final_state_bf16.device, + ), + ], + dim=0, + ) + new_rec_state = new_rec_state + self.recurrent_state_buffer * 0 + else: + new_rec_state = final_state_bf16 + self.recurrent_state_buffer * 0 + else: + new_rec_state = self.recurrent_state_buffer * 1 + + if ( + is_for_context_encoding + and not static_hybrid_cache_active + and valid_mask_1d is not None + and hasattr(valid_mask_1d, "numel") + and valid_mask_1d.numel() > 0 + ): + active_rows = _qwen36_active_state_rows(valid_mask_1d, seq_ids) + new_conv_state = _qwen36_preserve_inactive_state_rows( + new_conv_state, + self.conv_state_buffer, + active_rows, + ) + new_rec_state = _qwen36_preserve_inactive_state_rows( + new_rec_state, + self.recurrent_state_buffer, + active_rows, + ) + + # Output: norm, gate, project + output = output.to(hidden_states.dtype) + output = output.transpose(1, 2).contiguous() + output = output.reshape(batch_size, seq_len, self.num_v_heads, self.head_v_dim) + output = self.norm(output) + z_gate = z.reshape(batch_size, seq_len, self.num_v_heads, self.head_v_dim) + output = output * F.silu(z_gate) + output = output.reshape(batch_size, seq_len, self.value_dim) + output = self.out_proj(output) + + if static_hybrid_cache_active: + return output, (new_rec_state, new_conv_state), new_rec_state, new_conv_state + + # Return dummy KV for KVCacheManager + dummy_k = torch.zeros( + batch_size, + self.kv_heads_per_rank, + seq_len, + self.head_dim, + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + dummy_v = torch.zeros_like(dummy_k) + + return output, (dummy_k, dummy_v), new_rec_state, new_conv_state + + +# ============================================================ +# InferenceConfig (Dense -- no MoE) +# ============================================================ + + +class Qwen35InferenceConfig(InferenceConfig): + """Config for Qwen3.5/3.6-27B (dense) with hybrid DeltaNet + Attention.""" + + @classmethod + def from_pretrained(cls, model_path: str, **kwargs) -> "Qwen35InferenceConfig": + """Load Qwen3.5/Qwen3.6 text config from a pretrained model directory. + + Qwen3.6 stores the decoder settings under the top-level multimodal + `text_config`. NxDI's text-only inference config expects those fields + flattened onto the inference config itself. + """ + neuron_config = kwargs.pop("neuron_config", None) + if neuron_config is None: + neuron_config = NeuronConfig( + tp_degree=1, + batch_size=1, + seq_len=128, + torch_dtype=torch.bfloat16, + save_sharded_checkpoint=True, + ) + + config_path = os.path.join(model_path, "config.json") + if not os.path.exists(config_path): + raise FileNotFoundError(f"Configuration file not found at {config_path}") + + with open(config_path, "r", encoding="utf-8") as handle: + config_dict = json.load(handle) + + text_config = config_dict.get("text_config", config_dict) + rope_parameters = text_config.get("rope_parameters") or {} + inference_config = dict(text_config) + inference_config.setdefault("_name_or_path", model_path) + # Preserve original model_type (may be qwen3_5_moe_text for MoE variants). + inference_config.setdefault("model_type", text_config.get("model_type", "qwen3_5_text")) + inference_config.setdefault("architectures", config_dict.get("architectures", [])) + inference_config.setdefault("tie_word_embeddings", config_dict.get("tie_word_embeddings", False)) + # For MoE variants: expose intermediate_size (needed by initialize_moe_module). + # HF stores moe_intermediate_size and shared_expert_intermediate_size; we use + # moe_intermediate_size as the per-expert MLP dim. + if "moe_intermediate_size" in text_config and "intermediate_size" not in inference_config: + inference_config["intermediate_size"] = text_config["moe_intermediate_size"] + if "rope_theta" not in inference_config and "rope_theta" in rope_parameters: + inference_config["rope_theta"] = rope_parameters["rope_theta"] + if ( + "partial_rotary_factor" not in inference_config + and "partial_rotary_factor" in rope_parameters + ): + inference_config["partial_rotary_factor"] = rope_parameters[ + "partial_rotary_factor" + ] + inference_config.update(kwargs) + return cls(neuron_config=neuron_config, **inference_config) + + def __init__(self, *args, **kwargs): + # Set defaults BEFORE super().__init__() because it calls validate_config() + # which checks get_required_attributes(). These can be overridden by + # kwargs or load_config. + + # ── MoE detection & defaults ────────────────────────────────────────── + # Qwen3.5-35B-A3B (model_type=qwen3_5_moe_text) has: + # num_experts, num_experts_per_tok, moe_intermediate_size, + # shared_expert_intermediate_size, router_aux_loss_coef. + # For those variants we expose the MoE fields that initialize_moe_module() + # requires: num_local_experts, n_shared_experts, and use + # moe_intermediate_size as the routed-experts intermediate dim. + self._is_moe = str(kwargs.get("model_type", "")) == "qwen3_5_moe_text" or ( + "num_experts" in kwargs and "moe_intermediate_size" in kwargs + ) + if self._is_moe: + kwargs.setdefault("num_local_experts", kwargs.get("num_experts", 256)) + # Qwen3.5-MoE has a single shared expert with a per-token sigmoid gate. + kwargs.setdefault("n_shared_experts", 1) + # Route the moe_intermediate_size into the "intermediate_size" field + # expected by NxDI's initialize_moe_module (which sizes the routed + # experts). Shared expert uses shared_expert_intermediate_size, which + # for A3B equals moe_intermediate_size (both 512). + if "intermediate_size" not in kwargs and "moe_intermediate_size" in kwargs: + kwargs["intermediate_size"] = kwargs["moe_intermediate_size"] + kwargs.setdefault( + "shared_expert_intermediate_size", + kwargs.get("moe_intermediate_size", kwargs["intermediate_size"]), + ) + + # Layer types for hybrid dispatch: [3 DeltaNet + 1 GQA] repeated. + if "layer_types" not in kwargs and not any( + hasattr(a, "layer_types") for a in args if hasattr(a, "__dict__") + ): + num_layers = kwargs.get("num_hidden_layers", 64) + if num_layers % 4 != 0: + raise ValueError( + f"Qwen3.5 hybrid layer count must be divisible by 4, got {num_layers}" + ) + layer_types = [] + for _ in range(num_layers // 4): + layer_types.extend( + [ + "linear_attention", + "linear_attention", + "linear_attention", + "full_attention", + ] + ) + kwargs.setdefault("layer_types", layer_types) + + # DeltaNet-specific config defaults + kwargs.setdefault("linear_num_value_heads", 48) + kwargs.setdefault("linear_num_key_heads", 16) + kwargs.setdefault("linear_key_head_dim", 128) + kwargs.setdefault("linear_value_head_dim", 128) + kwargs.setdefault("linear_conv_kernel_dim", 4) + kwargs.setdefault("use_hybrid_cache_manager", False) + kwargs.setdefault("use_hybrid_apc_manager", False) + kwargs.setdefault("use_qwen_hybrid_chunked_prefill", False) + kwargs.setdefault("use_qwen_hybrid_chunked_prefill_nki", False) + kwargs.setdefault("use_qwen_deltanet_decode_nki", False) + kwargs.setdefault("gdn_checkpoint_interval", 256) + kwargs.setdefault("max_gdn_checkpoint_slots", 8) + kwargs.setdefault("hybrid_apc_layout_version", 1) + kwargs.setdefault("hybrid_apc_allow_residual_replay", False) + kwargs.setdefault("hybrid_apc_cache_salt", None) + use_hybrid_apc_manager = bool(kwargs.get("use_hybrid_apc_manager", False)) + kwargs.setdefault( + "hybrid_apc_require_vllm_metadata", use_hybrid_apc_manager + ) + kwargs.setdefault( + "hybrid_apc_allow_local_hash_fallback", not use_hybrid_apc_manager + ) + kwargs.setdefault( + "hybrid_apc_require_attention_block_refs", use_hybrid_apc_manager + ) + kwargs.setdefault("hybrid_apc_reject_unbacked_attention_hits", True) + kwargs.setdefault("hybrid_apc_disable_unbacked_prefix_reads", False) + kwargs.setdefault("hybrid_apc_enable_backed_prefix_reads", False) + kwargs.setdefault( + "hybrid_apc_model_revision", + kwargs.get("_name_or_path", kwargs.get("model_revision", "unknown")), + ) + kwargs.setdefault( + "hybrid_recurrent_cache_dtype", + kwargs.get("gdn_recurrent_cache_dtype", "float32"), + ) + kwargs.setdefault( + "hybrid_conv_cache_dtype", + kwargs.get("gdn_conv_cache_dtype", "bfloat16"), + ) + kwargs.setdefault( + "gdn_recurrent_cache_dtype", kwargs["hybrid_recurrent_cache_dtype"] + ) + kwargs.setdefault("gdn_conv_cache_dtype", kwargs["hybrid_conv_cache_dtype"]) + kwargs.setdefault("hybrid_cache_mode", "all") + kwargs.setdefault( + "hybrid_cache_prefix_boundary_only", + kwargs.get("hybrid_cache_block_boundary_only", True), + ) + kwargs.setdefault( + "hybrid_cache_block_boundary_only", + kwargs["hybrid_cache_prefix_boundary_only"], + ) + kwargs.setdefault("hybrid_cache_validate_exact", False) + kwargs.setdefault("use_text_only_cte_inputs", True) + kwargs.setdefault("use_compact_cte_attention_mask", True) + kwargs.setdefault("use_cold_zero_conv_fast_path", False) + kwargs.setdefault("disable_token_generation_wlo", False) + + super().__init__(*args, **kwargs) + + self.gdn_checkpoint_interval = int(self.gdn_checkpoint_interval) + if self.gdn_checkpoint_interval <= 0: + raise ValueError( + "gdn_checkpoint_interval must be positive, " + f"got {self.gdn_checkpoint_interval}" + ) + self.max_gdn_checkpoint_slots = int(self.max_gdn_checkpoint_slots) + if self.max_gdn_checkpoint_slots <= 0: + raise ValueError( + "max_gdn_checkpoint_slots must be positive, " + f"got {self.max_gdn_checkpoint_slots}" + ) + self.hybrid_apc_layout_version = int(self.hybrid_apc_layout_version) + self.hybrid_recurrent_cache_dtype = _normalize_hybrid_cache_dtype( + "hybrid_recurrent_cache_dtype", + self.hybrid_recurrent_cache_dtype, + "float32", + ) + self.hybrid_conv_cache_dtype = _normalize_hybrid_cache_dtype( + "hybrid_conv_cache_dtype", + self.hybrid_conv_cache_dtype, + "bfloat16", + ) + self.gdn_recurrent_cache_dtype = self.hybrid_recurrent_cache_dtype + self.gdn_conv_cache_dtype = self.hybrid_conv_cache_dtype + self.hybrid_cache_block_boundary_only = ( + self.hybrid_cache_prefix_boundary_only + ) + self.hybrid_apc_require_vllm_metadata = bool( + self.hybrid_apc_require_vllm_metadata + ) + self.hybrid_apc_allow_local_hash_fallback = bool( + self.hybrid_apc_allow_local_hash_fallback + ) + self.hybrid_apc_require_attention_block_refs = bool( + self.hybrid_apc_require_attention_block_refs + ) + self.hybrid_apc_reject_unbacked_attention_hits = bool( + self.hybrid_apc_reject_unbacked_attention_hits + ) + self.hybrid_apc_disable_unbacked_prefix_reads = bool( + self.hybrid_apc_disable_unbacked_prefix_reads + ) + if self.hybrid_apc_require_vllm_metadata: + self.hybrid_apc_allow_local_hash_fallback = False + self.hybrid_apc_require_attention_block_refs = True + self.hybrid_apc_reject_unbacked_attention_hits = True + if self.use_hybrid_cache_manager and self.use_hybrid_apc_manager: + raise ValueError( + "use_hybrid_cache_manager and use_hybrid_apc_manager are mutually exclusive" + ) + if self.use_hybrid_apc_manager and self.hybrid_cache_mode != "all": + raise ValueError("use_hybrid_apc_manager requires hybrid_cache_mode='all'") + if self.use_hybrid_apc_manager: + if self.hybrid_recurrent_cache_dtype != "float32": + raise ValueError( + "use_hybrid_apc_manager requires float32 recurrent GDN " + "checkpoint cache state; bf16 checkpoint roundtrips are not " + "coherent for all-mode prefix caching" + ) + pa_block_size = getattr(self.neuron_config, "pa_block_size", None) + if pa_block_size is not None and self.gdn_checkpoint_interval != int( + pa_block_size + ): + raise ValueError( + "use_hybrid_apc_manager v0 requires " + "gdn_checkpoint_interval == pa_block_size" + ) + if self.hybrid_apc_allow_residual_replay: + raise ValueError( + "hybrid_apc_allow_residual_replay is reserved for v1; " + "v0 restores only exact checkpoint boundaries" + ) + + # Attention output gate + self.attn_output_gate = getattr(self, "attn_output_gate", True) + + # Partial RoPE + self.partial_rotary_factor = getattr(self, "partial_rotary_factor", 0.25) + self.rope_dim = int(self.head_dim * self.partial_rotary_factor) # 64 + + # mRoPE (multimodal RoPE) for VL support + rope_params = getattr(self, "rope_parameters", {}) or {} + self.mrope_section = rope_params.get("mrope_section", [11, 11, 10]) + self.mrope_interleaved = rope_params.get("mrope_interleaved", True) + + # Standard HF config attributes expected by NxDI + if not hasattr(self, "output_attentions"): + self.output_attentions = False + if not hasattr(self, "output_hidden_states"): + self.output_hidden_states = False + + def get_required_attributes(self) -> List[str]: + return [ + "head_dim", + "hidden_act", + "hidden_size", + "intermediate_size", + "max_position_embeddings", + "num_attention_heads", + "num_hidden_layers", + "num_key_value_heads", + "rms_norm_eps", + "rope_theta", + "vocab_size", + # DeltaNet-specific + "linear_num_value_heads", + "linear_num_key_heads", + "linear_key_head_dim", + "linear_value_head_dim", + "linear_conv_kernel_dim", + "layer_types", + ] + + @classmethod + def get_neuron_config_cls(cls): + # MoE variants (qwen3_5_moe_text) use MoENeuronConfig so that + # initialize_moe_module() finds router_config / blockwise_matmul_config + # / moe_tp_degree / etc. Dense variants keep the plain NeuronConfig path. + # This is a class method but at call-time we don't have an instance; + # callers construct NeuronConfig explicitly, so we return the correct + # default for the common (dense) case here. + return NeuronConfig + + +# ============================================================ +# Attention (standard GQA for 16 of 64 layers) +# With output gate: q_proj is 2x sized, split into (query, gate) +# With partial RoPE: only first rope_dim dimensions get rotary +# ============================================================ + + +class Qwen35MRoPEEmbedding(nn.Module): + """Multimodal Rotary Position Embedding (mRoPE) for Qwen3.5. + + Handles 3D position information (temporal, height, width) for VL models. + Position IDs have shape (3, batch_size, seq_len) for T/H/W dimensions. + For text-only (2D position_ids), broadcasts to 3D with identical positions. + """ + + def __init__(self, config): + super().__init__() + self.head_dim = config.head_dim # 256 + self.rope_dim = config.rope_dim # 64 + self.mrope_section = config.mrope_section # [11, 11, 10] + self.mrope_interleaved = getattr(config, "mrope_interleaved", True) + self.rope_theta = config.rope_theta + + # Validate mrope_section sums to rope_dim // 2 = 32 + assert sum(self.mrope_section) == self.rope_dim // 2, ( + f"mrope_section {self.mrope_section} sums to {sum(self.mrope_section)}, " + f"expected {self.rope_dim // 2}" + ) + + def forward(self, x, position_ids_3d): + """Compute cos/sin from 3D position IDs. + + Args: + x: hidden_states (for device/dtype inference) + position_ids_3d: (3, batch_size, seq_len) -- T, H, W positions + + Returns: + cos: (batch_size, seq_len, rope_dim) + sin: (batch_size, seq_len, rope_dim) + """ + device = x.device + dtype = torch.float32 + + if position_ids_3d.ndim == 2: + position_ids_3d = position_ids_3d[None, ...].expand( + 3, position_ids_3d.shape[0], -1 + ) + + inv_freq = 1.0 / ( + self.rope_theta + ** ( + torch.arange(0, self.rope_dim, 2, dtype=dtype, device=device) + / self.rope_dim + ) + ) + inv_freq = inv_freq[None, None, :, None].expand( + 3, position_ids_3d.shape[1], -1, 1 + ) + positions = position_ids_3d[:, :, None, :].float() + freqs = (inv_freq.float() @ positions).transpose(2, 3) + + # Match HF Qwen3.6 mRoPE layout exactly: start from the temporal + # frequencies, then splice H/W frequencies into interleaved positions. + freqs_t = freqs[0] + if self.mrope_interleaved: + for dim, offset in enumerate((1, 2), start=1): + length = self.mrope_section[dim] * 3 + idx = slice(offset, length, 3) + freqs_t[..., idx] = freqs[dim, ..., idx] + + emb = torch.cat((freqs_t, freqs_t), dim=-1) + cos = emb.cos().to(dtype=x.dtype) + sin = emb.sin().to(dtype=x.dtype) + + return cos, sin + + +class NeuronQwen35Attention(NeuronAttentionBase): + """Standard GQA attention for Qwen3.5 with output gate and partial RoPE. + + 24 Q heads, 4 KV heads (6:1 GQA), head_dim=256 for 27B dense. + q_proj is doubled (query + gate), split at load time. + Only first rope_dim=64 of head_dim=256 gets rotary encoding. + + Uses NeuronAttentionBase infrastructure for QKV projection, KV cache, + RoPE, and attention computation. Overrides forward() to insert the + sigmoid output gate between attention output and o_proj. + """ + + def __init__(self, config): + # Partial RoPE: create mRoPE embedding with rope_dim (64) + self.rope_dim = config.rope_dim # 64 = head_dim * partial_rotary_factor + + # Create QK norm modules (will be passed to base class) + rms_norm_eps = config.rms_norm_eps + q_ln = get_rmsnorm_cls()(config.head_dim, rms_norm_eps) + k_ln = get_rmsnorm_cls()(config.head_dim, rms_norm_eps) + + # Partial RoPE: use standard RotaryEmbedding. + # For VL with 3D mRoPE positions, cos/sin are pre-computed externally in + # get_model_output() using Qwen35MRoPEEmbedding and passed as cos_cache/sin_cache. + rotary_emb = RotaryEmbedding( + self.rope_dim, # Only 64 dims get rotary embedding + max_position_embeddings=config.max_position_embeddings, + base=config.rope_theta, + ) + super().__init__( + config=config, + hidden_size=config.hidden_size, + num_attention_heads=config.num_attention_heads, + num_key_value_heads=config.num_key_value_heads, + head_dim=config.head_dim, + rotary_emb=rotary_emb, + rms_norm_eps=rms_norm_eps, + use_qk_norm=False, + q_layernorm=q_ln, + k_layernorm=k_ln, + ) + + # Separate mRoPE module for VL 3D position_ids + self.mrope_emb = Qwen35MRoPEEmbedding(config) + + # Output gate projection: hidden_size -> num_heads * head_dim + # Populated from the second half of q_proj during state dict conversion. + self.output_gate_proj = ColumnParallelLinear( + config.hidden_size, + config.num_attention_heads * config.head_dim, + bias=False, + gather_output=False, + ) + + self.qwen_output_gate_nki_kernel_enabled = bool( + getattr(config, "use_qwen_output_gate_nki", False) + or os.environ.get("QWEN36_OUTPUT_GATE_NKI", "0") == "1" + ) + self.qwen_qkv_gate_packed_enabled = bool( + getattr(config, "use_qwen_qkv_gate_packed", False) + or os.environ.get("QWEN36_QKV_GATE_PACKED", "0") == "1" + ) + self.qwen_gated_o_proj_nki_kernel_enabled = bool( + getattr(config, "use_qwen_gated_o_proj_nki", False) + or os.environ.get("QWEN36_GATED_OUT_PROJ_NKI", "0") == "1" + ) + if ( + self.qwen_output_gate_nki_kernel_enabled + and self.qwen_qkv_gate_packed_enabled + ): + raise ValueError( + "Qwen output-gate NKI and packed QKV+gate are mutually exclusive." + ) + if self.qwen_output_gate_nki_kernel_enabled: + if _qwen_gate_projection_kernel is None: + raise ImportError( + "QWEN36_OUTPUT_GATE_NKI requires nkilib.core.qkv.qkv" + ) + if getattr(config.neuron_config, "quantized", False): + setattr( + self.output_gate_proj, + "post_create_quantized_module_hook", + preprocess_quantized_linear_layer, + ) + else: + self.output_gate_proj.weight = transpose_parallel_linear_layer( + self.output_gate_proj.weight + ) + + if self.qwen_qkv_gate_packed_enabled: + if _qwen_gate_projection_kernel is None: + raise ImportError( + "QWEN36_QKV_GATE_PACKED requires nkilib.core.qkv.qkv" + ) + if not self.fused_qkv: + raise ValueError("QWEN36_QKV_GATE_PACKED requires fused_qkv=True") + self._enable_qwen_qkv_gate_packed_projection(config) + + self.qwen_qk_norm_rope_nki_kernel_enabled = bool( + getattr(config, "use_qwen_qk_norm_rope_nki", False) + or os.environ.get("QWEN36_QK_NORM_ROPE_NKI", "0") == "1" + ) + if self.qwen_qk_norm_rope_nki_kernel_enabled: + if _qwen_qk_norm_partial_rope_kernel is None: + raise ImportError( + "QWEN36_QK_NORM_ROPE_NKI requires src.nki_kernels." + "qwen_qk_norm_rope" + ) + if self.head_dim != 256 or self.rope_dim != 64: + raise ValueError( + "Qwen Q/K norm+RoPE NKI kernel currently supports only " + f"head_dim=256 and rope_dim=64, got head_dim={self.head_dim}, " + f"rope_dim={self.rope_dim}" + ) + + self.qkv_tkg_nki_kernel_enabled = bool( + getattr(config.neuron_config, "qkv_tkg_nki_kernel_enabled", False) + ) and not bool(getattr(config.neuron_config, "is_prefill_stage", False)) + if self.qkv_tkg_nki_kernel_enabled: + if _qkv_tkg_nki_kernel is None: + raise ImportError( + "qkv_tkg_nki_kernel_enabled requires " + "neuronxcc.nki._pre_prod_kernels.qkv_tkg_impl" + ) + if self.fused_qkv: + raise ValueError( + "qkv_tkg_nki_kernel_enabled uses split q/k/v projections " + "and must not be combined with fused_qkv" + ) + if self.qkv_proj_sp_enabled: + raise ValueError( + "qkv_tkg_nki_kernel_enabled does not support sequence-parallel " + "QKV projection" + ) + qkv_proj = self.get_qkv_proj() + split_qkv_projections = ( + qkv_proj.q_proj, + qkv_proj.k_proj, + qkv_proj.v_proj, + ) + for projection in split_qkv_projections: + if not getattr(config.neuron_config, "quantized", False): + projection.weight = transpose_parallel_linear_layer(projection.weight) + + def _enable_qwen_qkv_gate_packed_projection(self, config): + for attr_name in ("qkv_proj", "cte_qkv_proj", "tkg_qkv_proj"): + qkv_proj = getattr(self, attr_name, None) + if qkv_proj is not None and getattr(qkv_proj, "fused_qkv", False): + self._replace_qkv_projection_with_qwen_qkvgate(qkv_proj, config) + + def _replace_qkv_projection_with_qwen_qkvgate(self, qkv_proj, config): + if not hasattr(qkv_proj, "Wqkv"): + raise ValueError("QWEN36_QKV_GATE_PACKED requires a fused Wqkv module") + if not isinstance(qkv_proj.Wqkv, ColumnParallelLinear): + raise ValueError( + "QWEN36_QKV_GATE_PACKED currently supports ColumnParallelLinear Wqkv" + ) + + packed_q_heads = qkv_proj.num_attention_heads * 2 + packed_output_size = ( + packed_q_heads + 2 * qkv_proj.num_key_value_heads + ) * qkv_proj.head_dim + packed_wqkv = ColumnParallelLinear( + qkv_proj.hidden_size, + packed_output_size, + bias=qkv_proj.bias, + gather_output=qkv_proj.gather_output, + dtype=qkv_proj.dtype, + sequence_parallel_enabled=False, + tensor_model_parallel_group=qkv_proj.tensor_model_parallel_group, + rank_ordering=qkv_proj.rank_ordering, + ) + if ( + (qkv_proj.qkv_kernel_enabled or qkv_proj.qkv_nki_kernel_enabled) + and getattr(config.neuron_config, "quantized", False) + ): + setattr( + packed_wqkv, + "post_create_quantized_module_hook", + preprocess_quantized_linear_layer, + ) + elif qkv_proj.qkv_kernel_enabled or qkv_proj.qkv_nki_kernel_enabled: + packed_wqkv.weight = transpose_parallel_linear_layer(packed_wqkv.weight) + + for param in ( + [packed_wqkv.weight, packed_wqkv.scale] + if hasattr(packed_wqkv, "scale") + else [packed_wqkv.weight] + ): + setattr(param, "fused_qkv", True) + setattr(param, "num_attention_heads", packed_q_heads) + setattr(param, "num_key_value_heads", qkv_proj.num_key_value_heads) + setattr(param, "head_dim", qkv_proj.head_dim) + if qkv_proj.bias: + setattr(packed_wqkv.bias, "fused_qkv", True) + setattr(packed_wqkv.bias, "num_attention_heads", packed_q_heads) + setattr(packed_wqkv.bias, "num_key_value_heads", qkv_proj.num_key_value_heads) + setattr(packed_wqkv.bias, "head_dim", qkv_proj.head_dim) + + qkv_proj.Wqkv = packed_wqkv + qkv_proj.qwen_qkv_gate_packed = True + qkv_proj.qwen_real_num_attention_heads = qkv_proj.num_attention_heads + qkv_proj.qwen_packed_num_attention_heads = packed_q_heads + + @staticmethod + def _apply_projection_scale(output, projection): + scale = getattr(projection, "scale", None) + if scale is None: + return output + scale_tensor = scale.data if hasattr(scale, "data") else scale + if ( + scale_tensor.ndim == 2 + and scale_tensor.shape[0] == 128 + and scale_tensor.shape[1] == output.shape[-1] + ): + scale_tensor = scale_tensor[0] + else: + scale_tensor = scale_tensor.reshape(-1) + if scale_tensor.numel() != output.shape[-1]: + raise ValueError( + "QKV TKG projection scale shape does not match output width: " + f"scale={tuple(scale.shape)}, output={tuple(output.shape)}" + ) + return output * scale_tensor.reshape(1, 1, output.shape[-1]).to(output.dtype) + + @staticmethod + def _prepare_qkv_tkg_scale(scale_tensor, output_width): + if ( + scale_tensor.ndim == 2 + and scale_tensor.shape[0] == 128 + and scale_tensor.shape[1] == output_width + ): + return scale_tensor.contiguous() + if ( + scale_tensor.ndim == 2 + and scale_tensor.shape[0] == output_width + and scale_tensor.shape[1] == 1 + ): + return torch.broadcast_to( + scale_tensor.transpose(0, 1), + (128, output_width), + ).contiguous() + if ( + scale_tensor.ndim == 2 + and scale_tensor.shape[0] == 1 + and scale_tensor.shape[1] == output_width + ): + return torch.broadcast_to(scale_tensor, (128, output_width)).contiguous() + if scale_tensor.numel() == output_width: + return torch.broadcast_to( + scale_tensor.reshape(1, output_width), + (128, output_width), + ).contiguous() + raise ValueError( + "QKV TKG projection scale shape does not match output width: " + f"scale={tuple(scale_tensor.shape)}, output_width={output_width}" + ) + + def _run_split_qkv_tkg_projection(self, hidden_states, projection, local_heads): + bias = ( + projection.bias.data.unsqueeze(0) + if getattr(projection, "bias", None) is not None + else None + ) + weight = projection.weight.data + if weight.shape[0] != self.hidden_size and weight.shape[1] == self.hidden_size: + weight = weight.transpose(0, 1).contiguous() + # The preprod QKV TKG kernel's LNC2 path reduces across pi0 and then + # stores both programs to the same shared-HBM slice, which the current + # NKI verifier rejects as an output dependency. Use the single-LNC + # variant for this split projection until that kernel store is fixed. + kernel = _qkv_tkg_nki_kernel[1] + scale = getattr(projection, "scale", None) + if scale is not None: + scale_tensor = scale.data if hasattr(scale, "data") else scale + qkv_w_scales = self._prepare_qkv_tkg_scale( + scale_tensor, + weight.shape[1], + ) + quantization_type = getattr(_QKVQuantizationType, "ROW", None) + if quantization_type is None: + raise ValueError( + "qkv_tkg_nki_kernel_enabled requires ROW quantization support " + "when running quantized split-QKV projections" + ) + else: + qkv_w_scales = None + quantization_type = _QKVQuantizationType.NONE + + output = kernel( + hidden=hidden_states, + qkv_w=weight, + norm_w=None, + fused_add=False, + mlp_prev=None, + attn_prev=None, + d_head=self.head_dim, + output_layout=_QKVOutputLayout.BSD, + eps=self.rms_norm_eps, + norm_type=_QKVNormType.NO_NORM, + qkvInSB=False, + qkv_bias=bias, + norm_bias=None, + hidden_actual=self.hidden_size, + B=hidden_states.shape[0], + S=hidden_states.shape[1], + H=self.hidden_size, + num_q_heads=local_heads, + num_kv_heads=local_heads, + quantization_type=quantization_type, + qkv_w_scales=qkv_w_scales, + qkv_in_scales=None, + ) + if qkv_w_scales is not None: + return output + return self._apply_projection_scale(output, projection) + + def _prep_split_qkv_tkg_tensors( + self, + position_ids, + hidden_states, + past_key_value, + adapter_ids=None, + cos_cache=None, + sin_cache=None, + use_polar_compatible_rope=False, + ): + # NxDI traces a placeholder adapter_ids tensor even when no LoRA + # adapters are active. Qwen3.6 serving here is non-LoRA, so the split + # projection path intentionally ignores the placeholder. + qkv_proj = self.get_qkv_proj() + Q = self._run_split_qkv_tkg_projection( + hidden_states, + qkv_proj.q_proj, + self.num_heads, + ) + K = self._run_split_qkv_tkg_projection( + hidden_states, + qkv_proj.k_proj, + self.num_key_value_heads, + ) + V = self._run_split_qkv_tkg_projection( + hidden_states, + qkv_proj.v_proj, + self.num_key_value_heads, + ) + + bsz, q_len, _ = hidden_states.size() + Q = move_heads_front( + Q, + bsz, + q_len, + self.num_heads, + self.head_dim, + layernorm=self.q_layernorm, + post_transpose_layernorm=self.post_transpose_layernorm, + ) + K = move_heads_front( + K, + bsz, + q_len, + self.num_key_value_heads, + self.head_dim, + layernorm=self.k_layernorm, + post_transpose_layernorm=self.post_transpose_layernorm, + ) + V = move_heads_front( + V, + bsz, + q_len, + self.num_key_value_heads, + self.head_dim, + layernorm=None, + ) + + Q, K, cos_cache, sin_cache = self.apply_rotary_embedding( + Q, + K, + V, + position_ids, + cos_cache, + sin_cache, + use_polar_compatible_rope, + ) + return Q, K, V, cos_cache, sin_cache, None + + def _should_use_qwen_output_gate_nki(self, q_len): + return self.qwen_output_gate_nki_kernel_enabled + + def _should_use_qwen_qkv_gate_packed(self, q_len): + return ( + self.qwen_qkv_gate_packed_enabled + and not self.qkv_proj_sp_enabled + and _qwen_gate_projection_kernel is not None + ) + + def _should_use_qwen_gated_o_proj_nki(self, q_len): + o_proj = self.get_o_proj() + return ( + self.qwen_gated_o_proj_nki_kernel_enabled + and q_len > 1 + and hasattr(o_proj, "forward_gated") + ) + + def _output_gate_proj_nki(self, hidden_states): + weight = self.output_gate_proj.weight.data + bias = ( + self.output_gate_proj.bias.data.unsqueeze(0) + if getattr(self.output_gate_proj, "bias", None) is not None + else None + ) + + qkv_w_scale = None + qkv_in_scale = None + quantization_type = _NkilibQuantizationType.NONE + gate_scale = getattr(self.output_gate_proj, "scale", None) + if gate_scale is not None: + qkv_w_scale = gate_scale.data + gate_input_scale = getattr(self.output_gate_proj, "input_scale", None) + qkv_in_scale = gate_input_scale.data if gate_input_scale is not None else None + quantization_type = _NkilibQuantizationType.ROW + elif getattr(self.config.neuron_config, "quantized", False): + raise RuntimeError( + "Qwen output-gate NKI path requires output_gate_proj.scale " + "when running a quantized artifact." + ) + + return _qwen_gate_projection_kernel[self.logical_nc_config]( + input=hidden_states, + fused_qkv_weights=weight, + output_layout=_NkilibQKVOutputLayout.BSD, + bias=bias, + quantization_type=quantization_type, + qkv_w_scale=qkv_w_scale, + qkv_in_scale=qkv_in_scale, + ) + + def _qkv_gate_packed_projection_nki(self, hidden_states): + qkv_proj = self.get_qkv_proj() + weight = qkv_proj.Wqkv.weight.data + bias = ( + qkv_proj.Wqkv.bias.data.unsqueeze(0) + if getattr(qkv_proj.Wqkv, "bias", None) is not None + else None + ) + + qkv_w_scale = None + qkv_in_scale = None + quantization_type = _NkilibQuantizationType.NONE + qkv_scale = getattr(qkv_proj.Wqkv, "scale", None) + if qkv_scale is not None: + qkv_w_scale = qkv_scale.data + qkv_input_scale = getattr(qkv_proj.Wqkv, "input_scale", None) + qkv_in_scale = qkv_input_scale.data if qkv_input_scale is not None else None + quantization_type = _NkilibQuantizationType.ROW + elif getattr(self.config.neuron_config, "quantized", False): + raise RuntimeError( + "Qwen packed QKV+gate path requires Wqkv.scale when running " + "a quantized artifact." + ) + + packed = _qwen_gate_projection_kernel[self.logical_nc_config]( + input=hidden_states, + fused_qkv_weights=weight, + output_layout=_NkilibQKVOutputLayout.BSD, + bias=bias, + fused_residual_add=False, + mlp_prev=None, + attention_prev=None, + fused_norm_type=_NkilibNormType.NO_NORM, + gamma_norm_weights=None, + norm_eps=self.rms_norm_eps, + fused_rope=False, + cos_cache=None, + sin_cache=None, + quantization_type=quantization_type, + qkv_w_scale=qkv_w_scale, + qkv_in_scale=qkv_in_scale, + d_head=self.head_dim, + num_q_heads=self.num_heads * 2, + num_kv_heads=self.num_key_value_heads, + ) + + q_width = self.num_heads * self.head_dim + gate_end = q_width * 2 + k_end = gate_end + self.num_key_value_heads * self.head_dim + Q, gate, K, V = torch.tensor_split( + packed, + (q_width, gate_end, k_end), + dim=2, + ) + return Q, gate, K, V + + def _prep_qkv_gate_packed_tensors( + self, + position_ids, + hidden_states, + past_key_value, + adapter_ids=None, + cos_cache=None, + sin_cache=None, + use_polar_compatible_rope=False, + ): + Q, gate, K, V = self._qkv_gate_packed_projection_nki(hidden_states) + + bsz, q_len, _ = hidden_states.size() + V = move_heads_front( + V, + bsz, + q_len, + self.num_key_value_heads, + self.head_dim, + layernorm=None, + ) + if cos_cache is None or sin_cache is None: + cos_cache, sin_cache = self.rotary_emb(V, position_ids) + if ( + self._should_use_qwen_qk_norm_rope_nki(q_len) + and cos_cache is not None + and sin_cache is not None + ): + Q, K = _qwen_qk_norm_partial_rope_kernel[self.logical_nc_config]( + Q, + K, + self.q_layernorm.weight.data, + self.k_layernorm.weight.data, + cos_cache, + sin_cache, + self.rms_norm_eps, + ) + else: + Q = move_heads_front( + Q, + bsz, + q_len, + self.num_heads, + self.head_dim, + layernorm=self.q_layernorm, + post_transpose_layernorm=self.post_transpose_layernorm, + ) + K = move_heads_front( + K, + bsz, + q_len, + self.num_key_value_heads, + self.head_dim, + layernorm=self.k_layernorm, + post_transpose_layernorm=self.post_transpose_layernorm, + ) + Q, K, cos_cache, sin_cache = self.apply_rotary_embedding( + Q, + K, + V, + position_ids, + cos_cache, + sin_cache, + use_polar_compatible_rope, + ) + return Q, K, V, gate, cos_cache, sin_cache, None + + def _should_use_qwen_qk_norm_rope_nki(self, q_len): + return ( + self.qwen_qk_norm_rope_nki_kernel_enabled + and q_len > 1 + and self.q_layernorm is not None + and self.k_layernorm is not None + and not self.qkv_proj_sp_enabled + ) + + def _prep_qkv_tensors_qwen_qk_norm_rope_nki( + self, + position_ids, + hidden_states, + past_key_value, + adapter_ids=None, + cos_cache=None, + sin_cache=None, + rmsnorm=None, + ): + Q, K, V, residual = self.get_qkv_proj()( + hidden_states=hidden_states, + rmsnorm=rmsnorm, + adapter_ids=adapter_ids, + residual=None, + ) + + bsz, q_len, _ = hidden_states.size() + V = move_heads_front( + V, + bsz, + q_len, + self.num_key_value_heads, + self.head_dim, + layernorm=None, + ) + if cos_cache is None or sin_cache is None: + cos_cache, sin_cache = self.rotary_emb(V, position_ids) + + Q, K = _qwen_qk_norm_partial_rope_kernel[self.logical_nc_config]( + Q, + K, + self.q_layernorm.weight.data, + self.k_layernorm.weight.data, + cos_cache, + sin_cache, + self.rms_norm_eps, + ) + return Q, K, V, cos_cache, sin_cache, residual + + def apply_rotary_embedding( + self, Q, K, V, position_ids, cos_cache, sin_cache, use_polar_compatible_rope + ): + """Partial RoPE: only apply rotary embedding to first rope_dim dimensions. + + Q shape: (B, H, S, head_dim) where head_dim=256 + cos/sin shape: (B, S, rope_dim) where rope_dim=64 (from RotaryEmbedding(dim=64)) + + Split Q/K along last dim into: + q_rope (first 64 dims) -- apply RoPE + q_pass (remaining 192 dims) -- pass through unchanged + """ + from neuronx_distributed_inference.modules.attention.utils import ( + apply_rotary_pos_emb, + ) + + if self.rotary_emb is not None: + if cos_cache is None or sin_cache is None: + cos_cache, sin_cache = self.rotary_emb(V, position_ids) + + # Split into rope and pass-through portions + Q_orig_dtype = Q.dtype + q_rope = Q[..., : self.rope_dim] # (B, H, S, 64) + q_pass = Q[..., self.rope_dim :] # (B, H, S, 192) + k_rope = K[..., : self.rope_dim] + k_pass = K[..., self.rope_dim :] + + # Apply RoPE only to the rope portion + q_rope, k_rope = apply_rotary_pos_emb(q_rope, k_rope, cos_cache, sin_cache) + + # Concatenate back (ensure bf16 is maintained) + Q = torch.cat([q_rope, q_pass], dim=-1).to(Q_orig_dtype) + K = torch.cat([k_rope, k_pass], dim=-1).to(Q_orig_dtype) + + return Q, K, cos_cache, sin_cache + + def perform_prefill(self, Q, K, V, q_len, bsz, attention_mask=None): + """Prefill path with NKI flash attention for head_dim=256.""" + head_dim = Q.shape[-1] + + # Option B: nkilib flash attention for head_dim > 128 + if _nkilib_flash_attn is not None: + q_contig = Q.contiguous() + k_contig = K.contiguous() + v_contig = V.contiguous() + scale = 1.0 / math.sqrt(head_dim) + result = _nkilib_flash_attn( + q_contig, k_contig, v_contig, scale=scale, use_causal_mask=True + ) + return result, None + + # Option A: kernel patched globally + if NKILIB_PATCH_ACTIVE: + return _flash_fwd_call(Q, K, V, use_causal_mask=True), None + + # Fallback: softmax path (use 3D tensors to avoid compiler ICE with 4D patterns) + if head_dim > 128: + # GQA: expand K/V heads to match Q heads + num_q_heads = Q.shape[1] + num_kv_heads = K.shape[1] + if num_q_heads != num_kv_heads: + kv_rep = num_q_heads // num_kv_heads + K = ( + K.unsqueeze(2) + .expand(-1, -1, kv_rep, -1, -1) + .reshape(bsz, num_q_heads, q_len, head_dim) + ) + V = ( + V.unsqueeze(2) + .expand(-1, -1, kv_rep, -1, -1) + .reshape(bsz, num_q_heads, q_len, head_dim) + ) + # Reshape to 3D (B*H, S, d) to avoid neuronx-cc codegen ICE with 4D + # attention weight tensors (NCC_INLA001: Expected 2D tensor but got 4D AP) + Q_3d = Q.reshape(bsz * num_q_heads, q_len, head_dim) + K_3d = K.reshape(bsz * num_q_heads, q_len, head_dim) + V_3d = V.reshape(bsz * num_q_heads, q_len, head_dim) + attn_weights = torch.bmm(Q_3d, K_3d.transpose(-1, -2)) / math.sqrt(head_dim) + # Build causal mask for 3D: (1, S, S) broadcast over B*H + causal_mask = torch.triu( + torch.full( + (q_len, q_len), + -65504.0, + dtype=attn_weights.dtype, + device=attn_weights.device, + ), + diagonal=1, + ).unsqueeze(0) + attn_weights = attn_weights + causal_mask + attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to( + Q.dtype + ) + attn_output = torch.bmm(attn_weights, V_3d) + # Reshape back to 4D (B, H, S, d) + return attn_output.reshape(bsz, num_q_heads, q_len, head_dim), None + + return _flash_fwd_call(Q, K, V, use_causal_mask=True), None + + def perform_qwen_chunked_prefill( + self, + Q, + K, + V, + past_key_value, + position_ids, + attention_mask=None, + kv_mgr=None, + idx=None, + active_block_table=None, + computed_context_lens=None, + scatter_index=None, + kvcache_buffer=None, + ): + """Exact chunked CTE over full-cache or selected-prefix KV. + + For model-local chunked prefill, the current chunk K/V tensors are + scattered into the full cache at absolute position_ids. For vLLM prefix + reuse, BlockKVCacheManager returns selected prefix blocks already + arranged as logical positions, so concatenate the current suffix K/V + after that logical prefix. + """ + k_cache, v_cache = past_key_value + B, q_heads, q_len, head_dim = Q.shape + kv_heads = K.shape[1] + use_segmented_prefix_cte = ( + getattr( + self.config.neuron_config, + "prefix_cte_attention_backend", + "attention_cte", + ) + == "segmented_cte" + and active_block_table is not None + and getattr(active_block_table, "ndim", 0) > 1 + ) + if use_segmented_prefix_cte: + if kv_mgr is None or idx is None or scatter_index is None: + raise ValueError( + "segmented_cte Qwen prefix prefill requires kv_mgr, idx, " + "and scatter_index so active KV can be written to block KV." + ) + updated_kv = kv_mgr.update_kv_by_layer_id( + idx=idx, + kv_per_layer=(K.to(self.torch_dtype), V.to(self.torch_dtype)), + scatter_index=scatter_index, + kvcache_buffer=kvcache_buffer, + ) + attn_output, _flash_strategy = self.perform_prefix_prefill_segmented_cte( + Q, + q_len, + B, + updated_kv, + active_block_table, + computed_context_lens, + ) + return attn_output.permute(0, 1, 3, 2).contiguous(), updated_kv + + if k_cache.shape[0] != B: + # The cache is allocated at kv_cache_batch_size, while CTE can trace a + # smaller active batch. Keep attention reshapes on the active batch. + k_cache = k_cache[:B] + v_cache = v_cache[:B] + cache_len = k_cache.shape[2] + + pos = position_ids.long() + selected_prefix_cache = cache_len < int( + getattr(self.config.neuron_config, "seq_len", cache_len) + ) + if selected_prefix_cache: + k_cache = torch.cat([k_cache, K.to(k_cache.dtype)], dim=2) + v_cache = torch.cat([v_cache, V.to(v_cache.dtype)], dim=2) + prefix_positions = torch.arange( + cache_len, + device=position_ids.device, + dtype=pos.dtype, + ).view(1, -1).expand(B, -1) + cache_positions = torch.cat([prefix_positions, pos], dim=1).view( + B, + 1, + 1, + -1, + ) + prefix_valid = torch.ones( + (B, cache_len), + device=position_ids.device, + dtype=torch.bool, + ) + if ( + attention_mask is not None + and attention_mask.ndim == 2 + and attention_mask.shape[1] == q_len + ): + active_valid = attention_mask.to(torch.bool) + else: + active_valid = torch.ones( + (B, q_len), + device=position_ids.device, + dtype=torch.bool, + ) + key_valid_mask = torch.cat([prefix_valid, active_valid], dim=1).view( + B, + 1, + 1, + -1, + ) + cache_len = k_cache.shape[2] + else: + k_index = pos[:, None, :, None].expand(B, kv_heads, q_len, head_dim) + k_cache = torch.scatter( + k_cache, + dim=2, + index=k_index, + src=K.to(k_cache.dtype), + ) + v_cache = torch.scatter( + v_cache, + dim=2, + index=k_index, + src=V.to(v_cache.dtype), + ) + cache_positions = torch.arange( + cache_len, + device=position_ids.device, + dtype=pos.dtype, + ).view(1, 1, 1, -1) + key_valid_mask = None + + prefix_attention_impl = _qwen36_prefix_attention_impl() + if prefix_attention_impl == "grouped": + attn_output = _qwen35_grouped_prefix_attention( + Q, + k_cache, + v_cache, + pos, + cache_positions, + key_valid_mask, + ) + else: + attn_output = _qwen35_expanded_prefix_attention( + Q, + k_cache, + v_cache, + pos, + cache_positions, + key_valid_mask, + ) + return attn_output, None + + def forward( + self, + hidden_states, + attention_mask=None, + position_ids=None, + past_key_value=None, + cos_cache=None, + sin_cache=None, + rmsnorm=None, + adapter_ids=None, + active_mask=None, + **kwargs, + ): + """Forward with output gate applied BEFORE o_proj. + + Override NeuronAttentionBase.forward() to insert the sigmoid gate + between the attention output and o_proj, matching the HF reference: + gate = sigmoid(gate_proj(pre_attn_hidden)) + attn_output = attn_output * gate + attn_output = o_proj(attn_output) + """ + bsz, q_len, _ = hidden_states.shape + + # Use standard 2D position_ids for prep_qkv_tensors. + rope_pos_ids = position_ids + + use_split_qkv_tkg = ( + self.qkv_tkg_nki_kernel_enabled + and past_key_value is not None + and q_len == 1 + ) + if self._should_use_qwen_qkv_gate_packed(q_len): + Q, K, V, gate, cos_cache, sin_cache, _residual = ( + self._prep_qkv_gate_packed_tensors( + rope_pos_ids, + hidden_states, + past_key_value, + adapter_ids=adapter_ids, + cos_cache=cos_cache, + sin_cache=sin_cache, + ) + ) + elif use_split_qkv_tkg: + gate = ( + self._output_gate_proj_nki(hidden_states) + if self._should_use_qwen_output_gate_nki(q_len) + else self.output_gate_proj(hidden_states) + ) + Q, K, V, cos_cache, sin_cache, _residual = ( + self._prep_split_qkv_tkg_tensors( + rope_pos_ids, + hidden_states, + past_key_value, + adapter_ids=adapter_ids, + cos_cache=cos_cache, + sin_cache=sin_cache, + ) + ) + elif self.qkv_tkg_nki_kernel_enabled: + raise ValueError( + "qkv_tkg_nki_kernel_enabled is only valid for single-token " + f"decode, got past_key_value={past_key_value is not None}, " + f"q_len={q_len}" + ) + else: + # Compute gate from input hidden states (before QKV projection). + if self._should_use_qwen_output_gate_nki(q_len): + gate = self._output_gate_proj_nki(hidden_states) + else: + gate = self.output_gate_proj(hidden_states) + + # Standard QKV prep (projections, QK norm, RoPE) + if self._should_use_qwen_qk_norm_rope_nki(q_len): + Q, K, V, cos_cache, sin_cache, _residual = ( + self._prep_qkv_tensors_qwen_qk_norm_rope_nki( + rope_pos_ids, + hidden_states, + past_key_value, + adapter_ids=adapter_ids, + cos_cache=cos_cache, + sin_cache=sin_cache, + rmsnorm=rmsnorm, + ) + ) + else: + Q, K, V, cos_cache, sin_cache, _residual = self.prep_qkv_tensors( + rope_pos_ids, + hidden_states, + past_key_value, + adapter_ids=adapter_ids, + cos_cache=cos_cache, + sin_cache=sin_cache, + rmsnorm=rmsnorm, + ) + + qwen_chunked_prefill_active = ( + past_key_value is not None + and q_len > 1 + and getattr(self.config, "use_qwen_hybrid_chunked_prefill", False) + ) + + if past_key_value is None: + # Context encoding (prefill) + attn_output, _flash_strategy = self.perform_prefill( + Q, K, V, q_len, bsz, attention_mask + ) + elif qwen_chunked_prefill_active: + attn_output, present_key_value = self.perform_qwen_chunked_prefill( + Q, + K, + V, + past_key_value, + position_ids, + attention_mask, + kv_mgr=kwargs.get("kv_mgr"), + idx=kwargs.get("idx"), + active_block_table=kwargs.get("active_block_table"), + computed_context_lens=kwargs.get("computed_context_lens"), + scatter_index=kwargs.get("scatter_index"), + kvcache_buffer=kwargs.get("kvcache_buffer"), + ) + else: + # Token generation (decode) + tkg_mask = attention_mask + if tkg_mask is not None and tkg_mask.ndim == 2: + tkg_mask = tkg_mask.unsqueeze(1).unsqueeze(2) # (B, S) -> (B, 1, 1, S) + attn_output = self.compute_for_token_gen( + Q, K, V, position_ids, past_key_value, tkg_mask, active_mask + ) + + # attn_output is (B, H, S, head_dim) -- transpose to (B, S, H*head_dim) + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.head_dim) + + o_proj = self.get_o_proj() + if self._should_use_qwen_gated_o_proj_nki(q_len): + attn_output = o_proj.forward_gated(attn_output, gate, adapter_ids=adapter_ids) + else: + # Apply sigmoid output gate BEFORE o_proj (matching HF reference) + attn_output = attn_output * torch.sigmoid(gate) + attn_output = o_proj(attn_output, adapter_ids=adapter_ids) + + # Ensure K, V are in model dtype (bf16) for KV cache update + # (prevents mixed-precision dynamic-update-slice in neuronx-cc) + K = K.to(self.torch_dtype) + V = V.to(self.torch_dtype) + if "present_key_value" not in locals() or present_key_value is None: + present_key_value = (K, V) + past_key_value = present_key_value + return attn_output, past_key_value, cos_cache, sin_cache + + +# ============================================================ +# Dense MLP (replaces MoE) +# ============================================================ + + +class Qwen35MLP(nn.Module): + """Dense SwiGLU MLP for Qwen3.5/3.6-27B. + + gate_proj: hidden_size -> intermediate_size (5120 -> 17408) + up_proj: hidden_size -> intermediate_size (5120 -> 17408) + down_proj: intermediate_size -> hidden_size (17408 -> 5120) + + output = down_proj(silu(gate_proj(x)) * up_proj(x)) + """ + + def __init__(self, config): + super().__init__() + self.gate_proj = ColumnParallelLinear( + config.hidden_size, + config.intermediate_size, + bias=False, + gather_output=False, + ) + self.up_proj = ColumnParallelLinear( + config.hidden_size, + config.intermediate_size, + bias=False, + gather_output=False, + ) + self.down_proj = RowParallelLinear( + config.intermediate_size, + config.hidden_size, + bias=False, + input_is_parallel=True, + ) + + def forward(self, hidden_states): + gate = self.gate_proj(hidden_states) + up = self.up_proj(hidden_states) + hidden_states = F.silu(gate) * up + hidden_states = self.down_proj(hidden_states) + return hidden_states + + +# ============================================================ +# MoE MLP (Qwen3.5-MoE variants — qwen3_5_moe_text model_type) +# ============================================================ + + +class Qwen35MoEBlock(nn.Module): + """Qwen3.5-MoE sparse feed-forward with a shared expert gated by a per-token + sigmoid, matching the HF reference `Qwen3_5MoeSparseMoeBlock`. + + Uses NxDI's `initialize_moe_module` for the routed experts + a plain shared + MLP (SwiGLU). The shared-expert sigmoid gate lives here — NxDI's built-in + SharedExperts sums shared_output into the routed output directly, without + a per-token gate, so we override the shared path. + """ + + def __init__(self, config): + super().__init__() + if not HAS_MOE_V2: + raise RuntimeError( + "Qwen3.5-MoE variant needs neuronx_distributed_inference.modules.moe_v2, " + "which is not present in the installed NxDI." + ) + # NxDI's MoE module (routed experts + optional built-in shared experts). + # We DISABLE the built-in shared branch (n_shared_experts=0 for the module) + # and roll our own gated shared expert below. + moe_config = copy.copy(config) # shallow copy so we can override transiently + moe_config.n_shared_experts = 0 # keep initialize_moe_module's SharedExperts off + self.moe = initialize_moe_module(config=moe_config) + + # Sigmoid-gated shared expert with its own per-token gate. Uses column- + # parallel + row-parallel linears, sharded on the intermediate dim. + self.shared_intermediate_size = getattr( + config, "shared_expert_intermediate_size", + getattr(config, "moe_intermediate_size", config.intermediate_size), + ) + self.shared_gate_proj = ColumnParallelLinear( + config.hidden_size, self.shared_intermediate_size, + bias=False, gather_output=False, + ) + self.shared_up_proj = ColumnParallelLinear( + config.hidden_size, self.shared_intermediate_size, + bias=False, gather_output=False, + ) + self.shared_down_proj = RowParallelLinear( + self.shared_intermediate_size, config.hidden_size, + bias=False, input_is_parallel=True, + ) + # Per-token sigmoid gate for the shared expert output. + # gate_output = sigmoid(shared_expert_gate(x)) * shared_mlp(x). + # This is a scalar-output linear (1 output feature), which can't be + # column-sharded, so it lives replicated on every rank. + self.shared_expert_gate = nn.Linear(config.hidden_size, 1, bias=False) + + def forward(self, hidden_states): + # NxDI MoE forward returns (output, *aux) where output has same shape + # as hidden_states. + moe_output = self.moe(hidden_states) + if isinstance(moe_output, (tuple, list)): + moe_output = moe_output[0] + + # Shared expert (SwiGLU). + gate = self.shared_gate_proj(hidden_states) + up = self.shared_up_proj(hidden_states) + shared = self.shared_down_proj(F.silu(gate) * up) + + # Per-token sigmoid gate on the shared expert output. + shared_gate = torch.sigmoid(self.shared_expert_gate(hidden_states)) + shared = shared_gate * shared + + return moe_output + shared + + +# ============================================================ +# Decoder Layer (hybrid dispatch -- DeltaNet or GQA + Dense MLP) +# ============================================================ + + +class NeuronQwen35DecoderLayer(nn.Module): + """Hybrid decoder layer: dispatches to DeltaNet or standard attention. + Uses dense MLP for all layers (no MoE). + """ + + def __init__(self, config: Qwen35InferenceConfig, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + self.layer_type = config.layer_types[layer_idx] + self.layer_idx = layer_idx + self.config = config + + # Attention (DeltaNet or standard GQA) + if self.layer_type == "linear_attention": + self.linear_attn = NeuronGatedDeltaNet(config, layer_idx) + else: + self.self_attn = NeuronQwen35Attention(config=config) + + # Dense MLP (all layers). The reusable NxDI Llama MLP kernel supports + # both CTE and TKG; keep RMSNorm separate for CTE so normalization stays + # on the conservative high-precision path before FP8 GEMM quantization. + self.mlp_kernel_enabled = bool(config.neuron_config.mlp_kernel_enabled) + self.mlp_kernel_fused_rmsnorm = ( + self.mlp_kernel_enabled + and not config.neuron_config.sequence_parallel_enabled + ) + # MoE variants use a sparse feed-forward with routed experts + a + # sigmoid-gated shared expert. Dense variants keep the plain SwiGLU MLP. + self.is_moe = getattr(config, "_is_moe", False) + if self.is_moe: + self.mlp_kernel_enabled = False + self.mlp = Qwen35MoEBlock(config) + elif self.mlp_kernel_enabled: + tensor_model_parallel_group = ( + parallel_state.get_tensor_model_parallel_group() + if parallel_state.model_parallel_is_initialized() + else None + ) + self.mlp = NeuronLlamaMLP(config, tensor_model_parallel_group) + else: + self.mlp = Qwen35MLP(config) + + self.input_layernorm = get_rmsnorm_cls()( + config.hidden_size, eps=config.rms_norm_eps + ) + self.post_attention_layernorm = get_rmsnorm_cls()( + config.hidden_size, eps=config.rms_norm_eps + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask=None, + position_ids=None, + past_key_value=None, + padding_mask=None, + cos_cache=None, + sin_cache=None, + **kwargs, + ): + residual = hidden_states + + hidden_states = ModuleMarkerStartWrapper()(hidden_states) + hidden_states = self.input_layernorm(hidden_states) + + if self.layer_type == "linear_attention": + # DeltaNet path + attn_out, dummy_kv, new_rec_state, new_conv_state = self.linear_attn( + hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + **kwargs, + ) + hidden_states = residual + attn_out + present_key_value = dummy_kv + deltanet_states = ( + None + if getattr(self.config, "use_hybrid_cache_manager", False) + else (new_rec_state, new_conv_state) + ) + else: + deltanet_states = None + # Standard attention path + hidden_states, present_key_value, cos_cache, sin_cache = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + cos_cache=cos_cache, + sin_cache=sin_cache, + **kwargs, + ) + hidden_states = residual + hidden_states + + # Dense MLP FFN + residual = hidden_states + if self.mlp_kernel_enabled: + use_fused_mlp_rmsnorm = ( + self.mlp_kernel_fused_rmsnorm + and not bool(kwargs.get("is_for_context_encoding", False)) + and hidden_states.shape[1] == 1 + ) + if use_fused_mlp_rmsnorm: + mlp_fused_rmsnorm = self.post_attention_layernorm + else: + hidden_states = self.post_attention_layernorm(hidden_states) + mlp_fused_rmsnorm = None + hidden_states, _ = self.mlp(hidden_states, rmsnorm=mlp_fused_rmsnorm) + else: + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + hidden_states = ModuleMarkerEndWrapper()(hidden_states) + outputs = ( + hidden_states, + present_key_value, + cos_cache, + sin_cache, + None, + deltanet_states, + ) + return outputs + + +# ============================================================ +# Hybrid Cache Manager (opt-in) +# ============================================================ + + +class HybridDeltaNetCacheManager(KVCacheManager): + """Opt-in local/static cache manager for Qwen hybrid dense models. + + This manager stores DeltaNet recurrent/conv state by batch row and delegates + full-attention layers to the legacy KV manager. It is intentionally not a + production vLLM APC manager: block ownership, prefix hashes, refcounts, + eviction, continuous batching, and tenant isolation must remain in the + vLLM/NxDI block-cache lifecycle. + """ + + def __init__(self, config: Qwen35InferenceConfig, num_kv_head, **kwargs): + self.layer_types = list(config.layer_types) + self._validate_hybrid_config(config) + super().__init__(config, num_kv_head=num_kv_head, **kwargs) + + dtype = ( + config.neuron_config.attention_dtype + if config.neuron_config.attention_dtype is not None + else config.neuron_config.torch_dtype + ) + cache_dtype = getattr(self, "cache_dtype", dtype) + recurrent_cache_dtype = _torch_dtype_from_hybrid_cache_dtype( + config.hybrid_recurrent_cache_dtype + ) + conv_cache_dtype = _torch_dtype_from_hybrid_cache_dtype( + config.hybrid_conv_cache_dtype + ) + max_batch_size = ( + config.neuron_config.kv_cache_batch_size + + config.neuron_config.kv_cache_padding_size + ) + tp_degree = config.neuron_config.tp_degree + if config.linear_num_value_heads % tp_degree != 0: + raise ValueError( + f"linear_num_value_heads={config.linear_num_value_heads} must be divisible " + f"by tp_degree={tp_degree}" + ) + if config.linear_num_key_heads % tp_degree != 0: + raise ValueError( + f"linear_num_key_heads={config.linear_num_key_heads} must be divisible " + f"by tp_degree={tp_degree}" + ) + local_num_value_heads = config.linear_num_value_heads // tp_degree + local_num_key_heads = config.linear_num_key_heads // tp_degree + recurrent_shape = [ + max_batch_size, + local_num_value_heads, + config.linear_key_head_dim, + config.linear_value_head_dim, + ] + conv_dim = ( + 2 * local_num_key_heads * config.linear_key_head_dim + + local_num_value_heads * config.linear_value_head_dim + ) + conv_shape = [ + max_batch_size, + conv_dim, + config.linear_conv_kernel_dim - 1, + ] + + params = [] + for layer_idx, layer_type in enumerate(self.layer_types): + if layer_type == "linear_attention": + params.append( + nn.Parameter( + torch.zeros(recurrent_shape, dtype=recurrent_cache_dtype), + requires_grad=False, + ) + ) + params.append( + nn.Parameter( + torch.zeros(conv_shape, dtype=conv_cache_dtype), + requires_grad=False, + ) + ) + else: + k_shape = self.k_shapes[layer_idx] if hasattr(self, "k_shapes") else self.k_shape + v_shape = self.v_shapes[layer_idx] if hasattr(self, "v_shapes") else self.v_shape + params.append( + nn.Parameter(torch.zeros(k_shape, dtype=cache_dtype), requires_grad=False) + ) + params.append( + nn.Parameter(torch.zeros(v_shape, dtype=cache_dtype), requires_grad=False) + ) + + self.past_key_values = nn.ParameterList(params) + + @staticmethod + def _validate_hybrid_config(config: Qwen35InferenceConfig): + nc = config.neuron_config + unsupported = [] + if nc.is_block_kv_layout: + unsupported.append("block KV layout") + if getattr(nc, "kv_quant_config", None) is not None or getattr(nc, "kv_cache_quant", False): + unsupported.append("KV cache quantization") + if nc.enable_fused_speculation or nc.speculation_length > 0 or nc.is_medusa: + unsupported.append("speculative decoding") + if getattr(nc, "enable_eagle_speculation", False) or getattr(nc, "is_eagle_draft", False): + unsupported.append("EAGLE speculation") + if nc.flash_decoding_enabled: + unsupported.append("flash decoding") + if nc.attention_dp_degree > 1: + unsupported.append("attention data parallelism") + if nc.kv_cache_tiling: + unsupported.append("KV cache tiling") + if nc.padding_side != "right": + unsupported.append("left padding") + if nc.is_continuous_batching: + unsupported.append("continuous batching") + if unsupported: + raise ValueError( + "HybridDeltaNetCacheManager v1 does not support: " + + ", ".join(unsupported) + ) + + def _is_deltanet_layer(self, idx: int) -> bool: + return self.layer_types[idx] == "linear_attention" + + def get_seq_length(self, past_key_values=None): + for idx, layer_type in enumerate(self.layer_types): + if layer_type != "linear_attention": + if past_key_values is None: + _, v_cache = self._fetch_cache(idx) + elif len(past_key_values) == len(self.past_key_values): + v_cache = past_key_values[2 * idx + 1] + else: + v_cache = past_key_values[idx][1] + return v_cache.shape[2] + return 0 + + def get_deltanet_state_by_layer_id(self, idx, kvcache_buffer=None, seq_ids=None): + recurrent_state, conv_state = self._fetch_cache(idx, kvcache_buffer) + if seq_ids is not None: + cache_idx = self.get_cache_update_index_for_seq_ids(seq_ids) + recurrent_state = torch.index_select(recurrent_state, dim=0, index=cache_idx) + conv_state = torch.index_select(conv_state, dim=0, index=cache_idx) + elif self.kv_cache_padding_size > 0: + recurrent_state = recurrent_state[: -self.kv_cache_padding_size] + conv_state = conv_state[: -self.kv_cache_padding_size] + return recurrent_state, conv_state + + def get_cache( + self, + seq_len: int, + skip_slice=False, + kvcache_buffer=None, + seq_ids=None, + windowed_context_encoding_window_idx=-1, + **kwargs, + ): + past_key_values = [] + for idx in range(len(self.past_key_values) // 2): + if self._is_deltanet_layer(idx): + past_key_values.append( + list(self.get_deltanet_state_by_layer_id(idx, kvcache_buffer, seq_ids)) + ) + else: + past_key_values.append( + list( + self.get_kv_by_layer_id( + idx=idx, + skip_slice=skip_slice, + seq_len=seq_len, + kvcache_buffer=kvcache_buffer, + seq_ids=seq_ids, + windowed_context_encoding_window_idx=windowed_context_encoding_window_idx, + **kwargs, + ) + ) + ) + return past_key_values + + def update_cache( + self, + is_for_context_encoding: bool, + seq_ids: torch.Tensor, + position_ids: torch.Tensor, + new_key_values: List[torch.Tensor], + seq_len: int, + scatter_index=None, + kv_active_mask=None, + kvcache_buffer=None, + windowed_context_encoding_window_idx: int = -1, + **kwargs, + ): + updated_cache = [] + for idx, kv_per_layer in enumerate(new_key_values): + if self._is_deltanet_layer(idx): + recurrent_state, conv_state = self.update_deltanet_state_by_layer_id( + idx=idx, + seq_ids=seq_ids, + state_per_layer=kv_per_layer, + kvcache_buffer=kvcache_buffer, + ) + elif kwargs.get("qwen_chunked_prefill_update", False): + recurrent_state, conv_state = self.update_qwen_chunked_kv_by_layer_id( + idx=idx, + seq_ids=seq_ids, + position_ids=position_ids, + kv_per_layer=kv_per_layer, + kvcache_buffer=kvcache_buffer, + valid_mask=kwargs.get("qwen_chunked_valid_mask", None), + ) + else: + recurrent_state, conv_state = self.update_kv_by_layer_id( + idx=idx, + is_for_context_encoding=is_for_context_encoding, + seq_ids=seq_ids, + position_ids=position_ids, + kv_per_layer=kv_per_layer, + seq_len=seq_len, + scatter_index=scatter_index, + kv_active_mask=kv_active_mask, + kvcache_buffer=kvcache_buffer, + windowed_context_encoding_window_idx=windowed_context_encoding_window_idx, + **kwargs, + ) + updated_cache.append(recurrent_state) + updated_cache.append(conv_state) + return updated_cache + + def update_qwen_chunked_kv_by_layer_id( + self, + idx: int, + seq_ids: torch.Tensor, + position_ids: torch.Tensor, + kv_per_layer: Tuple[torch.Tensor, torch.Tensor], + kvcache_buffer=None, + valid_mask=None, + ): + latest_k, latest_v = kv_per_layer + k_cache, v_cache = self._fetch_cache(idx, kvcache_buffer) + latest_k = latest_k.to(k_cache.dtype) + latest_v = latest_v.to(v_cache.dtype) + + if seq_ids is not None: + cache_idx = self.get_cache_update_index_for_seq_ids(seq_ids) + selected_k = torch.index_select(k_cache, dim=0, index=cache_idx) + selected_v = torch.index_select(v_cache, dim=0, index=cache_idx) + else: + cache_idx = None + selected_k = k_cache[: latest_k.shape[0]] + selected_v = v_cache[: latest_v.shape[0]] + + pos = position_ids.long() + k_index = pos[:, None, :, None].expand_as(latest_k) + v_index = pos[:, None, :, None].expand_as(latest_v) + + if valid_mask is not None: + valid = valid_mask.to(torch.bool)[:, None, :, None] + old_k = torch.gather(selected_k, dim=2, index=k_index) + old_v = torch.gather(selected_v, dim=2, index=v_index) + latest_k = torch.where(valid, latest_k, old_k) + latest_v = torch.where(valid, latest_v, old_v) + + updated_k = torch.scatter(selected_k, dim=2, index=k_index, src=latest_k) + updated_v = torch.scatter(selected_v, dim=2, index=v_index, src=latest_v) + + if cache_idx is not None: + k_row_index = cache_idx.view(-1, 1, 1, 1).expand_as(updated_k) + v_row_index = cache_idx.view(-1, 1, 1, 1).expand_as(updated_v) + k_cache = torch.scatter(k_cache, dim=0, index=k_row_index, src=updated_k) + v_cache = torch.scatter(v_cache, dim=0, index=v_row_index, src=updated_v) + return k_cache, v_cache + + if updated_k.shape[0] == k_cache.shape[0]: + return updated_k + k_cache * 0, updated_v + v_cache * 0 + + pad_rows = k_cache.shape[0] - updated_k.shape[0] + if pad_rows > 0: + updated_k = torch.cat([updated_k, k_cache[updated_k.shape[0] :] * 0], dim=0) + updated_v = torch.cat([updated_v, v_cache[updated_v.shape[0] :] * 0], dim=0) + return updated_k + k_cache * 0, updated_v + v_cache * 0 + + def update_deltanet_state_by_layer_id( + self, + idx: int, + seq_ids: torch.Tensor, + state_per_layer: Tuple[torch.Tensor, torch.Tensor], + kvcache_buffer=None, + ): + latest_recurrent, latest_conv = state_per_layer + recurrent_cache, conv_cache = self._fetch_cache(idx, kvcache_buffer) + latest_recurrent = latest_recurrent.to(recurrent_cache.dtype) + latest_conv = latest_conv.to(conv_cache.dtype) + + if latest_recurrent.shape[0] == recurrent_cache.shape[0] and seq_ids is None: + return ( + latest_recurrent + recurrent_cache * 0, + latest_conv + conv_cache * 0, + ) + + if seq_ids is not None: + cache_idx = self.get_cache_update_index_for_seq_ids(seq_ids) + recurrent_index = cache_idx.view(-1, 1, 1, 1).expand_as(latest_recurrent) + conv_index = cache_idx.view(-1, 1, 1).expand_as(latest_conv) + recurrent_cache = torch.scatter( + input=recurrent_cache, + dim=0, + index=recurrent_index, + src=latest_recurrent, + ) + conv_cache = torch.scatter( + input=conv_cache, + dim=0, + index=conv_index, + src=latest_conv, + ) + return recurrent_cache, conv_cache + + pad_size = recurrent_cache.shape[0] - latest_recurrent.shape[0] + if pad_size > 0: + latest_recurrent = torch.cat( + [latest_recurrent, recurrent_cache[latest_recurrent.shape[0] :] * 0], + dim=0, + ) + latest_conv = torch.cat( + [latest_conv, conv_cache[latest_conv.shape[0] :] * 0], + dim=0, + ) + return latest_recurrent + recurrent_cache * 0, latest_conv + conv_cache * 0 + + +class QwenHybridBlockKVCacheManager(BlockKVCacheManager): + """Block KV manager that allocates real KV only for full-attention layers.""" + + _LINEAR_PLACEHOLDER_SHAPE = (1, 1, 1, 1) + + def __init__(self, config: Qwen35InferenceConfig, num_kv_head, **kwargs): + self.layer_types = list(config.layer_types) + super().__init__(config, num_kv_head=num_kv_head, **kwargs) + + params = [] + for layer_type in self.layer_types: + if layer_type == "full_attention": + params.append( + nn.Parameter( + torch.zeros(self.k_shape, dtype=self.cache_dtype), + requires_grad=False, + ) + ) + params.append( + nn.Parameter( + torch.zeros(self.v_shape, dtype=self.cache_dtype), + requires_grad=False, + ) + ) + else: + params.append( + nn.Parameter( + torch.zeros( + self._LINEAR_PLACEHOLDER_SHAPE, + dtype=self.cache_dtype, + ), + requires_grad=False, + ) + ) + params.append( + nn.Parameter( + torch.zeros( + self._LINEAR_PLACEHOLDER_SHAPE, + dtype=self.cache_dtype, + ), + requires_grad=False, + ) + ) + self.past_key_values = nn.ParameterList(params) + + def _is_attention_layer(self, idx: int) -> bool: + return self.layer_types[idx] == "full_attention" + + def get_seq_length(self, past_key_values=None): + for idx, layer_type in enumerate(self.layer_types): + if layer_type == "full_attention": + if past_key_values is None: + _, v_cache = self._fetch_cache(idx) + elif len(past_key_values) == len(self.past_key_values): + v_cache = past_key_values[2 * idx + 1] + else: + v_cache = past_key_values[idx][1] + if v_cache.ndim >= 4 and v_cache.shape[1] == self.pa_block_size: + return self.pa_num_blocks * self.pa_block_size + return v_cache.shape[2] + return 0 + + def get_cache(self, active_block_table=None, kvcache_buffer=None, **kwargs): + past_key_values = [] + use_segmented_prefix_cte = ( + kwargs.get("is_for_context_encoding", False) + and getattr( + self.neuron_config, + "prefix_cte_attention_backend", + "attention_cte", + ) + == "segmented_cte" + and active_block_table is not None + and getattr(active_block_table, "ndim", 0) > 1 + ) + for idx in range(len(self.past_key_values) // 2): + if self._is_attention_layer(idx): + if use_segmented_prefix_cte: + k_cache, v_cache = self.get_raw_kv_by_layer_id( + idx, + kvcache_buffer=kvcache_buffer, + ) + else: + k_cache, v_cache = self.get_kv_by_layer_id( + idx, + active_block_table, + kvcache_buffer=kvcache_buffer, + **kwargs, + ) + else: + k_cache, v_cache = self._fetch_cache( + idx, + kvcache_buffer=kvcache_buffer, + ) + past_key_values.append([k_cache, v_cache]) + return past_key_values + + def _is_raw_block_kv_pair(self, kv_per_layer: List[torch.Tensor]) -> bool: + if len(kv_per_layer) != 2: + return False + k_cache, v_cache = kv_per_layer + return ( + k_cache.ndim == 4 + and v_cache.ndim == 4 + and k_cache.shape[0] == self.pa_num_blocks + self._NUM_EXTRA_RESERVED_BLOCK + and v_cache.shape[0] == self.pa_num_blocks + self._NUM_EXTRA_RESERVED_BLOCK + and k_cache.shape[1] == self.pa_block_size + and v_cache.shape[1] == self.pa_block_size + ) + + def update_cache( + self, + new_key_values: List[torch.Tensor], + scatter_index=None, + kvcache_buffer=None, + **kwargs, + ): + updated_kv_cache = [] + for idx, kv_per_layer in enumerate(new_key_values): + if self._is_attention_layer(idx) and self._is_raw_block_kv_pair( + kv_per_layer + ): + k_cache, v_cache = kv_per_layer + elif self._is_attention_layer(idx): + k_cache, v_cache = self.update_kv_by_layer_id( + idx=idx, + kv_per_layer=kv_per_layer, + scatter_index=scatter_index, + kvcache_buffer=kvcache_buffer, + ) + else: + k_cache, v_cache = self._fetch_cache( + idx, + kvcache_buffer=kvcache_buffer, + ) + k_cache = k_cache * 1 + v_cache = v_cache * 1 + updated_kv_cache.append(k_cache) + updated_kv_cache.append(v_cache) + return updated_kv_cache + + +class HybridGDNCheckpointCache(nn.Module): + """Bounded device-side GDN prefix checkpoint bank. + + Metadata owns prefix hashes, refcounts, and eviction. This module only owns + recurrent/conv tensors addressed by checkpoint slot IDs supplied by the + scheduler/request-prep path. + """ + + def __init__(self, config: Qwen35InferenceConfig): + super().__init__() + self.gdn_layer_ids = tuple( + idx + for idx, layer_type in enumerate(config.layer_types) + if layer_type == "linear_attention" + ) + if not self.gdn_layer_ids: + raise ValueError("HybridGDNCheckpointCache requires GDN layers") + self.layer_to_bank_index = { + layer_id: bank_idx for bank_idx, layer_id in enumerate(self.gdn_layer_ids) + } + self.num_checkpoint_slots = int(config.max_gdn_checkpoint_slots) + if self.num_checkpoint_slots <= 0: + raise ValueError("max_gdn_checkpoint_slots must be positive") + + tp_degree = config.neuron_config.tp_degree + if config.linear_num_value_heads % tp_degree != 0: + raise ValueError("linear_num_value_heads must be divisible by tp_degree") + if config.linear_num_key_heads % tp_degree != 0: + raise ValueError("linear_num_key_heads must be divisible by tp_degree") + + self.local_num_value_heads = config.linear_num_value_heads // tp_degree + self.local_num_key_heads = config.linear_num_key_heads // tp_degree + self.key_dim = config.linear_key_head_dim + self.value_dim = config.linear_value_head_dim + self.conv_dim = ( + 2 * self.local_num_key_heads * config.linear_key_head_dim + + self.local_num_value_heads * config.linear_value_head_dim + ) + self.conv_state_len = config.linear_conv_kernel_dim - 1 + self.recurrent_dtype = _torch_dtype_from_hybrid_cache_dtype( + config.hybrid_recurrent_cache_dtype + ) + self.conv_dtype = _torch_dtype_from_hybrid_cache_dtype( + config.hybrid_conv_cache_dtype + ) + + recurrent_shape = ( + self.num_checkpoint_slots, + self.local_num_value_heads, + self.key_dim, + self.value_dim, + ) + conv_shape = ( + self.num_checkpoint_slots, + self.conv_dim, + self.conv_state_len, + ) + self.recurrent_slots = nn.ParameterList( + [ + nn.Parameter( + torch.zeros(recurrent_shape, dtype=self.recurrent_dtype), + requires_grad=False, + ) + for _ in self.gdn_layer_ids + ] + ) + self.conv_slots = nn.ParameterList( + [ + nn.Parameter( + torch.zeros(conv_shape, dtype=self.conv_dtype), + requires_grad=False, + ) + for _ in self.gdn_layer_ids + ] + ) + + @property + def checkpoint_params(self): + params = [] + for recurrent_slot, conv_slot in zip(self.recurrent_slots, self.conv_slots): + params.append(recurrent_slot) + params.append(conv_slot) + return params + + def bytes_per_checkpoint_per_rank(self) -> int: + recurrent_numel = ( + len(self.gdn_layer_ids) + * self.local_num_value_heads + * self.key_dim + * self.value_dim + ) + conv_numel = len(self.gdn_layer_ids) * self.conv_dim * self.conv_state_len + recurrent_bytes = 4 if self.recurrent_dtype == torch.float32 else 2 + conv_bytes = 4 if self.conv_dtype == torch.float32 else 2 + return recurrent_numel * recurrent_bytes + conv_numel * conv_bytes + + def _safe_slot_ids( + self, + slot_ids: torch.Tensor, + batch_size: int | None = None, + ) -> torch.Tensor: + slot_ids = slot_ids.reshape(-1).long().clamp( + min=0, + max=self.num_checkpoint_slots - 1, + ) + if batch_size is None: + return slot_ids + if slot_ids.shape[0] >= batch_size: + return slot_ids[:batch_size] + pad = torch.zeros( + (batch_size - slot_ids.shape[0],), + dtype=slot_ids.dtype, + device=slot_ids.device, + ) + return torch.cat([slot_ids, pad], dim=0) + + @staticmethod + def _safe_bool_vector( + mask: torch.Tensor, + batch_size: int, + device: torch.device, + ) -> torch.Tensor: + mask = mask.reshape(-1).to(device=device, dtype=torch.bool) + if mask.shape[0] >= batch_size: + return mask[:batch_size] + pad = torch.zeros( + (batch_size - mask.shape[0],), + dtype=torch.bool, + device=device, + ) + return torch.cat([mask, pad], dim=0) + + @staticmethod + def _active_rows( + state: torch.Tensor, + seq_ids: torch.Tensor | None, + batch_size: int, + ) -> torch.Tensor: + if seq_ids is not None and hasattr(seq_ids, "numel") and seq_ids.numel() > 0: + safe_seq_ids = seq_ids.reshape(-1)[:batch_size].to( + device=state.device, + dtype=torch.long, + ) + safe_seq_ids = safe_seq_ids.clamp(min=0, max=int(state.shape[0]) - 1) + return torch.index_select(state, 0, safe_seq_ids) + return state[:batch_size] + + def restore_to_active_rows( + self, + *, + layers: nn.ModuleList, + seq_ids: torch.Tensor | None, + checkpoint_slot_ids: torch.Tensor | None, + restore_mask: torch.Tensor | None, + zero_inactive: bool = False, + ) -> dict[int, tuple[torch.Tensor, torch.Tensor]] | None: + if checkpoint_slot_ids is None or restore_mask is None: + return None + batch_size = max( + int(checkpoint_slot_ids.reshape(-1).shape[0]), + int(restore_mask.reshape(-1).shape[0]), + ) + if batch_size <= 0: + return None + slot_ids = self._safe_slot_ids(checkpoint_slot_ids, batch_size) + restore_mask = self._safe_bool_vector( + restore_mask, + batch_size, + slot_ids.device, + ) + slot_ids = torch.where(restore_mask, slot_ids, torch.zeros_like(slot_ids)) + rec_mask = restore_mask.view(batch_size, 1, 1, 1) + conv_mask = restore_mask.view(batch_size, 1, 1) + + restored = {} + for bank_idx, layer_id in enumerate(self.gdn_layer_ids): + linear_attn = layers[layer_id].linear_attn + active_recurrent = self._active_rows( + linear_attn.recurrent_state_buffer, seq_ids, batch_size + ) + active_conv = self._active_rows( + linear_attn.conv_state_buffer, seq_ids, batch_size + ) + if zero_inactive: + inactive_recurrent = torch.zeros_like(active_recurrent) + inactive_conv = torch.zeros_like(active_conv) + else: + inactive_recurrent = active_recurrent + inactive_conv = active_conv + slot_recurrent = torch.index_select( + self.recurrent_slots[bank_idx], 0, slot_ids + ).to(active_recurrent.dtype) + slot_conv = torch.index_select(self.conv_slots[bank_idx], 0, slot_ids).to( + active_conv.dtype + ) + _debug_qwen36_hybrid_gdn_state( + "restore_slot_recurrent", + slot_recurrent, + layer_id=layer_id, + bank_idx=bank_idx, + slot_ids=slot_ids, + mask=restore_mask, + seq_ids=seq_ids, + ) + _debug_qwen36_hybrid_gdn_state( + "restore_slot_conv", + slot_conv, + layer_id=layer_id, + bank_idx=bank_idx, + slot_ids=slot_ids, + mask=restore_mask, + seq_ids=seq_ids, + ) + restored_recurrent = torch.where( + rec_mask, slot_recurrent, inactive_recurrent + ) + restored_conv = torch.where(conv_mask, slot_conv, inactive_conv) + _debug_qwen36_hybrid_gdn_state( + "restore_active_recurrent", + restored_recurrent, + layer_id=layer_id, + bank_idx=bank_idx, + slot_ids=slot_ids, + mask=restore_mask, + seq_ids=seq_ids, + ) + _debug_qwen36_hybrid_gdn_state( + "restore_active_conv", + restored_conv, + layer_id=layer_id, + bank_idx=bank_idx, + slot_ids=slot_ids, + mask=restore_mask, + seq_ids=seq_ids, + ) + restored[layer_id] = (restored_recurrent, restored_conv) + return restored + + def commit_from_active_rows( + self, + *, + layer_state_pairs: list[tuple[int, torch.Tensor, torch.Tensor]], + seq_ids: torch.Tensor | None, + checkpoint_slot_ids: torch.Tensor | None, + commit_mask: torch.Tensor | None, + ) -> list[torch.Tensor]: + if checkpoint_slot_ids is None or commit_mask is None: + return self.identity_outputs() + batch_size = max( + int(checkpoint_slot_ids.reshape(-1).shape[0]), + int(commit_mask.reshape(-1).shape[0]), + ) + if batch_size <= 0: + return self.identity_outputs() + slot_ids = self._safe_slot_ids(checkpoint_slot_ids, batch_size) + commit_mask = self._safe_bool_vector( + commit_mask, + batch_size, + slot_ids.device, + ) + slot_ids = torch.where(commit_mask, slot_ids, torch.zeros_like(slot_ids)) + rec_mask = commit_mask.view(batch_size, 1, 1, 1) + conv_mask = commit_mask.view(batch_size, 1, 1) + + state_by_layer = { + layer_id: (recurrent_state, conv_state) + for layer_id, recurrent_state, conv_state in layer_state_pairs + } + + def _commit_rows(slots, rows, row_mask): + output = slots * 1 + slot_axis = torch.arange( + slots.shape[0], dtype=slot_ids.dtype, device=slot_ids.device + ) + broadcast_shape = (slots.shape[0],) + (1,) * (slots.ndim - 1) + for row_idx in range(batch_size): + write_mask = torch.logical_and( + row_mask[row_idx], + slot_axis == slot_ids[row_idx], + ).view(broadcast_shape) + row_value = rows[row_idx : row_idx + 1].expand_as(output) + output = torch.where(write_mask, row_value, output) + return output + + outputs = [] + for bank_idx, layer_id in enumerate(self.gdn_layer_ids): + recurrent_slots = self.recurrent_slots[bank_idx] + conv_slots = self.conv_slots[bank_idx] + if layer_id not in state_by_layer: + outputs.append(recurrent_slots * 1) + outputs.append(conv_slots * 1) + continue + + recurrent_state, conv_state = state_by_layer[layer_id] + recurrent_rows = self._active_rows(recurrent_state, seq_ids, batch_size).to( + recurrent_slots.dtype + ) + conv_rows = self._active_rows(conv_state, seq_ids, batch_size).to( + conv_slots.dtype + ) + _debug_qwen36_hybrid_gdn_state( + "commit_input_recurrent", + recurrent_rows, + layer_id=layer_id, + bank_idx=bank_idx, + slot_ids=slot_ids, + mask=commit_mask, + seq_ids=seq_ids, + ) + _debug_qwen36_hybrid_gdn_state( + "commit_input_conv", + conv_rows, + layer_id=layer_id, + bank_idx=bank_idx, + slot_ids=slot_ids, + mask=commit_mask, + seq_ids=seq_ids, + ) + + committed_recurrent = _commit_rows( + recurrent_slots, recurrent_rows, commit_mask + ) + committed_conv = _commit_rows(conv_slots, conv_rows, commit_mask) + committed_recurrent_rows = torch.index_select( + committed_recurrent, 0, slot_ids + ) + committed_conv_rows = torch.index_select(committed_conv, 0, slot_ids) + _debug_qwen36_hybrid_gdn_state( + "commit_slot_recurrent", + committed_recurrent_rows, + layer_id=layer_id, + bank_idx=bank_idx, + slot_ids=slot_ids, + mask=commit_mask, + seq_ids=seq_ids, + ) + _debug_qwen36_hybrid_gdn_state( + "commit_slot_conv", + committed_conv_rows, + layer_id=layer_id, + bank_idx=bank_idx, + slot_ids=slot_ids, + mask=commit_mask, + seq_ids=seq_ids, + ) + + outputs.append(committed_recurrent) + outputs.append(committed_conv) + return outputs + + def identity_outputs(self) -> list[torch.Tensor]: + return [param * 1 for param in self.checkpoint_params] + + +# ============================================================ +# Model +# ============================================================ + + +def _effective_lm_head_pad_size(lm_head, logits, config): + pad_size = getattr(lm_head, "pad_size", None) + if not pad_size: + return pad_size + + if getattr(lm_head, "gather_output", False): + vocab_size = getattr(config, "vocab_size", None) + if vocab_size is not None: + return max(int(logits.shape[-1]) - int(vocab_size), 0) + + return pad_size + + +def _debug_tensor_minmax(tensor): + if tensor is None or not hasattr(tensor, "numel") or tensor.numel() == 0: + return "empty" + flat = tensor.reshape(-1) + return f"{int(flat.min().item())}:{int(flat.max().item())}" + + +def _debug_tensor_values(tensor, limit=8): + if tensor is None or not hasattr(tensor, "numel") or tensor.numel() == 0: + return [] + return tensor.reshape(-1)[:limit].tolist() + + +def _debug_tensor_shape(tensor): + if tensor is None or not hasattr(tensor, "shape"): + return None + return tuple(tensor.shape) + + +def _normalize_qwen36_slot_mapping(slot_mapping, batch_size: int, active_tokens: int): + if ( + slot_mapping is None + or not hasattr(slot_mapping, "numel") + or slot_mapping.numel() == 0 + or not hasattr(slot_mapping, "ndim") + ): + return slot_mapping + if slot_mapping.ndim != 1: + return slot_mapping + + batch_size = int(batch_size) + active_tokens = int(active_tokens) + total_slots = int(slot_mapping.numel()) + if batch_size > 0 and active_tokens > 0 and total_slots == batch_size * active_tokens: + return slot_mapping.reshape(batch_size, active_tokens) + if batch_size == 1: + return slot_mapping.reshape(1, total_slots) + return slot_mapping + + +def _use_legacy_tkg_args() -> bool: + return os.environ.get("QWEN36_TKG_LEGACY_ARGS") == "1" + + +def _qwen36_config_flag(config, neuron_config, name: str, default: bool = False) -> bool: + for owner in (config, neuron_config, getattr(config, "neuron_config", None)): + value = getattr(owner, name, None) + if value is not None: + return bool(value) + return bool(default) + + +def _use_expanded_hybrid_args_for_tag(config, tag: str) -> bool: + if not _qwen36_config_flag(config, None, "use_hybrid_apc_manager"): + return False + # The legacy ABI experiment intentionally keeps both traced stages on the + # older prefix-cache contract. Neuron prunes the extra CTE hybrid metadata + # inputs from the serialized trace, so runtime must not send them either. + if _use_legacy_tkg_args(): + return False + if tag == CONTEXT_ENCODING_MODEL_TAG: + return True + if tag == TOKEN_GENERATION_MODEL_TAG: + return True + return False + + +def _qwen36_shape_entry_arg_count(entry) -> int | None: + if isinstance(entry, str): + try: + entry = json.loads(entry) + except Exception: + return None + if isinstance(entry, (list, tuple)): + return len(entry) + return None + + +def _qwen36_compiled_arg_count(model_wrapper) -> int | None: + counts = [] + for owner in ( + model_wrapper, + getattr(model_wrapper, "model", None), + getattr(getattr(model_wrapper, "model", None), "nxd_model", None), + ): + shape_map = getattr(owner, "input_shape_map", None) + keys = getattr(shape_map, "keys", None) + if not callable(keys): + continue + try: + iterable = keys() + except Exception: + continue + for entry in iterable: + count = _qwen36_shape_entry_arg_count(entry) + if count is not None: + counts.append(count) + return max(counts) if counts else None + + +def _use_expanded_hybrid_args_for_wrapper(model_wrapper, tag: str) -> bool: + compiled_arg_count = _qwen36_compiled_arg_count(model_wrapper) + if compiled_arg_count is not None: + return compiled_arg_count >= 29 + return _use_expanded_hybrid_args_for_tag(model_wrapper.config, tag) + + +def _qwen36_expected_arg_count(config, tag: str) -> int: + return 29 if _use_expanded_hybrid_args_for_tag(config, tag) else 24 + + +def _assert_qwen36_arg_count(stage: str, args, expected: int) -> None: + actual = len(args) + if actual != expected: + raise RuntimeError( + f"Qwen3.6 {stage} argument contract mismatch: " + f"expected {expected} tensors, got {actual}" + ) + + +_QWEN36_PREFIX_ARG_NAMES = ( + "input_ids", + "attention_mask", + "position_ids", + "seq_ids", + "sampling_params", + "prev_hidden", + "adapter_ids", + "accepted_indices", + "current_length", + "medusa_mask", + "scatter_index", + "slot_mapping", + "block_table", + "num_queries", + "computed_context_lens", + "tile_q_indices", + "tile_block_tables", + "tile_masks", + "inputs_embeds", + "kv_cache", + "active_mask", +) +_QWEN36_MROPE_VISION_ARG_NAMES = ( + "rotary_position_ids", + "vision_embeddings", + "vision_mask", +) +_QWEN36_HYBRID_APC_ARG_NAMES = ( + "hybrid_restore_slot_ids", + "hybrid_restore_mask", + "hybrid_restore_prefix_lens", + "hybrid_commit_slot_ids", + "hybrid_commit_mask", +) + + +def _empty_qwen36_arg(): + return torch.empty(0) + + +def _qwen36_arg_names(config, tag: str): + names = list(_QWEN36_PREFIX_ARG_NAMES + _QWEN36_MROPE_VISION_ARG_NAMES) + if _use_expanded_hybrid_args_for_tag(config, tag): + names.extend(_QWEN36_HYBRID_APC_ARG_NAMES) + return names + + +def _normalize_qwen36_prefix_args(prefix_args): + args = list(prefix_args) + if len(args) > len(_QWEN36_PREFIX_ARG_NAMES): + raise RuntimeError( + "Qwen3.6 prefix argument contract mismatch: " + f"expected at most {len(_QWEN36_PREFIX_ARG_NAMES)} base tensors, " + f"got {len(args)}" + ) + while len(args) < len(_QWEN36_PREFIX_ARG_NAMES): + args.append(_empty_qwen36_arg()) + return args + + +def _normalize_qwen36_hybrid_args(hybrid_args, batch_size): + args = list(hybrid_args or ()) + while len(args) < len(_QWEN36_HYBRID_APC_ARG_NAMES): + args.append(torch.zeros((batch_size,), dtype=torch.int32)) + if len(args) > len(_QWEN36_HYBRID_APC_ARG_NAMES): + raise RuntimeError( + "Qwen3.6 Hybrid APC argument contract mismatch: " + f"expected {len(_QWEN36_HYBRID_APC_ARG_NAMES)} tensors, got {len(args)}" + ) + return args + + +def _build_qwen36_stage_args( + config, + tag: str, + prefix_args, + mrope_position_ids, + vision_embeddings, + vision_mask, + hybrid_args=None, +): + args = _normalize_qwen36_prefix_args(prefix_args) + args.extend([mrope_position_ids, vision_embeddings, vision_mask]) + if _use_expanded_hybrid_args_for_tag(config, tag): + batch_size = args[0].shape[0] + args.extend(_normalize_qwen36_hybrid_args(hybrid_args, batch_size)) + _assert_qwen36_arg_count(tag, args, _qwen36_expected_arg_count(config, tag)) + return args + + +def build_cte_args( + config, + prefix_args, + mrope_position_ids, + vision_embeddings, + vision_mask, + hybrid_args=None, +): + return _build_qwen36_stage_args( + config, + CONTEXT_ENCODING_MODEL_TAG, + prefix_args, + mrope_position_ids, + vision_embeddings, + vision_mask, + hybrid_args=hybrid_args, + ) + + +def build_tkg_args( + config, + prefix_args, + mrope_position_ids, + vision_embeddings, + vision_mask, + hybrid_args=None, +): + return _build_qwen36_stage_args( + config, + TOKEN_GENERATION_MODEL_TAG, + prefix_args, + mrope_position_ids, + vision_embeddings, + vision_mask, + hybrid_args=hybrid_args, + ) + + +def _debug_qwen36_arg_contract(stage: str, tag: str, config, args) -> None: + if ( + os.environ.get("QWEN36_ARG_CONTRACT_DEBUG") != "1" + and os.environ.get("QWEN36_HYBRID_APC_DEBUG") != "1" + ): + return + + names = _qwen36_arg_names(config, tag) + print( + f"[qwen36_arg_contract] stage={stage} tag={tag} argc={len(args)}", + flush=True, + ) + for idx, (name, value) in enumerate(zip(names, args)): + shape = _debug_tensor_shape(value) + dtype = getattr(value, "dtype", None) + min_value = "empty" + max_value = "empty" + if value is not None and hasattr(value, "numel") and value.numel() > 0: + try: + flat = value.detach().reshape(-1) if hasattr(value, "detach") else value.reshape(-1) + min_value = flat.min().item() + max_value = flat.max().item() + except Exception as exc: + min_value = f"error:{type(exc).__name__}" + max_value = f"error:{type(exc).__name__}" + print( + "[qwen36_arg_contract] " + f"stage={stage} tag={tag} index={idx} name={name} " + f"shape={shape} dtype={dtype} min={min_value} max={max_value}", + flush=True, + ) + + +def _debug_qwen36_flat_values(value) -> str: + if value is None: + return "None" + if not hasattr(value, "reshape"): + return repr(value) + try: + flat = value.detach().reshape(-1) if hasattr(value, "detach") else value.reshape(-1) + return repr(flat.tolist()) + except Exception as exc: + return f"error:{type(exc).__name__}" + + +def _debug_qwen36_hybrid_gdn_state( + tag: str, + tensor: torch.Tensor, + *, + layer_id: int, + bank_idx: int, + slot_ids: torch.Tensor, + mask: torch.Tensor, + seq_ids: torch.Tensor | None, +) -> None: + if os.environ.get("QWEN36_HYBRID_GDN_STATE_DEBUG") != "1": + return + shape = _debug_tensor_shape(tensor) + dtype = getattr(tensor, "dtype", None) + total = 0 + finite_count = "error" + nan_count = "error" + posinf_count = "error" + neginf_count = "error" + max_abs = "error" + mean_abs = "error" + try: + flat = tensor.detach().float().reshape(-1) + total = int(flat.numel()) + if total > 0: + finite = torch.isfinite(flat) + finite_i = finite.to(torch.int32) + finite_count = int(finite_i.sum().item()) + nan_count = int(torch.isnan(flat).to(torch.int32).sum().item()) + posinf_count = int(torch.isposinf(flat).to(torch.int32).sum().item()) + neginf_count = int(torch.isneginf(flat).to(torch.int32).sum().item()) + safe = torch.where(finite, flat, torch.zeros_like(flat)).abs() + max_abs = float(safe.max().item()) + mean_abs = float((safe.sum() / max(finite_count, 1)).item()) + else: + finite_count = 0 + nan_count = 0 + posinf_count = 0 + neginf_count = 0 + max_abs = "empty" + mean_abs = "empty" + except Exception as exc: + finite_count = f"error:{type(exc).__name__}" + nan_count = f"error:{type(exc).__name__}" + posinf_count = f"error:{type(exc).__name__}" + neginf_count = f"error:{type(exc).__name__}" + max_abs = f"error:{type(exc).__name__}" + mean_abs = f"error:{type(exc).__name__}" + + print( + "[qwen36_hybrid_gdn_state] " + f"tag={tag} layer={layer_id} bank={bank_idx} " + f"slot_ids={_debug_qwen36_flat_values(slot_ids)} " + f"mask={_debug_qwen36_flat_values(mask)} " + f"seq_ids={_debug_qwen36_flat_values(seq_ids)} " + f"shape={shape} dtype={dtype} finite={finite_count}/{total} " + f"nan={nan_count} posinf={posinf_count} neginf={neginf_count} " + f"max_abs={max_abs} mean_abs={mean_abs}", + flush=True, + ) + + +def _validate_qwen36_tkg_input_ids(input_ids, vocab_size) -> None: + if input_ids is None or not hasattr(input_ids, "numel") or input_ids.numel() == 0: + raise ValueError("Qwen3.6 TKG input_ids must be a non-empty tensor") + if input_ids.dtype not in (torch.int32, torch.int64): + raise ValueError( + "Qwen3.6 TKG input_ids must be int32 or int64, " + f"got {input_ids.dtype}" + ) + min_id = int(input_ids.min().item()) + max_id = int(input_ids.max().item()) + if min_id < 0: + raise ValueError(f"Qwen3.6 TKG input_ids contains negative token id {min_id}") + if vocab_size is not None and max_id >= int(vocab_size): + raise ValueError( + "Qwen3.6 TKG input_ids contains out-of-vocab token id " + f"{max_id}; vocab_size={int(vocab_size)}" + ) + + +def _qwen36_query_lengths(full_context_lens, computed_context_lens) -> list[int] | None: + if ( + full_context_lens is None + or computed_context_lens is None + or not hasattr(full_context_lens, "numel") + or not hasattr(computed_context_lens, "numel") + or full_context_lens.numel() == 0 + or computed_context_lens.numel() == 0 + ): + return None + full_values = full_context_lens.reshape(-1).to(torch.int64) + computed_values = computed_context_lens.reshape(-1).to(torch.int64) + count = min(int(full_values.numel()), int(computed_values.numel())) + if count <= 0: + return None + return [ + max(0, int(full_values[idx].item()) - int(computed_values[idx].item())) + for idx in range(count) + ] + + +def _qwen36_prefill_has_incomplete_row(prefill_completion_state) -> bool: + if prefill_completion_state is None: + return False + if hasattr(prefill_completion_state, "numel"): + if prefill_completion_state.numel() == 0: + return False + return not bool(prefill_completion_state.reshape(-1).to(torch.bool).all().item()) + try: + values = list(prefill_completion_state) + except TypeError: + return not bool(prefill_completion_state) + return any(not bool(value) for value in values) + + +def _qwen36_hybrid_apc_mask_has_active_row(mask) -> bool: + if mask is None: + return False + if hasattr(mask, "numel"): + if mask.numel() == 0: + return False + try: + return bool(mask.reshape(-1).to(torch.bool).any().item()) + except (RuntimeError, TypeError, ValueError): + # If a non-empty control tensor cannot be inspected on the host, keep + # the existing controls and avoid preparing the request twice. + return True + try: + values = list(mask) + except TypeError: + return bool(mask) + return any(bool(value) for value in values) + + +def _qwen36_hybrid_apc_controls_need_prepare( + hybrid_restore_mask, + hybrid_commit_mask, +) -> bool: + return not ( + _qwen36_hybrid_apc_mask_has_active_row(hybrid_restore_mask) + or _qwen36_hybrid_apc_mask_has_active_row(hybrid_commit_mask) + ) + + +def _qwen36_hybrid_apc_controls_materialized( + hybrid_restore_mask, + hybrid_restore_prefix_lens, + hybrid_commit_mask, +) -> bool: + return not _qwen36_hybrid_apc_controls_need_prepare( + hybrid_restore_mask, + hybrid_commit_mask, + ) or _qwen36_hybrid_apc_mask_has_active_row(hybrid_restore_prefix_lens) + + +def _qwen36_is_prefill_request( + input_ids, + position_ids, + *, + full_context_lens=None, + computed_context_lens=None, + prefill_completion_state=None, +) -> bool: + if _qwen36_prefill_has_incomplete_row(prefill_completion_state): + return True + + query_lengths = _qwen36_query_lengths(full_context_lens, computed_context_lens) + if ( + query_lengths is not None + and len(query_lengths) > 1 + and input_ids.ndim >= 2 + and input_ids.shape[0] == 1 + and input_ids.shape[-1] == len(query_lengths) + ): + return any(query_len > 1 for query_len in query_lengths) + + # Warm prefix-cache suffixes may start at a nonzero position, but they are + # still multi-token CTE requests. TKG must remain a one-token decode path. + if input_ids.shape[-1] > 1: + return True + return position_ids.min().item() == 0 + + +def _qwen36_deltanet_padding_mask( + *, + input_ids, + inputs_embeds, + attention_mask, + padding_idx, + is_for_context_encoding, + hybrid_restore_mask=None, + num_queries=None, +): + if padding_idx is None: + token_padding_mask = torch.ones( + (*input_ids.shape, 1), + dtype=inputs_embeds.dtype, + device=inputs_embeds.device, + ) + else: + token_padding_mask = ( + (input_ids != padding_idx).unsqueeze(-1).to(inputs_embeds.dtype) + ) + + query_padding_mask = None + if ( + is_for_context_encoding + and num_queries is not None + and hasattr(num_queries, "numel") + and num_queries.numel() >= input_ids.shape[0] + ): + query_lens = num_queries.reshape(-1)[: input_ids.shape[0]].to( + device=inputs_embeds.device, + dtype=torch.long, + ) + positions = torch.arange( + input_ids.shape[1], + device=inputs_embeds.device, + dtype=torch.long, + ) + query_padding_mask = ( + positions.unsqueeze(0) < query_lens.unsqueeze(1) + ).unsqueeze(-1).to(inputs_embeds.dtype) + + if ( + is_for_context_encoding + and query_padding_mask is not None + ): + deltanet_padding_mask = query_padding_mask + elif ( + is_for_context_encoding + and attention_mask is not None + and attention_mask.ndim == 2 + ): + attention_padding_mask = attention_mask.unsqueeze(-1).to(inputs_embeds.dtype) + if attention_padding_mask.shape[1] == inputs_embeds.shape[1]: + deltanet_padding_mask = attention_padding_mask + else: + deltanet_padding_mask = token_padding_mask + else: + deltanet_padding_mask = token_padding_mask + + if ( + is_for_context_encoding + and hybrid_restore_mask is not None + and hasattr(hybrid_restore_mask, "numel") + and hybrid_restore_mask.numel() > 0 + ): + restore_active = hybrid_restore_mask.reshape(-1).to(torch.bool) + if restore_active.numel() < input_ids.shape[0]: + restore_active = torch.cat( + [ + restore_active, + torch.zeros( + input_ids.shape[0] - restore_active.numel(), + dtype=torch.bool, + device=restore_active.device, + ), + ], + dim=0, + ) + restore_active = restore_active[: input_ids.shape[0]].to( + device=inputs_embeds.device + ).view(-1, 1, 1) + deltanet_padding_mask = torch.where( + restore_active, + token_padding_mask, + deltanet_padding_mask, + ) + return deltanet_padding_mask + + +def _qwen36_unpack_packed_decode_batch( + *, + input_ids, + attention_mask, + position_ids, + seq_ids, + adapter_ids, + slot_mapping, + full_context_lens, + computed_context_lens, +): + query_lengths = _qwen36_query_lengths(full_context_lens, computed_context_lens) + if ( + query_lengths is None + or len(query_lengths) <= 1 + or any(query_len > 1 for query_len in query_lengths) + or input_ids.ndim < 2 + or input_ids.shape[0] != 1 + or input_ids.shape[-1] != len(query_lengths) + ): + return input_ids, attention_mask, position_ids, seq_ids, adapter_ids, slot_mapping + + batch_size = len(query_lengths) + + def _unpack_token_rows(value): + if ( + value is not None + and hasattr(value, "ndim") + and value.ndim >= 2 + and value.shape[0] == 1 + and value.shape[1] == batch_size + ): + return value.reshape(batch_size, 1, *value.shape[2:]).contiguous() + return value + + def _repair_batch_vector(value, *, fill_from_index: bool = False): + if value is None or not hasattr(value, "numel") or value.numel() == 0: + return value + flattened = value.reshape(-1) + if flattened.numel() == batch_size: + return flattened + if flattened.numel() == 1 and batch_size > 1: + if fill_from_index: + return torch.arange( + batch_size, + dtype=value.dtype, + device=value.device, + ) + return flattened[:1].expand(batch_size).contiguous() + return value + + input_ids = _unpack_token_rows(input_ids) + position_ids = _unpack_token_rows(position_ids) + slot_mapping = _unpack_token_rows(slot_mapping) + if ( + attention_mask is not None + and hasattr(attention_mask, "ndim") + and attention_mask.ndim >= 2 + and attention_mask.shape[0] == 1 + and attention_mask.shape[1] == batch_size + and computed_context_lens is not None + and hasattr(computed_context_lens, "numel") + and computed_context_lens.numel() >= batch_size + ): + context_lens = computed_context_lens.reshape(-1).to(torch.int64)[:batch_size] + max_context_len = max(1, int(context_lens.max().item())) + repaired_mask = torch.zeros( + (batch_size, max_context_len), + dtype=attention_mask.dtype, + device=attention_mask.device, + ) + for row_idx, context_len in enumerate(context_lens): + active_len = max(0, min(int(context_len.item()), max_context_len)) + if active_len: + repaired_mask[row_idx, :active_len] = 1 + attention_mask = repaired_mask + if ( + slot_mapping is not None + and hasattr(slot_mapping, "ndim") + and slot_mapping.ndim == 1 + and int(slot_mapping.numel()) == batch_size + ): + slot_mapping = slot_mapping.reshape(batch_size, 1).contiguous() + seq_ids = _repair_batch_vector(seq_ids, fill_from_index=True) + adapter_ids = _repair_batch_vector(adapter_ids) + return input_ids, attention_mask, position_ids, seq_ids, adapter_ids, slot_mapping + + +def _qwen36_hashable_request_id(request_id: Any) -> Hashable: + if isinstance(request_id, list): + return tuple(request_id) + try: + hash(request_id) + except TypeError: + return repr(request_id) + return request_id + + +def _qwen36_metadata_for_request( + metadata_by_request_id, + request_id, +) -> dict[str, Any] | None: + if not isinstance(metadata_by_request_id, dict): + return None + normalized = _qwen36_hashable_request_id(request_id) + metadata = metadata_by_request_id.get(normalized) + if metadata is None and request_id is not None: + metadata = metadata_by_request_id.get(str(request_id)) + return metadata if isinstance(metadata, dict) else None + + +def _qwen36_request_metadata_values( + metadata_by_request_id, + request_ids, + key: str, +): + if request_ids is None: + return None + if isinstance(request_ids, list): + request_ids = tuple(request_ids) + elif not isinstance(request_ids, tuple): + request_ids = (request_ids,) + + values = [] + found = False + for request_id in request_ids: + metadata = _qwen36_metadata_for_request(metadata_by_request_id, request_id) + value = metadata.get(key) if metadata is not None else None + values.append(value) + found = found or value is not None + if not found: + return None + return values[0] if len(values) == 1 else tuple(values) + + +def _qwen36_request_ids_have_metadata( + metadata_by_request_id, + request_ids, +) -> bool: + return any( + _qwen36_request_metadata_values( + metadata_by_request_id, + request_ids, + key, + ) + is not None + for key in ( + "cumulative_hashes_by_prefix_len", + "attention_block_refs_by_prefix_len", + "request_prefix_len", + "vllm_attention_hit_len", + ) + ) + + +def _qwen36_select_vllm_hybrid_apc_request_ids( + metadata_by_request_id, + *request_id_groups, +): + first_present = None + for request_ids in request_id_groups: + if request_ids is None: + continue + if first_present is None: + first_present = request_ids + if _qwen36_request_ids_have_metadata(metadata_by_request_id, request_ids): + return request_ids + return first_present + + +def _qwen36_flat_item_count(value: Any) -> int: + if value is None: + return 0 + if hasattr(value, "numel"): + try: + return int(value.reshape(-1).numel()) + except Exception: + return 0 + if isinstance(value, (list, tuple)): + return len(value) + return 1 + + +def _qwen36_pad_batch_repeat_first(value, target_batch): + if value is None or not hasattr(value, "numel") or value.numel() == 0: + return value + if value.ndim == 0 or value.shape[0] >= target_batch: + return value + pad_n = target_batch - value.shape[0] + return torch.cat([value, value[:1].expand(pad_n, *value.shape[1:])], dim=0) + + +def _qwen36_pad_batch_with_value(value, target_batch, fill_value): + if value is None or not hasattr(value, "numel") or value.numel() == 0: + return value + if value.ndim == 0 or value.shape[0] >= target_batch: + return value + pad_shape = (target_batch - value.shape[0],) + tuple(value.shape[1:]) + pad = torch.full(pad_shape, fill_value, dtype=value.dtype, device=value.device) + return torch.cat([value, pad], dim=0) + + +def _qwen36_pad_hybrid_restore_controls_for_dummy_cte_rows( + restore_slot_ids, + restore_mask, + restore_prefix_lens, + target_batch, +): + return ( + _qwen36_pad_batch_with_value(restore_slot_ids, target_batch, 0), + _qwen36_pad_batch_with_value(restore_mask, target_batch, 0), + _qwen36_pad_batch_with_value(restore_prefix_lens, target_batch, 0), + ) + + +def _qwen36_update_state_rows_by_seq_ids(previous_state, new_rows, seq_ids): + if ( + previous_state is None + or new_rows is None + or seq_ids is None + or not hasattr(previous_state, "shape") + or not hasattr(new_rows, "shape") + or not hasattr(seq_ids, "numel") + or previous_state.ndim != new_rows.ndim + or previous_state.shape[1:] != new_rows.shape[1:] + or previous_state.shape[0] <= 0 + or new_rows.shape[0] <= 0 + or seq_ids.numel() == 0 + ): + return new_rows + + row_count = min(int(new_rows.shape[0]), int(seq_ids.reshape(-1).shape[0])) + if row_count <= 0: + return previous_state * 1 + + output = previous_state * 1 + seq_ids_flat = seq_ids.reshape(-1)[:row_count].to( + device=previous_state.device, + dtype=torch.long, + ) + slot_axis = torch.arange( + int(previous_state.shape[0]), + dtype=torch.long, + device=previous_state.device, + ) + broadcast_shape = (int(previous_state.shape[0]),) + ( + 1, + ) * (previous_state.ndim - 1) + typed_rows = new_rows[:row_count].to(previous_state.dtype) + for row_idx in range(row_count): + seq_id = seq_ids_flat[row_idx] + valid_seq = torch.logical_and( + seq_id >= 0, + seq_id < int(previous_state.shape[0]), + ) + write_mask = torch.logical_and(valid_seq, slot_axis == seq_id).view( + broadcast_shape + ) + row_value = typed_rows[row_idx : row_idx + 1].expand_as(output) + output = torch.where(write_mask, row_value, output) + return output + + +def _qwen36_preserve_inactive_state_rows(new_state, previous_state, active_rows): + if ( + new_state is None + or previous_state is None + or active_rows is None + or not hasattr(new_state, "shape") + or not hasattr(previous_state, "shape") + or not hasattr(active_rows, "numel") + or new_state.shape != previous_state.shape + or active_rows.numel() == 0 + ): + return new_state + active_rows = active_rows.reshape(-1).to(device=new_state.device, dtype=torch.bool) + row_count = min(int(active_rows.numel()), int(new_state.shape[0])) + if row_count <= 0: + return new_state + if row_count < int(new_state.shape[0]): + active_rows = torch.cat( + [ + active_rows[:row_count], + torch.ones( + int(new_state.shape[0]) - row_count, + dtype=torch.bool, + device=new_state.device, + ), + ], + dim=0, + ) + else: + active_rows = active_rows[: int(new_state.shape[0])] + view_shape = (int(new_state.shape[0]),) + (1,) * (new_state.ndim - 1) + active_rows = active_rows.view(view_shape) + return torch.where(active_rows, new_state, previous_state) + + +def _qwen36_active_state_rows(valid_mask_1d, seq_ids): + if ( + valid_mask_1d is None + or not hasattr(valid_mask_1d, "numel") + or valid_mask_1d.numel() == 0 + ): + return None + active_rows = valid_mask_1d.squeeze(-1).to(torch.bool).any(dim=-1) + if seq_ids is not None and hasattr(seq_ids, "numel") and seq_ids.numel() > 0: + seq_active = seq_ids.reshape(-1).to( + device=active_rows.device, + dtype=torch.long, + )[: active_rows.numel()] >= 0 + active_rows = active_rows & seq_active + return active_rows + + +def _qwen36_request_ids_tuple(request_ids): + if request_ids is None: + return None + if isinstance(request_ids, list): + return tuple(request_ids) + if isinstance(request_ids, tuple): + return request_ids + return (request_ids,) + + +def _qwen36_request_ids_from_hybrid_apc_records(records): + if records is None: + return None + if isinstance(records, dict): + records = (records,) + elif isinstance(records, list): + records = tuple(records) + if not isinstance(records, tuple): + return None + request_ids = [] + for record in records: + if not isinstance(record, dict): + return None + request_id = record.get("request_id") + if request_id is None: + return None + request_ids.append(request_id) + return tuple(request_ids) if request_ids else None + + +def _qwen36_max_seq_slots_for_request_ids(model, seq_ids, request_count): + max_slots = int(request_count or 0) + for owner in ( + model, + getattr(model, "neuron_config", None), + getattr(getattr(model, "context_encoding_model", None), "neuron_config", None), + getattr(getattr(model, "token_generation_model", None), "neuron_config", None), + ): + for attr in ("batch_size", "max_batch_size", "max_num_seqs"): + value = getattr(owner, attr, None) + if value is None: + continue + try: + max_slots = max(max_slots, int(value)) + except (TypeError, ValueError): + pass + if seq_ids is not None and hasattr(seq_ids, "numel") and seq_ids.numel() > 0: + flat = seq_ids.reshape(-1) + try: + non_negative = flat[flat >= 0] + if non_negative.numel() > 0: + max_slots = max(max_slots, int(non_negative.max().item()) + 1) + except Exception: + pass + return max(1, max_slots) + + +def _qwen36_stable_seq_ids_for_request_ids(model, seq_ids, request_ids): + request_ids = _qwen36_request_ids_tuple(request_ids) + if not request_ids: + return seq_ids + + normalized_request_ids = tuple( + _qwen36_hashable_request_id(request_id) for request_id in request_ids + ) + if any(request_id is None for request_id in normalized_request_ids): + return seq_ids + + slot_by_request = getattr(model, "_qwen36_hybrid_seq_slot_by_request", None) + request_by_slot = getattr(model, "_qwen36_hybrid_request_by_seq_slot", None) + if not isinstance(slot_by_request, dict) or not isinstance(request_by_slot, dict): + slot_by_request = {} + request_by_slot = {} + setattr(model, "_qwen36_hybrid_seq_slot_by_request", slot_by_request) + setattr(model, "_qwen36_hybrid_request_by_seq_slot", request_by_slot) + + max_slots = _qwen36_max_seq_slots_for_request_ids( + model, + seq_ids, + len(normalized_request_ids), + ) + active_request_ids = set(normalized_request_ids) + for stale_slot, stale_owner in list(request_by_slot.items()): + if stale_owner in active_request_ids: + continue + request_by_slot.pop(stale_slot, None) + slot_by_request.pop(stale_owner, None) + + assigned_slots = [] + for request_id in normalized_request_ids: + slot = slot_by_request.get(request_id) + if slot is None or slot < 0 or slot >= max_slots: + free_slots = [ + candidate + for candidate in range(max_slots) + if candidate not in request_by_slot + ] + if not free_slots: + return seq_ids + slot = free_slots[0] + slot_by_request[request_id] = slot + request_by_slot[slot] = request_id + assigned_slots.append(slot) + + dtype = seq_ids.dtype if hasattr(seq_ids, "dtype") else torch.int32 + if seq_ids is not None and hasattr(seq_ids, "device"): + device = seq_ids.device + else: + device = None + kwargs = {"dtype": dtype} + if device is not None: + kwargs["device"] = device + return torch.tensor(assigned_slots, **kwargs) + + +def _qwen36_select_vllm_hybrid_apc_request_ids_for_input( + metadata_by_request_id, + *, + all_request_ids, + new_request_ids, + full_context_lens, + computed_context_lens, + prefill_completion_state, +): + all_request_ids_tuple = _qwen36_request_ids_tuple(all_request_ids) + logical_request_count = max( + _qwen36_flat_item_count(full_context_lens), + _qwen36_flat_item_count(computed_context_lens), + _qwen36_flat_item_count(prefill_completion_state), + ) + if ( + logical_request_count > 1 + and all_request_ids_tuple is not None + and len(all_request_ids_tuple) == logical_request_count + ): + # Keep request identity aligned with the model row order. In mixed + # cached/new prefill batches, scheduler "new" ids can be a strict + # subset, but the metadata vectors still describe every model row. + return all_request_ids_tuple + return _qwen36_select_vllm_hybrid_apc_request_ids( + metadata_by_request_id, + new_request_ids, + all_request_ids, + ) + + +def _qwen36_add_vllm_hybrid_apc_metadata( + hybrid_apc_request_dict: dict[str, Any], + *, + request_ids, + metadata_by_request_id, +) -> None: + for key in ( + "cumulative_hashes_by_prefix_len", + "attention_block_refs_by_prefix_len", + "request_prefix_len", + "vllm_attention_hit_len", + "active_suffix_len", + "full_input_ids", + ): + value = _qwen36_request_metadata_values( + metadata_by_request_id, + request_ids, + key, + ) + if value is not None: + if key == "full_input_ids" and not isinstance(value, torch.Tensor): + input_ids = hybrid_apc_request_dict.get("input_ids") + dtype = ( + input_ids.dtype + if isinstance(input_ids, torch.Tensor) + else torch.int64 + ) + device = ( + input_ids.device + if isinstance(input_ids, torch.Tensor) + else None + ) + value = torch.tensor([list(value)], dtype=dtype, device=device) + hybrid_apc_request_dict[key] = value + + +def _debug_logits_stage(stage: str, tensor) -> None: + if os.environ.get("QWEN36_LOGIT_STAGE_DEBUG") != "1": + return + if tensor is None or not hasattr(tensor, "numel"): + print( + f"[qwen36_logits_debug] stage={stage} tensor=none", + flush=True, + ) + return + if tensor.numel() == 0: + print( + f"[qwen36_logits_debug] stage={stage} " + f"shape={tuple(tensor.shape)} dtype={tensor.dtype} device={tensor.device} empty", + flush=True, + ) + return + + try: + with torch.no_grad(): + flat = tensor.detach().reshape(-1) + if torch.is_floating_point(flat): + finite_mask = torch.isfinite(flat) + finite_count = int(finite_mask.sum().item()) + nan_count = int(torch.isnan(flat).sum().item()) + posinf_count = int( + torch.logical_and(torch.isinf(flat), flat > 0).sum().item() + ) + neginf_count = int( + torch.logical_and(torch.isinf(flat), flat < 0).sum().item() + ) + if finite_count: + finite_flat = flat[finite_mask].float() + finite_min = float(finite_flat.min().item()) + finite_max = float(finite_flat.max().item()) + else: + finite_min = "none" + finite_max = "none" + print( + "[qwen36_logits_debug] " + f"stage={stage} shape={tuple(tensor.shape)} dtype={tensor.dtype} " + f"device={tensor.device} numel={tensor.numel()} finite={finite_count} " + f"nan={nan_count} posinf={posinf_count} neginf={neginf_count} " + f"finite_min={finite_min} finite_max={finite_max}", + flush=True, + ) + else: + print( + "[qwen36_logits_debug] " + f"stage={stage} shape={tuple(tensor.shape)} dtype={tensor.dtype} " + f"device={tensor.device} numel={tensor.numel()} " + f"minmax={_debug_tensor_minmax(tensor)}", + flush=True, + ) + except Exception as exc: + print( + "[qwen36_logits_debug] " + f"stage={stage} summary_error={type(exc).__name__}: {exc}", + flush=True, + ) + + +def _qwen36_output_logits_for_return(logits, lm_head, neuron_config): + if not ( + getattr(neuron_config, "output_logits", False) + and getattr(neuron_config, "on_device_sampling_config", None) is not None + and not getattr(lm_head, "gather_output", True) + ): + return logits + return _gather_along_dim( + logits, + partition_dim=2, + process_group=getattr(lm_head, "tensor_parallel_group", None), + ) + + +class NeuronQwen35Model(NeuronBaseModel): + def setup_attr_for_model(self, config: Qwen35InferenceConfig): + self.on_device_sampling = ( + config.neuron_config.on_device_sampling_config is not None + ) + self.tp_degree = config.neuron_config.tp_degree + self.hidden_size = config.hidden_size + self.num_attention_heads = config.num_attention_heads + self.num_key_value_heads = config.num_key_value_heads + self.max_batch_size = config.neuron_config.max_batch_size + self.buckets = config.neuron_config.buckets + + def init_model(self, config: Qwen35InferenceConfig): + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = ParallelEmbedding( + config.vocab_size, + config.hidden_size, + self.padding_idx, + dtype=config.neuron_config.torch_dtype, + shard_across_embedding=True, + ) + self.layers = nn.ModuleList( + [ + NeuronQwen35DecoderLayer(config, layer_idx) + for layer_idx in range(config.num_hidden_layers) + ] + ) + self.norm = get_rmsnorm_cls()(self.hidden_size, eps=config.rms_norm_eps) + self.lm_head = ColumnParallelLinear( + config.hidden_size, + config.vocab_size, + gather_output=False if self.on_device_sampling else True, + bias=False, + ) + + # mRoPE embedding for VL + self.mrope_emb = Qwen35MRoPEEmbedding(config) + + def init_inference_optimization(self, config: Qwen35InferenceConfig): + super().init_inference_optimization(config) + if getattr(config, "use_hybrid_apc_manager", False): + if getattr(config.neuron_config, "is_block_kv_layout", False): + self.kv_mgr = QwenHybridBlockKVCacheManager( + config, + num_kv_head=self.num_key_value_heads, + ) + self.hybrid_gdn_checkpoint_cache = HybridGDNCheckpointCache(config) + elif getattr(config, "use_hybrid_cache_manager", False): + self.kv_mgr = HybridDeltaNetCacheManager( + config, + num_kv_head=self.num_key_value_heads, + global_rank=self.rank_util, + attention_chunk_size=self.attention_chunk_size, + sliding_window=self.sliding_window, + windowed_context_encoding_size=self.windowed_context_encoding_size, + layer_to_cache_size_mapping=self.layer_to_cache_size_mapping, + ) + + @property + def _deltanet_state_params(self): + """Return DeltaNet state nn.Parameters in alias order.""" + params = [] + for layer in self.layers: + if hasattr(layer, "linear_attn"): + params.append(layer.linear_attn.recurrent_state_buffer) + params.append(layer.linear_attn.conv_state_buffer) + return params + + @property + def _hybrid_gdn_checkpoint_params(self): + if not hasattr(self, "hybrid_gdn_checkpoint_cache"): + return [] + return self.hybrid_gdn_checkpoint_cache.checkpoint_params + + def encode_vision_to_input(self, inputs_embeds, vision_embeddings, vision_mask): + """Scatter vision embeddings into text input embeddings at image-token + positions, using exactly the Qwen3-VL upstream pattern. + + vision_embeddings: (1, seq_len, hidden). Real vision embeddings live in + slots i < n_vis; pad slots i >= n_vis are ZEROS. + vision_mask: (1, seq_len, 1) int32. Real slots (i < n_vis) hold image-token + positions in [0, seq_len). Pad slots (i >= n_vis) hold seq_len-1 + (the last padding position of input_ids), so scatter writes zero to + a single padding slot — safe because that position has + attention_mask == 0. + + The scatter uses PyTorch index_put_(accumulate=False), matching + `scatter_by_index_put` in neuronx_distributed_inference.models.llama4. + """ + _, max_positions, embedding_dim = inputs_embeds.shape + h_new = inputs_embeds.clone() + vision_flat = vision_embeddings.reshape(-1, embedding_dim) + positions_flat = vision_mask.reshape(-1) + num_positions = positions_flat.shape[0] + vision_flat = vision_flat[:num_positions] + h_new.view(-1, embedding_dim).index_put_( + (positions_flat,), vision_flat, accumulate=False + ) + return h_new + + def get_model_output( + self, + input_ids=None, + seq_ids=None, + attention_mask=None, + position_ids=None, + past_key_values=None, + active_mask=None, + inputs_embeds=None, + prev_hidden=None, + adapter_ids=None, + rotary_position_ids=None, + update_cache=False, + is_for_context_encoding=False, + vision_embeddings=None, + vision_mask=None, + hybrid_restore_slot_ids=None, + hybrid_restore_mask=None, + hybrid_restore_prefix_lens=None, + hybrid_commit_slot_ids=None, + hybrid_commit_mask=None, + local_attn_mask=None, + windowed_context_encoding_window_idx=-1, + padding_mask=None, + **kwargs, + ): + """Override to collect DeltaNet state tensors from decoder layers.""" + batch_size, seq_length = input_ids.shape[:2] + if self.config.neuron_config.layer_boundary_markers: + input_ids = ModuleMarkerStartWrapper()(input_ids) + + past_key_values_length = 0 + if past_key_values is not None: + if hasattr(self.kv_mgr, "get_seq_length"): + past_key_values_length = self.kv_mgr.get_seq_length(past_key_values) + else: + past_key_values_length = past_key_values[0][1].shape[2] + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + # CRITICAL: Zero out embeddings for padding tokens so DeltaNet recurrence + # is not polluted. DeltaNet has no attention mask -- it processes all + # sequence positions through a linear recurrence. Padding tokens have + # real embedding vectors which corrupt the recurrence state. + # The mask is [B, S, 1] float with 1.0 for real tokens, 0.0 for padding. + deltanet_padding_mask = _qwen36_deltanet_padding_mask( + input_ids=input_ids, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + padding_idx=self.padding_idx, + is_for_context_encoding=is_for_context_encoding, + hybrid_restore_mask=hybrid_restore_mask, + num_queries=kwargs.get("num_queries"), + ) + if is_for_context_encoding: + inputs_embeds = inputs_embeds * deltanet_padding_mask + + # Vision embedding injection. When use_text_only_cte_inputs=False we + # always trace the scatter into the graph. The input generator makes + # sure "dummy" (text-only) vision inputs are IDEMPOTENT: every pad slot + # points to the same target position and carries the same value, so + # scattering them repeatedly does not corrupt real embeddings. + # + # NOTE: the gate MUST be Python-static (trace-time) since the two + # branches produce different graphs. `shape[1] != seq_length` is not a + # reliable proxy because the input generator pads vision inputs to + # seq_length; instead we look at the compile-time config flag. + if (vision_embeddings is not None) and (vision_mask is not None): + if vision_embeddings.dtype != self.config.neuron_config.torch_dtype: + vision_embeddings = vision_embeddings.to( + self.config.neuron_config.torch_dtype + ) + traced_with_vision = ( + not getattr(self.config, "use_text_only_cte_inputs", True) + and vision_embeddings.ndim == 3 + and vision_mask.ndim == 3 + and vision_embeddings.shape[1] == seq_length + and vision_mask.shape[1] == seq_length + ) + if is_for_context_encoding and traced_with_vision: + inputs_embeds = self.encode_vision_to_input( + inputs_embeds, vision_embeddings, vision_mask + ) + elif is_for_context_encoding and vision_embeddings.numel() > 0: + inputs_embeds = inputs_embeds + vision_embeddings.sum() * 0 + inputs_embeds = ( + inputs_embeds + vision_mask.sum().to(inputs_embeds.dtype) * 0 + ) + + if position_ids is None: + device = input_ids.device if input_ids is not None else inputs_embeds.device + position_ids = torch.arange( + past_key_values_length, + seq_length + past_key_values_length, + dtype=torch.long, + device=device, + ) + position_ids = position_ids.unsqueeze(0).view(-1, seq_length) + else: + position_ids = position_ids.view(-1, seq_length).long() + + hidden_states = inputs_embeds + + # Get KV cache for TKG and for model-local chunked CTE. + use_qwen_chunked_prefill = ( + is_for_context_encoding + and getattr(self.config, "use_qwen_hybrid_chunked_prefill", False) + ) + active_block_table = kwargs.get("active_block_table", None) + cte_has_prefix_blocks = ( + is_for_context_encoding + and use_qwen_chunked_prefill + and active_block_table is not None + and getattr(active_block_table, "ndim", 0) > 1 + ) + cache_size = ( + self.config.neuron_config.seq_len + if use_qwen_chunked_prefill + else self.n_positions + ) + if (not is_for_context_encoding) or cte_has_prefix_blocks: + if self.kv_mgr is not None: + past_key_values = self.kv_mgr.get_cache( + seq_ids=seq_ids, + seq_len=cache_size, + is_for_context_encoding=is_for_context_encoding, + windowed_context_encoding_window_idx=windowed_context_encoding_window_idx, + **kwargs, + ) + + # Decoder layers + next_decoder_cache = () + deltanet_state_tensors = [] + deltanet_layer_state_pairs = [] + cos_cache = None + sin_cache = None + restored_gdn_states = None + if getattr(self.config, "use_hybrid_apc_manager", False) and hasattr( + self, "hybrid_gdn_checkpoint_cache" + ): + if hybrid_restore_prefix_lens is not None and position_ids is not None: + # Host-side request prep must set suffix position_ids to the + # restored cumulative-prefix boundary. This is a no-op on + # default zero masks, but it keeps the contract explicit. + if ( + not torch.jit.is_tracing() + and hybrid_restore_mask is not None + and bool(hybrid_restore_mask.to(torch.bool).any().item()) + ): + expected = hybrid_restore_prefix_lens.long() + observed = position_ids[:, 0].long() + if not torch.equal(observed, expected): + raise ValueError( + "hybrid APC restore prefix lens must match " + "position_ids[:, 0]" + ) + restored_gdn_states = ( + self.hybrid_gdn_checkpoint_cache.restore_to_active_rows( + layers=self.layers, + seq_ids=seq_ids, + checkpoint_slot_ids=hybrid_restore_slot_ids, + restore_mask=hybrid_restore_mask, + zero_inactive=( + is_for_context_encoding + and not _qwen36_hybrid_apc_mask_has_active_row( + hybrid_restore_prefix_lens + ) + ), + ) + ) + + # Keep CTE masks compact on the Neuron paths. Qwen attention prefill + # applies causal masking inside the attention kernel/path, while DeltaNet + # consumes deltanet_padding_mask separately. Dense SxS masks are only a + # small fallback path and are not viable for long-context CTE. + use_compact_cte_attention_mask = getattr( + self.config, "use_compact_cte_attention_mask", True + ) + use_neuron_cte_attention = use_qwen_chunked_prefill or getattr( + self.config.neuron_config, "is_block_kv_layout", False + ) + # Convert 2D attention_mask to 4D causal mask for the small fallback path. + if ( + attention_mask is not None + and attention_mask.ndim == 2 + and is_for_context_encoding + and not use_compact_cte_attention_mask + and not use_neuron_cte_attention + ): + causal = torch.ones( + (seq_length, seq_length), + dtype=torch.bool, + device=attention_mask.device, + ).tril() + padding_4d = attention_mask[:, None, None, :].to(torch.bool) + attention_mask = (causal[None, None, :, :] & padding_4d).to( + attention_mask.dtype + ) + + # Pre-compute mRoPE cos/sin + if rotary_position_ids is not None and rotary_position_ids.ndim == 3: + cos_cache, sin_cache = self.mrope_emb(inputs_embeds, rotary_position_ids) + + for idx, decoder_layer in enumerate(self.layers): + past_key_value = ( + past_key_values[idx] if past_key_values is not None else None + ) + if restored_gdn_states is not None and idx in restored_gdn_states: + past_key_value = restored_gdn_states[idx] + + layer_outputs = decoder_layer( + hidden_states, + seq_ids=seq_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + active_mask=active_mask, + adapter_ids=adapter_ids, + cos_cache=cos_cache, + sin_cache=sin_cache, + rotary_position_ids=rotary_position_ids, + kv_mgr=self.kv_mgr, + get_kv_per_layer=False, + update_kv_per_layer=False, + idx=idx, + is_for_context_encoding=is_for_context_encoding, + seq_len=cache_size, + residual=None, + local_mask=local_attn_mask, + windowed_context_encoding_window_idx=windowed_context_encoding_window_idx, + padding_mask=padding_mask, + deltanet_padding_mask=deltanet_padding_mask, + qwen_chunked_prefill_update=use_qwen_chunked_prefill, + qwen_chunked_valid_mask=deltanet_padding_mask.squeeze(-1) + if use_qwen_chunked_prefill + else None, + **kwargs, + ) + + hidden_states = layer_outputs[0] + kv = layer_outputs[1] + next_decoder_cache += (kv,) + cos_cache, sin_cache = layer_outputs[2:4] + + # Collect DeltaNet state tensors + deltanet_states = layer_outputs[5] if len(layer_outputs) > 5 else None + if deltanet_states is not None: + deltanet_state_tensors.append(deltanet_states[0]) + deltanet_state_tensors.append(deltanet_states[1]) + deltanet_layer_state_pairs.append( + (idx, deltanet_states[0], deltanet_states[1]) + ) + + # Update KV cache + if update_cache: + next_decoder_cache = self.kv_mgr.update_cache( + is_for_context_encoding=is_for_context_encoding, + seq_ids=seq_ids, + position_ids=position_ids, + new_key_values=next_decoder_cache, + seq_len=cache_size, + windowed_context_encoding_window_idx=windowed_context_encoding_window_idx, + qwen_chunked_prefill_update=use_qwen_chunked_prefill, + qwen_chunked_valid_mask=deltanet_padding_mask.squeeze(-1) + if use_qwen_chunked_prefill + else None, + **kwargs, + ) + + if getattr(self.config, "use_hybrid_apc_manager", False) and hasattr( + self, "hybrid_gdn_checkpoint_cache" + ): + commit_during_tkg = bool( + getattr(self.config, "hybrid_apc_commit_during_token_generation", False) + ) + if not is_for_context_encoding and not commit_during_tkg: + self._hybrid_gdn_checkpoint_updated_states = [] + else: + self._hybrid_gdn_checkpoint_updated_states = ( + self.hybrid_gdn_checkpoint_cache.commit_from_active_rows( + layer_state_pairs=deltanet_layer_state_pairs, + seq_ids=seq_ids, + checkpoint_slot_ids=hybrid_commit_slot_ids, + commit_mask=hybrid_commit_mask, + ) + ) + + _debug_logits_stage("before_final_norm", hidden_states) + hidden_states = self.norm(hidden_states) + _debug_logits_stage("after_final_norm_full", hidden_states) + + self._deltanet_updated_states = deltanet_state_tensors + + return (hidden_states, next_decoder_cache) + + def forward( + self, + input_ids, + attention_mask, + position_ids, + seq_ids, + sampling_params, + prev_hidden=None, + adapter_ids=None, + accepted_indices=None, + current_length=None, + medusa_mask=None, + scatter_index=None, + slot_mapping=None, + active_block_table=None, + num_queries=None, + computed_context_lens=None, + tile_q_indices=None, + tile_block_tables=None, + tile_masks=None, + inputs_embeds=None, + kv_cache=None, + active_mask=None, + rotary_position_id=None, + vision_embeddings=None, + vision_mask=None, + hybrid_restore_slot_ids=None, + hybrid_restore_mask=None, + hybrid_restore_prefix_lens=None, + hybrid_commit_slot_ids=None, + hybrid_commit_mask=None, + ): + """Override base forward to append DeltaNet state tensors to output.""" + prev_hidden = self.set_none_if_empty(prev_hidden) + adapter_ids = self.set_none_if_empty(adapter_ids) + accepted_indices = self.set_none_if_empty(accepted_indices) + current_length = self.set_none_if_empty(current_length) + medusa_mask = self.set_none_if_empty(medusa_mask) + scatter_index = self.set_none_if_empty(scatter_index) + slot_mapping = self.set_none_if_empty(slot_mapping) + active_block_table = self.set_none_if_empty(active_block_table) + num_queries = self.set_none_if_empty(num_queries) + computed_context_lens = self.set_none_if_empty(computed_context_lens) + tile_q_indices = self.set_none_if_empty(tile_q_indices) + tile_block_tables = self.set_none_if_empty(tile_block_tables) + tile_masks = self.set_none_if_empty(tile_masks) + inputs_embeds = self.set_none_if_empty(inputs_embeds) + kv_cache = self.set_none_if_empty(kv_cache) + active_mask = self.set_none_if_empty(active_mask) + rotary_position_id = self.set_none_if_empty(rotary_position_id) + vision_embeddings = self.set_none_if_empty(vision_embeddings) + vision_mask = self.set_none_if_empty(vision_mask) + hybrid_restore_slot_ids = self.set_none_if_empty(hybrid_restore_slot_ids) + hybrid_restore_mask = self.set_none_if_empty(hybrid_restore_mask) + hybrid_restore_prefix_lens = self.set_none_if_empty(hybrid_restore_prefix_lens) + hybrid_commit_slot_ids = self.set_none_if_empty(hybrid_commit_slot_ids) + hybrid_commit_mask = self.set_none_if_empty(hybrid_commit_mask) + + is_for_context_encoding = position_ids.shape[-1] != 1 and not ( + hasattr(self.neuron_config, "speculation_length") + and position_ids.shape[-1] == self.neuron_config.speculation_length + ) + + seq_ids = seq_ids.to(torch.int32) + attn_mask = attention_mask + + hidden_states, updated_kv_cache = self.get_model_output( + input_ids=input_ids, + seq_ids=seq_ids, + attention_mask=attn_mask, + position_ids=position_ids, + active_mask=active_mask, + inputs_embeds=inputs_embeds, + adapter_ids=adapter_ids, + rotary_position_ids=rotary_position_id, + update_cache=True, + is_for_context_encoding=is_for_context_encoding, + padding_mask=None, + active_block_table=active_block_table, + scatter_index=slot_mapping + if getattr(self, "is_block_kv_layout", False) + else scatter_index, + vision_embeddings=vision_embeddings, + vision_mask=vision_mask, + hybrid_restore_slot_ids=hybrid_restore_slot_ids, + hybrid_restore_mask=hybrid_restore_mask, + hybrid_restore_prefix_lens=hybrid_restore_prefix_lens, + hybrid_commit_slot_ids=hybrid_commit_slot_ids, + hybrid_commit_mask=hybrid_commit_mask, + num_queries=num_queries, + computed_context_lens=computed_context_lens, + ) + + batch_size = input_ids.shape[0] + if not getattr(self, "sliced_hidden", False): + if not is_for_context_encoding: + pass + else: + if getattr(self.config, "use_qwen_hybrid_chunked_prefill", False): + query_index = None + if ( + num_queries is not None + and hasattr(num_queries, "numel") + and num_queries.numel() >= batch_size + ): + query_index = ( + num_queries.reshape(-1)[:batch_size] + .to(device=input_ids.device, dtype=torch.long) + .view(batch_size, 1) + - 1 + ).clamp(min=0) + token_index = None + if self.padding_idx is not None: + token_index = ( + (input_ids != self.padding_idx) + .sum(dim=1, keepdim=True) + .long() + - 1 + ).clamp(min=0) + if query_index is not None: + index = query_index + elif attention_mask is not None and attention_mask.ndim == 2: + attention_index = ( + attention_mask.to(torch.long).sum(dim=1, keepdim=True) + - 1 + ).clamp(min=0) + if ( + hybrid_restore_mask is not None + and hasattr(hybrid_restore_mask, "numel") + and hybrid_restore_mask.numel() > 0 + ): + restore_active = ( + hybrid_restore_mask.reshape(-1).to(torch.bool).any() + ) + index = torch.where( + restore_active, + token_index if token_index is not None else attention_index, + attention_index, + ) + else: + index = attention_index + else: + index = ( + token_index + if token_index is not None + else torch.full( + (batch_size, 1), + max(0, input_ids.shape[1] - 1), + dtype=torch.long, + device=input_ids.device, + ) + ) + else: + index = torch.max(position_ids, dim=1, keepdim=True).indices + index = index.unsqueeze(1).expand(batch_size, 1, self.hidden_size) + hidden_states = torch.gather(hidden_states, dim=1, index=index) + + _debug_logits_stage("after_final_norm", hidden_states) + _debug_logits_stage("selected_hidden_before_lm_head", hidden_states) + _debug_logits_stage("lm_head_weight", getattr(self.lm_head, "weight", None)) + logits = self.lm_head(hidden_states) + _debug_logits_stage("after_lm_head_pre_float", logits) + logits = logits.float() + _debug_logits_stage("after_lm_head", logits) + + if hasattr(self.lm_head, "pad_size"): + if self.lm_head.gather_output: + rank_id = torch.tensor(0, device=logits.device, dtype=torch.int32) + world_size = 1 + else: + from neuronx_distributed.parallel_layers import parallel_state + + rank_id = self.rank_util.get_rank() + world_size = torch.distributed.get_world_size( + group=self.lm_head.tensor_parallel_group + ) + from neuronx_distributed_inference.models.model_base import ( + mask_padded_logits, + ) + + logits = mask_padded_logits( + logits, + rank_id, + world_size, + pad_size=_effective_lm_head_pad_size( + self.lm_head, logits, self.config + ), + ) + _debug_logits_stage("after_mask_padded_logits", logits) + + if self.on_device_sampling: + res = self._sample_on_device( + logits, sampling_params, False, is_for_context_encoding + ) + else: + res = logits + + _debug_logits_stage("before_return_logits", logits) + outputs = [res] + if self.neuron_config.output_logits and self.on_device_sampling: + outputs += [ + _qwen36_output_logits_for_return( + logits, + self.lm_head, + self.neuron_config, + ) + ] + _qwen36_validate_alias_output_counts( + self, + updated_kv_cache, + is_for_context_encoding=is_for_context_encoding, + ) + outputs += updated_kv_cache + + # Append DeltaNet state tensors (for input_output_aliases) + if ( + not getattr(self.config, "use_hybrid_cache_manager", False) + and hasattr(self, "_deltanet_updated_states") + ): + outputs += self._deltanet_updated_states + if ( + getattr(self.config, "use_hybrid_apc_manager", False) + and hasattr(self, "_hybrid_gdn_checkpoint_updated_states") + ): + outputs += self._hybrid_gdn_checkpoint_updated_states + + return outputs + + +# ============================================================ +# State Dict Converter (Dense -- no MoE weight handling) +# ============================================================ + + +_QWEN36_FP8_DTYPES = tuple( + dtype + for dtype in ( + getattr(torch, "float8_e4m3fn", None), + getattr(torch, "float8_e5m2", None), + ) + if dtype is not None +) + + +def _qwen36_cat(tensors, dim=0): + """Concatenate tensors, including FP8 tensors on builds without FP8 cat.""" + if tensors and tensors[0].dtype in _QWEN36_FP8_DTYPES: + return torch.cat( + [tensor.contiguous().view(torch.int8) for tensor in tensors], + dim=dim, + ).view(tensors[0].dtype) + return torch.cat(tensors, dim=dim) + + +def convert_qwen35_hf_to_neuron_state_dict(neuron_state_dict, config): + """Convert HF Qwen3.5/3.6-27B weights to NxDI format. + + Weight mappings per layer type: + + DeltaNet layers (linear_attention): + HF: layers.X.linear_attn.{in_proj_qkv, in_proj_z, in_proj_a, in_proj_b, + conv1d, A_log, dt_bias, norm, out_proj} + NxDI: projections keep names; conv1d/A_log/dt_bias are remapped into + ColumnParallelLinear parameter containers so NxD can shard them. + + Full attention layers: + HF: layers.X.self_attn.q_proj.weight: (12288, 5120) -- doubled for gate + NxDI: layers.X.self_attn.Wqkv.weight (fused Q+K+V, gate separated) + layers.X.self_attn.output_gate_proj.weight (gate part) + HF: layers.X.self_attn.{k_proj, v_proj, o_proj, q_norm, k_norm} + NxDI: layers.X.self_attn.{..., q_layernorm, k_layernorm} + + Dense MLP (all layers): + HF: layers.X.mlp.{gate_proj, up_proj, down_proj}.weight + NxDI: layers.X.mlp.{gate_proj, up_proj, down_proj}.weight (same names) + + FP8 quantized checkpoints carry one scale tensor next to each quantized + weight. NxDI normalizes saved ``.weight_scale`` keys to model ``.scale`` + keys before this converter runs, so any Qwen-specific weight split/reorder/ + fusion below must apply the same transformation to the matching scale. + """ + # Add rank_util + neuron_state_dict["rank_util.rank"] = torch.arange( + 0, + config.neuron_config.tp_degree, + dtype=torch.int32, + ) + + # ── Qwen3.5-MoE weight renames (only fires for MoE variants) ────────── + # HF stores per-layer MoE weights as: + # mlp.gate.weight — router (num_experts, hidden) + # mlp.experts.gate_up_proj — routed (num_experts, 2*I, H) + # mlp.experts.down_proj — routed (num_experts, H, I) + # mlp.shared_expert.{gate,up,down}_proj.weight + # mlp.shared_expert_gate.weight + # Our Qwen35MoEBlock uses: + # mlp.moe.router.linear_router.weight + # mlp.moe.expert_mlps.mlp_op.gate_up_proj.weight + # mlp.moe.expert_mlps.mlp_op.down_proj.weight + # mlp.shared_{gate,up,down}_proj.weight + # mlp.shared_expert_gate.weight (unchanged) + if getattr(config, "_is_moe", False): + for l in range(config.num_hidden_layers): + p = f"layers.{l}." + # Transpose stacked-expert weights: HF stores + # gate_up_proj: (num_experts, 2*I, H) + # down_proj: (num_experts, H, I) + # NxDI's ExpertMLPsV2 stores + # gate_up_proj: (num_experts, H, 2*I) + # down_proj: (num_experts, I, H) + gu_key = p + "mlp.experts.gate_up_proj" + if gu_key in neuron_state_dict: + w = neuron_state_dict[gu_key] + if w.ndim == 3: + neuron_state_dict[gu_key] = w.transpose(1, 2).contiguous() + dp_key = p + "mlp.experts.down_proj" + if dp_key in neuron_state_dict: + w = neuron_state_dict[dp_key] + if w.ndim == 3: + neuron_state_dict[dp_key] = w.transpose(1, 2).contiguous() + + renames = { + p + "mlp.gate.weight": + p + "mlp.moe.router.linear_router.weight", + p + "mlp.experts.gate_up_proj": + p + "mlp.moe.expert_mlps.mlp_op.gate_up_proj.weight", + p + "mlp.experts.down_proj": + p + "mlp.moe.expert_mlps.mlp_op.down_proj.weight", + p + "mlp.shared_expert.gate_proj.weight": + p + "mlp.shared_gate_proj.weight", + p + "mlp.shared_expert.up_proj.weight": + p + "mlp.shared_up_proj.weight", + p + "mlp.shared_expert.down_proj.weight": + p + "mlp.shared_down_proj.weight", + # shared_expert_gate name is already correct + } + for old, new in renames.items(): + if old in neuron_state_dict: + neuron_state_dict[new] = neuron_state_dict.pop(old) + + def _reorder_deltanet_qkv_for_tp(qkv_weight: torch.Tensor) -> torch.Tensor: + """Pack [Q_all | K_all | V_all] into per-rank Q/K/V blocks. + + ColumnParallelLinear slices the first dimension contiguously. DeltaNet + needs each rank to receive its local query, key, and value heads + together, so the full HF tensor is repacked as: + [rank0 Q | rank0 K | rank0 V | rank1 Q | rank1 K | rank1 V | ...]. + """ + tp_degree = config.neuron_config.tp_degree + num_k_heads = config.linear_num_key_heads + num_v_heads = config.linear_num_value_heads + head_k_dim = config.linear_key_head_dim + head_v_dim = config.linear_value_head_dim + if num_k_heads % tp_degree != 0: + raise ValueError( + f"linear_num_key_heads={num_k_heads} must be divisible by tp_degree={tp_degree}" + ) + if num_v_heads % tp_degree != 0: + raise ValueError( + f"linear_num_value_heads={num_v_heads} must be divisible by tp_degree={tp_degree}" + ) + + key_dim = num_k_heads * head_k_dim + value_dim = num_v_heads * head_v_dim + q_weight = qkv_weight[:key_dim].reshape(num_k_heads, head_k_dim, -1) + k_weight = qkv_weight[key_dim : 2 * key_dim].reshape(num_k_heads, head_k_dim, -1) + v_weight = qkv_weight[2 * key_dim : 2 * key_dim + value_dim].reshape( + num_v_heads, head_v_dim, -1 + ) + local_k_heads = num_k_heads // tp_degree + local_v_heads = num_v_heads // tp_degree + blocks = [] + for rank in range(tp_degree): + blocks.append( + q_weight[ + rank * local_k_heads : (rank + 1) * local_k_heads + ].reshape(-1, qkv_weight.shape[1]) + ) + blocks.append( + k_weight[ + rank * local_k_heads : (rank + 1) * local_k_heads + ].reshape(-1, qkv_weight.shape[1]) + ) + blocks.append( + v_weight[ + rank * local_v_heads : (rank + 1) * local_v_heads + ].reshape(-1, qkv_weight.shape[1]) + ) + return _qwen36_cat(blocks, dim=0).contiguous() + + def _reorder_deltanet_qkv_channels_for_tp(channel_tensor: torch.Tensor) -> torch.Tensor: + """Repack a first-dimension Q/K/V channel tensor into TP rank blocks.""" + tp_degree = config.neuron_config.tp_degree + num_k_heads = config.linear_num_key_heads + num_v_heads = config.linear_num_value_heads + head_k_dim = config.linear_key_head_dim + head_v_dim = config.linear_value_head_dim + key_dim = num_k_heads * head_k_dim + value_dim = num_v_heads * head_v_dim + q_tensor = channel_tensor[:key_dim] + k_tensor = channel_tensor[key_dim : 2 * key_dim] + v_tensor = channel_tensor[2 * key_dim : 2 * key_dim + value_dim] + local_key_dim = key_dim // tp_degree + local_value_dim = value_dim // tp_degree + blocks = [] + for rank in range(tp_degree): + blocks.append(q_tensor[rank * local_key_dim : (rank + 1) * local_key_dim]) + blocks.append(k_tensor[rank * local_key_dim : (rank + 1) * local_key_dim]) + blocks.append( + v_tensor[rank * local_value_dim : (rank + 1) * local_value_dim] + ) + return _qwen36_cat(blocks, dim=0).contiguous() + + def _split_interleaved_q_proj_tensor( + tensor: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Split interleaved Qwen q_proj tensor into query and output gate.""" + num_heads = config.num_attention_heads + head_dim = config.head_dim + trailing_shape = tensor.shape[1:] + tensor = tensor.reshape(num_heads, head_dim * 2, *trailing_shape) + query_tensor = tensor[:, :head_dim, ...].reshape( + num_heads * head_dim, + *trailing_shape, + ) + gate_tensor = tensor[:, head_dim:, ...].reshape( + num_heads * head_dim, + *trailing_shape, + ) + return query_tensor.contiguous(), gate_tensor.contiguous() + + # CRITICAL: Convert (1+weight) RMSNorm weights to standard RMSNorm weights. + # Qwen3.5 uses RMSNorm with `output = norm(x) * (1 + weight)` where weight + # is initialized to zeros. Standard NxDI RMSNorm uses `output = norm(x) * weight` + # where weight is initialized to ones. To convert: new_weight = old_weight + 1.0 + norm_keys_to_convert = [] + for l in range(config.num_hidden_layers): + norm_keys_to_convert.append(f"layers.{l}.input_layernorm.weight") + norm_keys_to_convert.append(f"layers.{l}.post_attention_layernorm.weight") + if config.layer_types[l] == "full_attention": + norm_keys_to_convert.append(f"layers.{l}.self_attn.q_norm.weight") + norm_keys_to_convert.append(f"layers.{l}.self_attn.k_norm.weight") + norm_keys_to_convert.append("norm.weight") + + for nk in norm_keys_to_convert: + if nk in neuron_state_dict: + old_val = neuron_state_dict[nk] + neuron_state_dict[nk] = old_val.float() + 1.0 + if "layers.0." in nk or nk == "norm.weight": + logger.debug( + f"[NORM FIX] {nk}: mean {old_val.float().mean():.4f} -> {neuron_state_dict[nk].mean():.4f}" + ) + else: + if "layers.0." in nk or nk == "norm.weight": + logger.warning(f"[NORM FIX] key not found: {nk}") + + for l in range(config.num_hidden_layers): + layer_type = config.layer_types[l] + + # === DeltaNet layers === + if layer_type == "linear_attention": + qkv_key = f"layers.{l}.linear_attn.in_proj_qkv.weight" + if qkv_key in neuron_state_dict and config.neuron_config.tp_degree > 1: + neuron_state_dict[qkv_key] = _reorder_deltanet_qkv_for_tp( + neuron_state_dict[qkv_key] + ) + qkv_scale_key = f"layers.{l}.linear_attn.in_proj_qkv.scale" + if qkv_scale_key in neuron_state_dict and config.neuron_config.tp_degree > 1: + neuron_state_dict[qkv_scale_key] = _reorder_deltanet_qkv_channels_for_tp( + neuron_state_dict[qkv_scale_key] + ) + + conv_key = f"layers.{l}.linear_attn.conv1d.weight" + conv_weight_key = f"layers.{l}.linear_attn.conv1d_weight.weight" + conv_scale_key = f"layers.{l}.linear_attn.conv1d.scale" + conv_weight_scale_key = f"layers.{l}.linear_attn.conv1d_weight.scale" + if conv_key in neuron_state_dict: + conv_weight = neuron_state_dict.pop(conv_key) + if config.neuron_config.tp_degree > 1: + conv_weight = _reorder_deltanet_qkv_channels_for_tp(conv_weight) + neuron_state_dict[conv_weight_key] = conv_weight.squeeze(1).contiguous() + if conv_scale_key in neuron_state_dict: + conv_scale = neuron_state_dict.pop(conv_scale_key) + if config.neuron_config.tp_degree > 1: + conv_scale = _reorder_deltanet_qkv_channels_for_tp(conv_scale) + neuron_state_dict[conv_weight_scale_key] = conv_scale.contiguous() + + for vector_name in ("A_log", "dt_bias"): + vector_key = f"layers.{l}.linear_attn.{vector_name}" + vector_weight_key = f"layers.{l}.linear_attn.{vector_name}_weight.weight" + if vector_key in neuron_state_dict: + neuron_state_dict[vector_weight_key] = ( + neuron_state_dict.pop(vector_key).reshape(-1, 1).contiguous() + ) + + # === Attention layers === + if layer_type == "full_attention": + neuron_state_dict[f"layers.{l}.self_attn.rank_util.rank"] = torch.arange( + 0, + config.neuron_config.tp_degree, + dtype=torch.int32, + ) + + # QK norms: q_norm -> q_layernorm, k_norm -> k_layernorm + q_norm_key = f"layers.{l}.self_attn.q_norm.weight" + k_norm_key = f"layers.{l}.self_attn.k_norm.weight" + if q_norm_key in neuron_state_dict: + neuron_state_dict[f"layers.{l}.self_attn.q_layernorm.weight"] = ( + neuron_state_dict.pop(q_norm_key).detach().clone() + ) + if k_norm_key in neuron_state_dict: + neuron_state_dict[f"layers.{l}.self_attn.k_layernorm.weight"] = ( + neuron_state_dict.pop(k_norm_key).detach().clone() + ) + + # q_proj is doubled: (12288, 5120) = (num_heads * head_dim * 2, hidden) + # INTERLEAVED: [head0_query(256) | head0_gate(256) | head1_query(256) | ...] + q_proj_key = f"layers.{l}.self_attn.q_proj.weight" + q_proj_scale_key = f"layers.{l}.self_attn.q_proj.scale" + if q_proj_key in neuron_state_dict: + q_proj_w = neuron_state_dict.pop(q_proj_key) + query_w, gate_w = _split_interleaved_q_proj_tensor(q_proj_w) + + neuron_state_dict[q_proj_key] = query_w + neuron_state_dict[f"layers.{l}.self_attn.output_gate_proj.weight"] = ( + gate_w + ) + if q_proj_scale_key in neuron_state_dict: + q_proj_scale = neuron_state_dict.pop(q_proj_scale_key) + query_scale, gate_scale = _split_interleaved_q_proj_tensor( + q_proj_scale + ) + neuron_state_dict[q_proj_scale_key] = query_scale + neuron_state_dict[f"layers.{l}.self_attn.output_gate_proj.scale"] = ( + gate_scale + ) + + # Fuse QKV + if config.neuron_config.fused_qkv: + q_key = f"layers.{l}.self_attn.q_proj.weight" + k_key = f"layers.{l}.self_attn.k_proj.weight" + v_key = f"layers.{l}.self_attn.v_proj.weight" + gate_key = f"layers.{l}.self_attn.output_gate_proj.weight" + pack_gate_in_qkv = bool( + getattr(config, "use_qwen_qkv_gate_packed", False) + ) + if q_key in neuron_state_dict: + qkv_weight_parts = [neuron_state_dict[q_key]] + if pack_gate_in_qkv: + if gate_key not in neuron_state_dict: + raise ValueError( + f"Missing output-gate tensor for packed QKV: {gate_key}" + ) + qkv_weight_parts.append(neuron_state_dict[gate_key]) + qkv_weight_parts.extend( + [neuron_state_dict[k_key], neuron_state_dict[v_key]] + ) + neuron_state_dict[f"layers.{l}.self_attn.Wqkv.weight"] = _qwen36_cat( + qkv_weight_parts + ) + q_scale_key = f"layers.{l}.self_attn.q_proj.scale" + gate_scale_key = f"layers.{l}.self_attn.output_gate_proj.scale" + k_scale_key = f"layers.{l}.self_attn.k_proj.scale" + v_scale_key = f"layers.{l}.self_attn.v_proj.scale" + scale_keys = [q_scale_key] + if pack_gate_in_qkv: + scale_keys.append(gate_scale_key) + scale_keys.extend([k_scale_key, v_scale_key]) + scale_keys_present = [key in neuron_state_dict for key in scale_keys] + if any(scale_keys_present): + if not all(scale_keys_present): + missing = [ + key + for key, present in zip(scale_keys, scale_keys_present) + if not present + ] + raise ValueError( + f"Missing FP8 fused-QKV scale tensor(s): {missing}" + ) + neuron_state_dict[f"layers.{l}.self_attn.Wqkv.scale"] = _qwen36_cat( + [neuron_state_dict[key] for key in scale_keys] + ) + del neuron_state_dict[q_scale_key] + del neuron_state_dict[k_scale_key] + del neuron_state_dict[v_scale_key] + del neuron_state_dict[q_key] + del neuron_state_dict[k_key] + del neuron_state_dict[v_key] + + # Dense MLP: no weight conversion needed -- HF and NxDI use same names + # HF: layers.X.mlp.{gate_proj, up_proj, down_proj}.weight + # NxDI: layers.X.mlp.{gate_proj, up_proj, down_proj}.weight + + gc.collect() + + return neuron_state_dict + + +# ============================================================ +# Custom ModelWrapper and DecoderModelInstance for DeltaNet state aliasing +# ============================================================ + + +def _reassert_hybrid_gdn_checkpoint_param_dtypes(module): + config = getattr(module, "config", None) + if config is None: + return + + recurrent_dtype = _torch_dtype_from_hybrid_cache_dtype( + getattr(config, "hybrid_recurrent_cache_dtype", "float32") + ) + conv_dtype = _torch_dtype_from_hybrid_cache_dtype( + getattr(config, "hybrid_conv_cache_dtype", "bfloat16") + ) + + def _retarget(params, dtype): + for param in params: + if param.dtype != dtype: + param.data = param.data.to(dtype) + + for layer in getattr(module, "layers", []): + linear_attn = getattr(layer, "linear_attn", None) + if linear_attn is None: + continue + recurrent_buffer = getattr(linear_attn, "recurrent_state_buffer", None) + conv_buffer = getattr(linear_attn, "conv_state_buffer", None) + if recurrent_buffer is not None and recurrent_buffer.dtype != recurrent_dtype: + recurrent_buffer.data = recurrent_buffer.data.to(recurrent_dtype) + if conv_buffer is not None and conv_buffer.dtype != conv_dtype: + conv_buffer.data = conv_buffer.data.to(conv_dtype) + + cache = getattr(module, "hybrid_gdn_checkpoint_cache", None) + if cache is not None: + _retarget(cache.recurrent_slots, recurrent_dtype) + _retarget(cache.conv_slots, conv_dtype) + cache.recurrent_dtype = recurrent_dtype + cache.conv_dtype = conv_dtype + + +def _qwen36_is_context_encoding_trace( + n_active_tokens: int | None, + neuron_config, +) -> bool: + n_active_tokens = int(n_active_tokens or 0) + speculation_length = getattr(neuron_config, "speculation_length", None) + return n_active_tokens != 1 and not ( + speculation_length is not None and n_active_tokens == speculation_length + ) + + +def _qwen36_include_hybrid_gdn_checkpoint_outputs( + config, + *, + is_for_context_encoding: bool | None = None, + n_active_tokens: int | None = None, + neuron_config=None, +) -> bool: + if is_for_context_encoding is None: + is_for_context_encoding = _qwen36_is_context_encoding_trace( + n_active_tokens, + neuron_config, + ) + if not getattr(config, "use_hybrid_apc_manager", False): + return True + if is_for_context_encoding: + return True + return bool(getattr(config, "hybrid_apc_commit_during_token_generation", False)) + + +def _qwen36_validate_alias_output_counts( + module, + updated_kv_cache, + *, + is_for_context_encoding: bool, +): + kv_mgr = getattr(module, "kv_mgr", None) + if kv_mgr is not None: + expected_kv = len(kv_mgr.past_key_values) + else: + expected_kv = 0 + actual_kv = len(updated_kv_cache) + if actual_kv != expected_kv: + raise RuntimeError( + "Qwen3.6 output alias count mismatch: " + f"updated_kv_cache has {actual_kv} tensors but kv_mgr.past_key_values " + f"has {expected_kv}" + ) + + expected_states = 0 + if not getattr(module.config, "use_hybrid_cache_manager", False): + expected_states = len(getattr(module, "_deltanet_state_params", [])) + actual_states = len(getattr(module, "_deltanet_updated_states", [])) + if actual_states != expected_states: + raise RuntimeError( + "Qwen3.6 output alias count mismatch: " + f"_deltanet_updated_states has {actual_states} tensors but " + f"_deltanet_state_params has {expected_states}" + ) + + checkpoint_outputs_expected = _qwen36_include_hybrid_gdn_checkpoint_outputs( + module.config, + is_for_context_encoding=is_for_context_encoding, + ) + expected_checkpoints = ( + len(getattr(module, "_hybrid_gdn_checkpoint_params", [])) + if checkpoint_outputs_expected + else 0 + ) + actual_checkpoints = len( + getattr(module, "_hybrid_gdn_checkpoint_updated_states", []) + ) + if actual_checkpoints != expected_checkpoints: + raise RuntimeError( + "Qwen3.6 output alias count mismatch: " + f"_hybrid_gdn_checkpoint_updated_states has {actual_checkpoints} tensors " + f"but _hybrid_gdn_checkpoint_params expects {expected_checkpoints}" + ) + + +class Qwen35DecoderModelInstance(DecoderModelInstance): + """Custom DecoderModelInstance that adds DeltaNet state buffers to input_output_aliases.""" + + def load_module(self): + super().load_module() + _reassert_hybrid_gdn_checkpoint_param_dtypes(self.module) + + @staticmethod + def _num_trace_outputs_before_aliases(neuron_config): + if ( + getattr(neuron_config, "output_logits", False) + and getattr(neuron_config, "on_device_sampling_config", None) is not None + ): + return 2 + return 1 + + def get(self, bucket_rank, **kwargs): + """Override to add DeltaNet state aliases after KV cache aliases.""" + module, input_output_aliases = super().get(bucket_rank, **kwargs) + + num_output_from_trace = self._num_trace_outputs_before_aliases( + self.neuron_config + ) + base_num_output_from_trace = 1 if not self.neuron_config.output_logits else 2 + if num_output_from_trace != base_num_output_from_trace: + alias_shift = base_num_output_from_trace - num_output_from_trace + for param in list(input_output_aliases.keys()): + input_output_aliases[param] -= alias_shift + + if module.kv_mgr is not None: + num_kv = len(module.kv_mgr.past_key_values) + else: + num_kv = 0 + + state_start_idx = num_output_from_trace + num_kv + + if ( + not getattr(module.config, "use_hybrid_cache_manager", False) + and hasattr(module, "_deltanet_state_params") + ): + for i, param in enumerate(module._deltanet_state_params): + input_output_aliases[param] = state_start_idx + i + + checkpoint_start_idx = state_start_idx + len(module._deltanet_state_params) + include_checkpoint_aliases = _qwen36_include_hybrid_gdn_checkpoint_outputs( + module.config, + n_active_tokens=getattr(module, "n_active_tokens", 0), + neuron_config=self.neuron_config, + ) + if include_checkpoint_aliases: + for i, param in enumerate( + getattr(module, "_hybrid_gdn_checkpoint_params", []) + ): + input_output_aliases[param] = checkpoint_start_idx + i + + return module, input_output_aliases + + +class Qwen35ModelWrapper(ModelWrapper): + """Custom ModelWrapper for VL support with mRoPE and vision inputs.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._qwen36_hybrid_apc_pending_input_dict = None + self.hybrid_apc_store = None + self.hybrid_apc_slot_allocator = None + self.hybrid_apc_bridge = None + self._init_hybrid_apc_scheduler_bridge() + + def _init_hybrid_apc_scheduler_bridge(self): + if not _qwen36_config_flag( + self.config, + self.neuron_config, + "use_hybrid_apc_manager", + ): + return + + required_gdn_layers = tuple( + idx + for idx, layer_type in enumerate(self.config.layer_types) + if layer_type == "linear_attention" + ) + if not required_gdn_layers: + raise ValueError("hybrid APC requires at least one GDN layer") + + tp_rank = 0 + try: + if torch.distributed.is_available() and torch.distributed.is_initialized(): + tp_rank = int(parallel_state.get_tensor_model_parallel_rank()) + except Exception: + tp_rank = 0 + + block_size = int( + getattr( + self.neuron_config, + "pa_block_size", + self.config.gdn_checkpoint_interval, + ) + ) + self.hybrid_apc_store = HybridAPCMetadataStore( + required_gdn_layers=required_gdn_layers, + block_size=block_size, + checkpoint_interval=self.config.gdn_checkpoint_interval, + max_checkpoints=self.config.max_gdn_checkpoint_slots, + layout_version=self.config.hybrid_apc_layout_version, + model_revision=self.config.hybrid_apc_model_revision, + tp_rank=tp_rank, + recurrent_dtype=self.config.hybrid_recurrent_cache_dtype, + conv_dtype=self.config.hybrid_conv_cache_dtype, + allow_residual_replay=self.config.hybrid_apc_allow_residual_replay, + ) + self.hybrid_apc_slot_allocator = HybridAPCSlotAllocator( + self.config.max_gdn_checkpoint_slots + ) + self.hybrid_apc_bridge = HybridAPCSchedulerBridge( + store=self.hybrid_apc_store, + slot_allocator=self.hybrid_apc_slot_allocator, + cache_salt=self.config.hybrid_apc_cache_salt, + model_revision=self.config.hybrid_apc_model_revision, + layout_version=self.config.hybrid_apc_layout_version, + tp_rank=tp_rank, + recurrent_dtype=self.config.hybrid_recurrent_cache_dtype, + conv_dtype=self.config.hybrid_conv_cache_dtype, + allow_local_hash_fallback=self.config.hybrid_apc_allow_local_hash_fallback, + require_attention_block_refs=self.config.hybrid_apc_require_attention_block_refs, + reject_unbacked_attention_hits=( + self.config.hybrid_apc_reject_unbacked_attention_hits + ), + ) + + def ensure_hybrid_apc_scheduler_bridge(self): + if not _qwen36_config_flag( + self.config, + self.neuron_config, + "use_hybrid_apc_manager", + ): + return None + if getattr(self, "hybrid_apc_bridge", None) is None: + self._init_hybrid_apc_scheduler_bridge() + return self.hybrid_apc_bridge + + def get_model_instance(self): + return Qwen35DecoderModelInstance( + model_cls=self.model_cls, + config=self.config, + **self.model_init_kwargs, + ) + + def input_generator(self): + """Generate inputs including mrope_position_ids, vision_embeddings, and vision_mask.""" + base_inputs = super().input_generator() + extended_inputs = [] + + for bucket_inputs in base_inputs: + input_ids = bucket_inputs[0] + batch_size = input_ids.shape[0] + n_active_tokens = input_ids.shape[1] + + is_cte = self.tag == CONTEXT_ENCODING_MODEL_TAG + + if is_cte: + mrope_position_ids = ( + torch.arange(0, n_active_tokens, dtype=torch.int32) + .unsqueeze(0) + .unsqueeze(0) + .expand(3, batch_size, -1) + .contiguous() + ) + + if getattr(self.config, "use_text_only_cte_inputs", True): + vision_embeddings = torch.zeros( + (0,), dtype=self.config.neuron_config.torch_dtype + ) + vision_mask = torch.zeros((0,), dtype=torch.int32) + else: + vision_embeddings = torch.zeros( + (batch_size, n_active_tokens, self.config.hidden_size), + dtype=self.config.neuron_config.torch_dtype, + ) + vision_mask = torch.full( + (batch_size, n_active_tokens, 1), + fill_value=n_active_tokens - 1, + dtype=torch.int32, + ) + else: + mrope_position_ids = torch.zeros((0,), dtype=torch.int32) + vision_embeddings = torch.zeros( + (0,), dtype=self.config.neuron_config.torch_dtype + ) + vision_mask = torch.zeros((0,), dtype=torch.int32) + + hybrid_args = None + if _use_expanded_hybrid_args_for_tag(self.config, self.tag): + hybrid_args = ( + torch.zeros((batch_size,), dtype=torch.int32), + torch.zeros((batch_size,), dtype=torch.int32), + torch.zeros((batch_size,), dtype=torch.int32), + torch.zeros((batch_size,), dtype=torch.int32), + torch.zeros((batch_size,), dtype=torch.int32), + ) + + if is_cte: + padded = build_cte_args( + self.config, + bucket_inputs, + mrope_position_ids, + vision_embeddings, + vision_mask, + hybrid_args=hybrid_args, + ) + else: + padded = build_tkg_args( + self.config, + bucket_inputs, + mrope_position_ids, + vision_embeddings, + vision_mask, + hybrid_args=hybrid_args, + ) + _debug_qwen36_arg_contract( + "compile", + self.tag, + self.config, + padded, + ) + extended_inputs.append(tuple(padded)) + + return extended_inputs + + def _prepare_hybrid_apc_pad_inputs(self, args): + if ( + self.tag != CONTEXT_ENCODING_MODEL_TAG + or len(args) < 29 + or not _qwen36_config_flag( + self.config, + self.neuron_config, + "use_hybrid_apc_manager", + ) + or _qwen36_hybrid_apc_controls_materialized( + args[25], + args[26], + args[28], + ) + ): + return args + + computed_context_lens = args[14] + num_queries = args[13] + full_context_lens = ( + computed_context_lens + num_queries + if hasattr(computed_context_lens, "shape") and hasattr(num_queries, "shape") + else None + ) + hybrid_apc_request_dict = { + "input_ids": args[0], + "attention_mask": args[1], + "position_ids": args[2], + "seq_ids": args[3], + "sampling_params": args[4], + "adapter_ids": args[6], + "slot_mapping": args[11], + "block_table": args[12], + "num_queries": num_queries, + "computed_context_lens": computed_context_lens, + } + if full_context_lens is not None: + hybrid_apc_request_dict["full_context_lens"] = full_context_lens + + request_records = getattr( + self, + "_qwen36_vllm_hybrid_apc_request_records", + None, + ) + request_ids = _qwen36_request_ids_from_hybrid_apc_records(request_records) + if request_records is not None: + hybrid_apc_request_dict["hybrid_request_records"] = request_records + if request_ids is None: + request_ids = getattr(self, "_qwen36_vllm_request_ids", None) + if request_ids is not None: + if isinstance(request_ids, list): + request_ids = tuple(request_ids) + if isinstance(request_ids, tuple) and len(request_ids) == 1: + hybrid_apc_request_dict["hybrid_request_id"] = request_ids[0] + else: + hybrid_apc_request_dict["hybrid_request_id"] = request_ids + cached_request_ids = getattr(self, "_qwen36_vllm_cached_request_ids", None) + if cached_request_ids is not None: + hybrid_apc_request_dict["hybrid_cached_request_ids"] = cached_request_ids + prefill_completion_state = getattr( + self, + "_qwen36_vllm_prefill_completion_state", + None, + ) + if prefill_completion_state is not None: + hybrid_apc_request_dict[ + "hybrid_prefill_completion_state" + ] = prefill_completion_state + _qwen36_add_vllm_hybrid_apc_metadata( + hybrid_apc_request_dict, + request_ids=request_ids, + metadata_by_request_id=getattr( + self, + "_qwen36_vllm_hybrid_apc_metadata_by_request_id", + None, + ), + ) + + prepared_inputs = prepare_hybrid_apc_request_for_execution( + self, + hybrid_apc_request_dict, + ) + hybrid_args = prepare_hybrid_apc_model_inputs(self, prepared_inputs) + if not hybrid_args: + return args + + updated_args = list(args) + for index, key in ( + (0, "input_ids"), + (1, "attention_mask"), + (2, "position_ids"), + (3, "seq_ids"), + (4, "sampling_params"), + (6, "adapter_ids"), + (11, "slot_mapping"), + (12, "block_table"), + (13, "num_queries"), + (14, "computed_context_lens"), + ): + if key in prepared_inputs: + updated_args[index] = prepared_inputs[key] + if len(hybrid_args) == 14 and len(updated_args) >= 29: + updated_args[15:29] = hybrid_args + else: + updated_args[24:29] = hybrid_args + self._qwen36_hybrid_apc_pending_input_dict = hybrid_apc_request_dict + return tuple(updated_args) + + def _forward_with_pad(self, *args): + self._qwen36_hybrid_apc_pending_input_dict = None + try: + outputs = super()._forward_with_pad(*args) + except Exception: + pending = self._qwen36_hybrid_apc_pending_input_dict + self._qwen36_hybrid_apc_pending_input_dict = None + if pending is not None: + cancel_hybrid_apc_request(pending) + raise + pending = self._qwen36_hybrid_apc_pending_input_dict + self._qwen36_hybrid_apc_pending_input_dict = None + if pending is not None: + finish_hybrid_apc_request(pending) + return outputs + + def pad_inputs(self, *args, pad_type="first_fit"): + """Override to pad mrope_position_ids and vision inputs to bucket size.""" + args = self._prepare_hybrid_apc_pad_inputs(args) + if ( + self.tag in (CONTEXT_ENCODING_MODEL_TAG, TOKEN_GENERATION_MODEL_TAG) + and len(args) == 15 + and self.is_prefix_caching + and not getattr( + getattr(self, "neuron_config", None), + "enable_fused_speculation", + False, + ) + and not getattr( + getattr(self, "neuron_config", None), + "enable_eagle_speculation", + False, + ) + ): + args = tuple( + _normalize_qwen36_prefix_args(args) + + [_empty_qwen36_arg(), _empty_qwen36_arg(), _empty_qwen36_arg()] + ) + orig_mrope = args[21] if len(args) >= 22 else None + orig_vis_emb = args[22] if len(args) >= 23 else None + orig_vis_mask = args[23] if len(args) >= 24 else None + if len(args) >= 29: + orig_restore_slots = args[24] + orig_restore_mask = args[25] + orig_restore_prefix = args[26] + orig_commit_slots = args[27] + orig_commit_mask = args[28] + elif ( + len(args) >= 20 + and _use_expanded_hybrid_args_for_tag(self.config, self.tag) + and self.is_prefix_caching + and not self.neuron_config.enable_fused_speculation + and not self.neuron_config.enable_eagle_speculation + ): + orig_restore_slots = args[15] + orig_restore_mask = args[16] + orig_restore_prefix = args[17] + orig_commit_slots = args[18] + orig_commit_mask = args[19] + else: + orig_restore_slots = None + orig_restore_mask = None + orig_restore_prefix = None + orig_commit_slots = None + orig_commit_mask = None + + # Pre-pad/truncate vision args to match the target CTE bucket so the + # upstream super().pad_inputs() shape check at model_wrapper.py:801 + # does not replace them with dummies. Bucket is determined by input + # length; we compute it the same way upstream does. + if ( + self.tag == CONTEXT_ENCODING_MODEL_TAG + and len(args) >= 24 + and orig_vis_mask is not None + and orig_vis_mask.ndim == 3 + ): + try: + target_bucket = self.get_target_bucket(*args, strategy=pad_type) + if isinstance(target_bucket, list): + target_bucket = target_bucket[1] + target_len = int(target_bucket) + except Exception: + target_len = None + if target_len is not None: + def _fit_seq(t, target, fill_dim0=False, fill_value=None): + if t is None or t.ndim != 3: + return t + cur = t.shape[1] + if cur == target: + return t + if cur < target: + pad_shape = list(t.shape) + pad_shape[1] = target - cur + if fill_value is not None: + pad = torch.full(pad_shape, fill_value=fill_value, dtype=t.dtype) + else: + pad = torch.zeros(pad_shape, dtype=t.dtype) + return torch.cat([t, pad], dim=1) + return t[:, :target].contiguous() + + new_vis_emb = _fit_seq(orig_vis_emb, target_len) + new_vis_mask = _fit_seq(orig_vis_mask, target_len, fill_value=target_len - 1) + if new_vis_emb is not None or new_vis_mask is not None: + args = list(args) + if new_vis_emb is not None: + args[22] = new_vis_emb + if new_vis_mask is not None: + args[23] = new_vis_mask + args = tuple(args) + + padded_args = super().pad_inputs(*args, pad_type=pad_type) + + if len(padded_args) >= 24 and orig_mrope is not None: + padded_seq_len = padded_args[0].shape[1] + batch_size = padded_args[0].shape[0] + is_cte = self.tag == CONTEXT_ENCODING_MODEL_TAG + + if is_cte: + current_mrope = orig_mrope + current_vis_emb = orig_vis_emb + current_vis_mask = orig_vis_mask + + if ( + current_mrope.ndim == 3 + and current_mrope.shape[-1] < padded_seq_len + ): + pad_size = padded_seq_len - current_mrope.shape[-1] + last_pos = current_mrope[:, :, -1:] + # Padded tokens are masked out of the active CTE, so do not + # advance mRoPE into fake future positions. + mrope_pad = last_pos.expand(3, batch_size, pad_size) + mrope_position_ids = torch.cat([current_mrope, mrope_pad], dim=-1) + elif ( + current_mrope.ndim == 3 + and current_mrope.shape[-1] > padded_seq_len + ): + # Bucket smaller than caller-provided mrope; truncate. + mrope_position_ids = current_mrope[:, :, :padded_seq_len].contiguous() + elif current_mrope.ndim == 3: + mrope_position_ids = current_mrope + else: + mrope_position_ids = ( + torch.arange(0, padded_seq_len, dtype=torch.int32) + .unsqueeze(0) + .unsqueeze(0) + .expand(3, batch_size, -1) + .contiguous() + ) + + if ( + current_vis_emb is not None + and current_vis_emb.ndim == 3 + and current_vis_emb.shape[1] < padded_seq_len + ): + # Qwen3-VL pad convention: pad slots of vision_embeddings + # are zeros; pad slots of vision_mask (below) point at + # padded_seq_len-1 which is guaranteed to be a padded + # (attention_mask==0) input position — so scatter writes + # zero to a masked slot with no downstream effect. + pad_len = padded_seq_len - current_vis_emb.shape[1] + pad_emb = torch.zeros( + (batch_size, pad_len, current_vis_emb.shape[2]), + dtype=current_vis_emb.dtype, + ) + vision_embeddings = torch.cat([current_vis_emb, pad_emb], dim=1) + elif current_vis_emb is not None and current_vis_emb.ndim == 3: + vision_embeddings = current_vis_emb[:, :padded_seq_len] + elif getattr(self.config, "use_text_only_cte_inputs", True): + vision_embeddings = torch.zeros( + (0,), dtype=self.config.neuron_config.torch_dtype + ) + else: + # Dummy vision inputs for text-only calls when graph was + # traced with vision inputs. Zeros are fine because mask + # sends them all to the padding-position at padded_seq_len-1. + vision_embeddings = torch.zeros( + (batch_size, padded_seq_len, self.config.hidden_size), + dtype=self.config.neuron_config.torch_dtype, + ) + + if ( + current_vis_mask is not None + and current_vis_mask.ndim == 3 + and current_vis_mask.shape[1] < padded_seq_len + ): + # Qwen3-VL pad convention: pad slots of vision_mask point + # at padded_seq_len-1 (a padded input slot). + pad_len = padded_seq_len - current_vis_mask.shape[1] + pad_mask = torch.full( + (batch_size, pad_len, 1), + fill_value=padded_seq_len - 1, + dtype=torch.int32, + ) + vision_mask = torch.cat([current_vis_mask, pad_mask], dim=1) + elif current_vis_mask is not None and current_vis_mask.ndim == 3: + vision_mask = current_vis_mask[:, :padded_seq_len] + elif getattr(self.config, "use_text_only_cte_inputs", True): + vision_mask = torch.zeros((0,), dtype=torch.int32) + else: + # Dummy mask for text-only forward on a vision-traced graph. + # All slots target padded_seq_len-1 (the last position, always + # a padded/eos slot). Combined with zero vision_emb the scatter + # overwrites that one position with zeros — harmless because + # attention_mask=0 there. + vision_mask = torch.full( + (batch_size, padded_seq_len, 1), + fill_value=padded_seq_len - 1, + dtype=torch.int32, + ) + + padded_args = ( + *padded_args[:21], + mrope_position_ids, + vision_embeddings, + vision_mask, + ) + + if vision_mask.ndim == 3: + padded_args = list(padded_args) + padded_args[23] = padded_args[23].clamp(max=padded_seq_len - 1) + padded_args = tuple(padded_args) + + if ( + len(padded_args) >= 24 + and _use_expanded_hybrid_args_for_tag(self.config, self.tag) + ): + padded_batch_size = padded_args[0].shape[0] + + def _pad_vector(value, dtype=torch.int32): + if value is None or not hasattr(value, "ndim") or value.ndim == 0: + return torch.zeros((padded_batch_size,), dtype=dtype) + value = value.to(dtype) + if value.shape[0] == padded_batch_size: + return value + if value.shape[0] > padded_batch_size: + return value[:padded_batch_size] + pad = torch.zeros( + (padded_batch_size - value.shape[0],), + dtype=value.dtype, + ) + return torch.cat([value, pad], dim=0) + + hybrid_args = ( + _pad_vector(orig_restore_slots), + _pad_vector(orig_restore_mask), + _pad_vector(orig_restore_prefix), + _pad_vector(orig_commit_slots), + _pad_vector(orig_commit_mask), + ) + if len(padded_args) >= 29: + padded_args = (*padded_args[:24], *hybrid_args) + else: + padded_args = (*padded_args, *hybrid_args) + + _assert_qwen36_arg_count( + self.tag, + padded_args, + _qwen36_expected_arg_count(self.config, self.tag), + ) + _debug_qwen36_arg_contract("pad", self.tag, self.config, padded_args) + return padded_args + + +# ============================================================ +# Top-Level Model +# ============================================================ + + +class NeuronQwen35ForCausalLM(NeuronBaseForCausalLM): + _model_cls = NeuronQwen35Model + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._init_hybrid_apc_scheduler_bridge() + + def _init_hybrid_apc_scheduler_bridge(self): + self.hybrid_apc_store = None + self.hybrid_apc_slot_allocator = None + self.hybrid_apc_bridge = None + if not _qwen36_config_flag( + self.config, + self.neuron_config, + "use_hybrid_apc_manager", + ): + return + + required_gdn_layers = tuple( + idx + for idx, layer_type in enumerate(self.config.layer_types) + if layer_type == "linear_attention" + ) + if not required_gdn_layers: + raise ValueError("hybrid APC requires at least one GDN layer") + + tp_rank = 0 + try: + if torch.distributed.is_available() and torch.distributed.is_initialized(): + tp_rank = int(parallel_state.get_tensor_model_parallel_rank()) + except Exception: + tp_rank = 0 + + block_size = int( + getattr( + self.neuron_config, + "pa_block_size", + self.config.gdn_checkpoint_interval, + ) + ) + self.hybrid_apc_store = HybridAPCMetadataStore( + required_gdn_layers=required_gdn_layers, + block_size=block_size, + checkpoint_interval=self.config.gdn_checkpoint_interval, + max_checkpoints=self.config.max_gdn_checkpoint_slots, + layout_version=self.config.hybrid_apc_layout_version, + model_revision=self.config.hybrid_apc_model_revision, + tp_rank=tp_rank, + recurrent_dtype=self.config.hybrid_recurrent_cache_dtype, + conv_dtype=self.config.hybrid_conv_cache_dtype, + allow_residual_replay=self.config.hybrid_apc_allow_residual_replay, + ) + self.hybrid_apc_slot_allocator = HybridAPCSlotAllocator( + self.config.max_gdn_checkpoint_slots + ) + self.hybrid_apc_bridge = HybridAPCSchedulerBridge( + store=self.hybrid_apc_store, + slot_allocator=self.hybrid_apc_slot_allocator, + cache_salt=self.config.hybrid_apc_cache_salt, + model_revision=self.config.hybrid_apc_model_revision, + layout_version=self.config.hybrid_apc_layout_version, + tp_rank=tp_rank, + recurrent_dtype=self.config.hybrid_recurrent_cache_dtype, + conv_dtype=self.config.hybrid_conv_cache_dtype, + allow_local_hash_fallback=self.config.hybrid_apc_allow_local_hash_fallback, + require_attention_block_refs=self.config.hybrid_apc_require_attention_block_refs, + reject_unbacked_attention_hits=( + self.config.hybrid_apc_reject_unbacked_attention_hits + ), + ) + + def ensure_hybrid_apc_scheduler_bridge(self): + if not _qwen36_config_flag( + self.config, + self.neuron_config, + "use_hybrid_apc_manager", + ): + return None + if getattr(self, "hybrid_apc_bridge", None) is None: + self._init_hybrid_apc_scheduler_bridge() + return self.hybrid_apc_bridge + + def on_attention_block_evicted(self, block_ref: int): + if self.hybrid_apc_store is None: + return [] + return self.hybrid_apc_store.on_attention_block_evicted(block_ref) + + def on_attention_blocks_evicted(self, block_refs): + invalidated = [] + if self.hybrid_apc_store is None: + return invalidated + for block_ref in block_refs: + invalidated.extend( + self.hybrid_apc_store.on_attention_block_evicted(block_ref) + ) + return invalidated + + def get_model_wrapper_cls(self): + """Return custom ModelWrapper with DeltaNet state aliasing.""" + return Qwen35ModelWrapper + + @staticmethod + def load_hf_model(model_path, **kwargs): + """Load HF model weights. + + The model is a VL model (Qwen3_5ForConditionalGeneration) but we + only need the text backbone. + """ + from transformers import AutoModelForCausalLM + + kwargs.setdefault("trust_remote_code", True) + return AutoModelForCausalLM.from_pretrained(model_path, **kwargs) + + @classmethod + def get_config_cls(cls): + return Qwen35InferenceConfig + + @staticmethod + def update_state_dict_for_tied_weights(state_dict): + # Qwen3.5-2B has tie_word_embeddings=True. HF only stores + # embed_tokens.weight and derives lm_head from it. NxDI's lm_head is a + # separate ColumnParallelLinear that needs its own weight tensor. + if "lm_head.weight" not in state_dict and "embed_tokens.weight" in state_dict: + state_dict["lm_head.weight"] = state_dict["embed_tokens.weight"].clone() + + @staticmethod + def convert_hf_to_neuron_state_dict(state_dict, config): + """Strip VL wrapper prefix and convert to NxDI format.""" + new_sd = {} + for k, v in state_dict.items(): + if k.startswith("language_model."): + new_k = k.replace("language_model.", "", 1) + new_sd[new_k] = v + elif k.startswith("model.language_model."): + new_k = k.replace("model.language_model.", "", 1) + new_sd[new_k] = v + elif k.startswith("model.visual") or k.startswith("visual"): + continue # Skip vision encoder + elif k.startswith("model."): + new_sd[k.replace("model.", "", 1)] = v + elif k.startswith("mtp."): + continue # Skip MTP + elif k.startswith("lm_head."): + new_sd[k] = v + else: + new_sd[k] = v + + return convert_qwen35_hf_to_neuron_state_dict(new_sd, config) + + def enable_context_encoding(self): + self.compile_tag = CONTEXT_ENCODING_MODEL_TAG + super().enable_context_encoding() + + def enable_token_generation(self): + self.compile_tag = TOKEN_GENERATION_MODEL_TAG + disable_wlo = bool( + getattr(self.config, "disable_token_generation_wlo", False) + ) or os.environ.get("QWEN36_DISABLE_TOKEN_GENERATION_WLO") == "1" + super().enable_token_generation(enable_wlt_optimization=not disable_wlo) + + def _copy_past_key_values(self, outputs): + """Override to also copy DeltaNet state buffers on CPU.""" + super()._copy_past_key_values(outputs) + if getattr(self.config, "use_hybrid_cache_manager", False): + return + + num_output_from_trace = Qwen35DecoderModelInstance._num_trace_outputs_before_aliases( + self.neuron_config + ) + + if ( + hasattr(self, "token_generation_model") + and self.token_generation_model is not None + ): + tkg_model = self.token_generation_model.model + cte_model = self.context_encoding_model.model + else: + return + + if tkg_model.kv_mgr is not None: + num_kv = len(tkg_model.kv_mgr.past_key_values) + else: + num_kv = 0 + + state_start = num_output_from_trace + num_kv + + tkg_params = getattr(tkg_model, "_deltanet_state_params", []) + cte_params = getattr(cte_model, "_deltanet_state_params", []) + + if len(tkg_params) > 0 and state_start + len(tkg_params) <= len(outputs): + for i, (tkg_param, cte_param) in enumerate(zip(tkg_params, cte_params)): + new_state = outputs[state_start + i] + tkg_param.data = new_state + cte_param.data = new_state + + checkpoint_start = state_start + len(tkg_params) + tkg_checkpoint_params = getattr(tkg_model, "_hybrid_gdn_checkpoint_params", []) + cte_checkpoint_params = getattr(cte_model, "_hybrid_gdn_checkpoint_params", []) + if ( + len(tkg_checkpoint_params) > 0 + and checkpoint_start + len(tkg_checkpoint_params) <= len(outputs) + ): + for i, (tkg_param, cte_param) in enumerate( + zip(tkg_checkpoint_params, cte_checkpoint_params) + ): + new_state = outputs[checkpoint_start + i] + tkg_param.data = new_state + cte_param.data = new_state + + def get_required_kwargs(self): + """Return extra kwargs for HF generation loop.""" + return ["llava_args"] + + def _get_model_outputs( + self, + input_ids, + attention_mask, + position_ids, + seq_ids, + sampling_params, + prev_hidden, + adapter_ids, + medusa_args, + llava_args, + slot_mapping=None, + block_table=None, + full_context_lens=None, + computed_context_lens=None, + tf_args=None, + hybrid_restore_slot_ids=None, + hybrid_restore_mask=None, + hybrid_restore_prefix_lens=None, + hybrid_commit_slot_ids=None, + hybrid_commit_mask=None, + ): + """Override to pass Qwen/vLLM positional args explicitly.""" + prefill_completion_state = getattr( + self, + "_qwen36_vllm_prefill_completion_state", + None, + ) + is_prefill = _qwen36_is_prefill_request( + input_ids, + position_ids, + full_context_lens=full_context_lens, + computed_context_lens=computed_context_lens, + prefill_completion_state=prefill_completion_state, + ) + metadata_by_request_id = getattr( + self, + "_qwen36_vllm_hybrid_apc_metadata_by_request_id", + None, + ) + request_records = getattr( + self, + "_qwen36_vllm_hybrid_apc_request_records", + None, + ) + request_ids = _qwen36_request_ids_from_hybrid_apc_records(request_records) + if request_ids is None: + request_ids = _qwen36_select_vllm_hybrid_apc_request_ids_for_input( + metadata_by_request_id, + all_request_ids=getattr(self, "_qwen36_vllm_request_ids", None), + new_request_ids=getattr(self, "_qwen36_vllm_new_request_ids", None), + full_context_lens=full_context_lens, + computed_context_lens=computed_context_lens, + prefill_completion_state=prefill_completion_state, + ) + if not is_prefill: + ( + input_ids, + attention_mask, + position_ids, + seq_ids, + adapter_ids, + slot_mapping, + ) = _qwen36_unpack_packed_decode_batch( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + seq_ids=seq_ids, + adapter_ids=adapter_ids, + slot_mapping=slot_mapping, + full_context_lens=full_context_lens, + computed_context_lens=computed_context_lens, + ) + seq_ids = _qwen36_stable_seq_ids_for_request_ids( + self, + seq_ids, + request_ids, + ) + + hybrid_apc_request_dict = None + if ( + is_prefill + and _qwen36_config_flag( + self.config, + self.neuron_config, + "use_hybrid_apc_manager", + ) + and getattr(self.neuron_config, "is_prefix_caching", False) + and _qwen36_hybrid_apc_controls_need_prepare( + hybrid_restore_mask, + hybrid_commit_mask, + ) + ): + hybrid_apc_request_dict = { + "input_ids": input_ids, + "attention_mask": attention_mask, + "position_ids": position_ids, + "seq_ids": seq_ids, + "sampling_params": sampling_params, + "adapter_ids": adapter_ids, + "slot_mapping": slot_mapping, + "block_table": block_table, + "full_context_lens": full_context_lens, + "computed_context_lens": computed_context_lens, + } + if llava_args: + hybrid_apc_request_dict["llava_args"] = llava_args + if len(llava_args) >= 3: + hybrid_apc_request_dict["rotary_position_ids"] = llava_args[2] + if request_records is not None: + hybrid_apc_request_dict["hybrid_request_records"] = request_records + if request_ids is not None: + if isinstance(request_ids, list): + request_ids = tuple(request_ids) + if isinstance(request_ids, tuple) and len(request_ids) == 1: + hybrid_apc_request_dict["hybrid_request_id"] = request_ids[0] + else: + hybrid_apc_request_dict["hybrid_request_id"] = request_ids + cached_request_ids = getattr( + self, + "_qwen36_vllm_cached_request_ids", + None, + ) + if cached_request_ids is not None: + hybrid_apc_request_dict["hybrid_cached_request_ids"] = ( + cached_request_ids + ) + if prefill_completion_state is not None: + hybrid_apc_request_dict[ + "hybrid_prefill_completion_state" + ] = prefill_completion_state + _qwen36_add_vllm_hybrid_apc_metadata( + hybrid_apc_request_dict, + request_ids=request_ids, + metadata_by_request_id=metadata_by_request_id, + ) + prepared_inputs = prepare_hybrid_apc_request_for_execution( + self, + hybrid_apc_request_dict, + ) + input_ids = prepared_inputs.get("input_ids", input_ids) + attention_mask = prepared_inputs.get("attention_mask", attention_mask) + position_ids = prepared_inputs.get("position_ids", position_ids) + seq_ids = prepared_inputs.get("seq_ids", seq_ids) + sampling_params = prepared_inputs.get("sampling_params", sampling_params) + adapter_ids = prepared_inputs.get("adapter_ids", adapter_ids) + slot_mapping = prepared_inputs.get("slot_mapping", slot_mapping) + block_table = prepared_inputs.get("block_table", block_table) + full_context_lens = prepared_inputs.get("full_context_lens", full_context_lens) + computed_context_lens = prepared_inputs.get( + "computed_context_lens", + computed_context_lens, + ) + num_queries = prepared_inputs.get("num_queries", num_queries) + hybrid_restore_slot_ids = prepared_inputs.get("hybrid_restore_slot_ids") + hybrid_restore_mask = prepared_inputs.get("hybrid_restore_mask") + hybrid_restore_prefix_lens = prepared_inputs.get( + "hybrid_restore_prefix_lens" + ) + hybrid_commit_slot_ids = prepared_inputs.get("hybrid_commit_slot_ids") + hybrid_commit_mask = prepared_inputs.get("hybrid_commit_mask") + prepared_mrope_position_ids = prepared_inputs.get( + "rotary_position_ids", + prepared_inputs.get("rotary_position_id"), + ) + if prepared_mrope_position_ids is not None and llava_args: + llava_args = list(llava_args) + if len(llava_args) >= 3: + llava_args[2] = prepared_mrope_position_ids + elif len(llava_args) >= 2: + llava_args.append(prepared_mrope_position_ids) + elif prepared_mrope_position_ids is not None: + mrope_position_ids = prepared_mrope_position_ids + else: + prepared_mrope_position_ids = None + + seq_len = input_ids.shape[1] + batch_size = input_ids.shape[0] + + if llava_args and len(llava_args) >= 2: + vision_embeddings = llava_args[0] + vision_mask = llava_args[1] + if len(llava_args) >= 3: + mrope_position_ids = llava_args[2] + else: + mrope_position_ids = None + elif is_prefill: + if getattr(self.config, "use_text_only_cte_inputs", True): + vision_embeddings = torch.zeros( + (0,), dtype=self.config.neuron_config.torch_dtype + ) + vision_mask = torch.zeros((0,), dtype=torch.int32) + else: + vision_embeddings = torch.zeros( + (batch_size, seq_len, self.config.hidden_size), + dtype=self.config.neuron_config.torch_dtype, + ) + vision_mask = torch.full( + (batch_size, seq_len, 1), + fill_value=seq_len - 1, + dtype=torch.int32, + ) + mrope_position_ids = prepared_mrope_position_ids + else: + vision_embeddings = torch.zeros((0,), dtype=torch.float32) + vision_mask = torch.zeros((0,), dtype=torch.int32) + mrope_position_ids = None + + if is_prefill: + if mrope_position_ids is None: + mrope_position_ids = ( + torch.arange(0, seq_len, dtype=torch.int32) + .unsqueeze(0) + .unsqueeze(0) + .expand(3, batch_size, -1) + .contiguous() + ) + else: + mrope_position_ids = torch.zeros((0,), dtype=torch.int32) + + def _empty(): + return torch.empty(0) + + def _optional_tensor(value): + return value if value is not None else _empty() + + def _length_matrix(value, default_value, batch=batch_size): + if value is None or not hasattr(value, "numel") or value.numel() == 0: + return torch.full((batch, 1), default_value, dtype=torch.int32) + value = value.to(torch.int32) + if value.ndim == 0: + return value.reshape(1, 1) + if value.ndim == 1: + return value.reshape(-1, 1) + return value + + def _slice_batch(value, start, end): + if value is None or not hasattr(value, "numel") or value.numel() == 0: + return _empty() + if value.ndim > 0 and value.shape[0] >= end: + return value[start:end] + return value + + def _pad_batch(value, target_batch, fill_value=0): + if value is None or not hasattr(value, "numel") or value.numel() == 0: + return value + if value.ndim == 0 or value.shape[0] >= target_batch: + return value + pad_shape = (target_batch - value.shape[0],) + tuple(value.shape[1:]) + pad = torch.full(pad_shape, fill_value, dtype=value.dtype) + return torch.cat([value, pad], dim=0) + + def _pad_batch_repeat_first(value, target_batch): + if value is None or not hasattr(value, "numel") or value.numel() == 0: + return value + if value.ndim == 0 or value.shape[0] >= target_batch: + return value + pad_n = target_batch - value.shape[0] + return torch.cat([value, value[:1].expand(pad_n, *value.shape[1:])], dim=0) + + if self.neuron_config.is_prefix_caching: + if is_prefill: + computed_context_lens_arg = _length_matrix(computed_context_lens, 0) + full_context_lens_arg = _length_matrix(full_context_lens, seq_len) + num_queries_arg = ( + full_context_lens_arg - computed_context_lens_arg + ).to(torch.int32) + else: + if seq_len != 1: + raise ValueError( + "Qwen3.6 TKG expects active decode length 1, " + f"got input_ids.shape[-1]={seq_len}" + ) + num_queries_arg = torch.full( + (batch_size, 1), seq_len, dtype=torch.int32 + ) + if ( + position_ids is not None + and hasattr(position_ids, "numel") + and position_ids.numel() > 0 + ): + computed_context_lens_arg = _length_matrix(position_ids, 0) + elif full_context_lens is not None: + computed_context_lens_arg = _length_matrix( + full_context_lens, seq_len + ) + else: + computed_context_lens_arg = _length_matrix( + computed_context_lens, + attention_mask.shape[-1] if attention_mask is not None else 0, + ) + slot_mapping_arg = _optional_tensor(slot_mapping) + slot_mapping_arg = _normalize_qwen36_slot_mapping( + slot_mapping_arg, + batch_size, + seq_len, + ) + block_table_arg = _optional_tensor(block_table) + else: + computed_context_lens_arg = _empty() + num_queries_arg = _empty() + slot_mapping_arg = _empty() + block_table_arg = _empty() + + if hybrid_restore_slot_ids is None: + hybrid_restore_slot_ids = torch.zeros((batch_size,), dtype=torch.int32) + if hybrid_restore_mask is None: + hybrid_restore_mask = torch.zeros((batch_size,), dtype=torch.int32) + if hybrid_restore_prefix_lens is None: + hybrid_restore_prefix_lens = torch.zeros((batch_size,), dtype=torch.int32) + if hybrid_commit_slot_ids is None: + hybrid_commit_slot_ids = torch.zeros((batch_size,), dtype=torch.int32) + if hybrid_commit_mask is None: + hybrid_commit_mask = torch.zeros((batch_size,), dtype=torch.int32) + + if is_prefill: + ctx_bs = self.context_encoding_model.neuron_config.batch_size + output_logits = [] + + for cb in range(0, batch_size, ctx_bs): + cb_end = min(cb + ctx_bs, batch_size) + actual_chunk = cb_end - cb + + chunk_input_ids = input_ids[cb:cb_end] + chunk_attn_mask = attention_mask[cb:cb_end] + chunk_pos_ids = position_ids[cb:cb_end] + chunk_seq_ids = seq_ids[cb:cb_end] + chunk_sampling = sampling_params[cb:cb_end] + chunk_slot_mapping = _slice_batch(slot_mapping_arg, cb, cb_end) + chunk_block_table = _slice_batch(block_table_arg, cb, cb_end) + chunk_num_queries = _slice_batch(num_queries_arg, cb, cb_end) + chunk_computed_context_lens = _slice_batch( + computed_context_lens_arg, cb, cb_end + ) + chunk_restore_slots = hybrid_restore_slot_ids[cb:cb_end] + chunk_restore_mask = hybrid_restore_mask[cb:cb_end] + chunk_restore_prefix = hybrid_restore_prefix_lens[cb:cb_end] + chunk_commit_slots = hybrid_commit_slot_ids[cb:cb_end] + chunk_commit_mask = hybrid_commit_mask[cb:cb_end] + chunk_prev_hidden = ( + prev_hidden[cb:cb_end] + if prev_hidden is not None + and hasattr(prev_hidden, "ndim") + and prev_hidden.ndim > 0 + and prev_hidden.shape[0] > 0 + else prev_hidden + ) + chunk_adapter_ids = ( + adapter_ids[cb:cb_end] + if adapter_ids is not None + and hasattr(adapter_ids, "ndim") + and adapter_ids.ndim > 0 + and adapter_ids.shape[0] > 0 + else adapter_ids + ) + + if mrope_position_ids.ndim == 3: + chunk_mrope = mrope_position_ids[:, cb:cb_end, :] + else: + chunk_mrope = mrope_position_ids + + if vision_embeddings.ndim == 3: + chunk_vis_emb = vision_embeddings[cb:cb_end] + chunk_vis_mask = vision_mask[cb:cb_end] + else: + chunk_vis_emb = vision_embeddings + chunk_vis_mask = vision_mask + + if actual_chunk < ctx_bs: + pad_n = ctx_bs - actual_chunk + chunk_input_ids = torch.cat( + [chunk_input_ids, chunk_input_ids[:1].expand(pad_n, -1)], dim=0 + ) + chunk_attn_mask = torch.cat( + [chunk_attn_mask, chunk_attn_mask[:1].expand(pad_n, -1)], dim=0 + ) + chunk_pos_ids = torch.cat( + [chunk_pos_ids, chunk_pos_ids[:1].expand(pad_n, -1)], dim=0 + ) + pad_seq = torch.full( + (pad_n,), -1, dtype=chunk_seq_ids.dtype + ) + chunk_seq_ids = torch.cat([chunk_seq_ids, pad_seq], dim=0) + chunk_sampling = torch.cat( + [chunk_sampling, chunk_sampling[:1].expand(pad_n, -1)], dim=0 + ) + chunk_slot_mapping = _pad_batch(chunk_slot_mapping, ctx_bs, -1) + chunk_block_table = _pad_batch_repeat_first( + chunk_block_table, ctx_bs + ) + chunk_num_queries = _pad_batch_repeat_first( + chunk_num_queries, ctx_bs + ) + chunk_computed_context_lens = _pad_batch( + chunk_computed_context_lens, ctx_bs, 0 + ) + # Dummy CTE rows repeat active token tensors to satisfy the + # compiled batch shape, but they must not advertise a + # prefix-cache restore. Their seq_ids are marked negative + # and the DeltaNet state update preserves negative rows, so + # recurrent state cannot leak into seq_ids later reused by + # real requests. + ( + chunk_restore_slots, + chunk_restore_mask, + chunk_restore_prefix, + ) = _qwen36_pad_hybrid_restore_controls_for_dummy_cte_rows( + chunk_restore_slots, + chunk_restore_mask, + chunk_restore_prefix, + ctx_bs, + ) + chunk_commit_slots = torch.cat( + [chunk_commit_slots, torch.zeros(pad_n, dtype=chunk_commit_slots.dtype)], + dim=0, + ) + chunk_commit_mask = torch.cat( + [chunk_commit_mask, torch.zeros(pad_n, dtype=chunk_commit_mask.dtype)], + dim=0, + ) + if ( + chunk_prev_hidden is not None + and hasattr(chunk_prev_hidden, "ndim") + and chunk_prev_hidden.ndim > 0 + and chunk_prev_hidden.shape[0] > 0 + ): + chunk_prev_hidden = torch.cat( + [ + chunk_prev_hidden, + chunk_prev_hidden[:1].expand(pad_n, -1), + ], + dim=0, + ) + if ( + chunk_adapter_ids is not None + and hasattr(chunk_adapter_ids, "ndim") + and chunk_adapter_ids.ndim > 0 + and chunk_adapter_ids.shape[0] > 0 + ): + chunk_adapter_ids = torch.cat( + [ + chunk_adapter_ids, + chunk_adapter_ids[:1].expand(pad_n, -1), + ], + dim=0, + ) + if chunk_mrope.ndim == 3: + chunk_mrope = torch.cat( + [chunk_mrope, chunk_mrope[:, :1, :].expand(-1, pad_n, -1)], + dim=1, + ) + if chunk_vis_emb.ndim == 3: + chunk_vis_emb = torch.cat( + [ + chunk_vis_emb, + torch.zeros( + (pad_n,) + chunk_vis_emb.shape[1:], + dtype=chunk_vis_emb.dtype, + ), + ], + dim=0, + ) + chunk_vis_mask = torch.cat( + [ + chunk_vis_mask, + torch.full( + (pad_n,) + chunk_vis_mask.shape[1:], + fill_value=seq_len - 1, + dtype=chunk_vis_mask.dtype, + ), + ], + dim=0, + ) + + if os.environ.get("QWEN36_HYBRID_APC_DEBUG") == "1": + def _dbg_minmax(tensor): + if not hasattr(tensor, "numel") or tensor.numel() == 0: + return "empty" + flat = tensor.reshape(-1) + return f"{int(flat.min().item())}:{int(flat.max().item())}" + + print( + "[hybrid_apc_debug] qwen-cte-call " + f"input_shape={tuple(chunk_input_ids.shape)} " + f"attention_shape={tuple(chunk_attn_mask.shape)} " + f"position_shape={tuple(chunk_pos_ids.shape)} " + f"position_minmax={_dbg_minmax(chunk_pos_ids)} " + f"seq_ids={chunk_seq_ids.reshape(-1).tolist() if hasattr(chunk_seq_ids, 'numel') and chunk_seq_ids.numel() else []} " + f"slot_shape={tuple(chunk_slot_mapping.shape)} " + f"slot_minmax={_dbg_minmax(chunk_slot_mapping)} " + f"block_shape={tuple(chunk_block_table.shape)} " + f"block_minmax={_dbg_minmax(chunk_block_table)} " + f"num_queries={chunk_num_queries.reshape(-1).tolist() if hasattr(chunk_num_queries, 'numel') and chunk_num_queries.numel() else []} " + f"computed={chunk_computed_context_lens.reshape(-1).tolist() if hasattr(chunk_computed_context_lens, 'numel') and chunk_computed_context_lens.numel() else []} " + f"restore_slots={chunk_restore_slots.reshape(-1).tolist()} " + f"restore_mask={chunk_restore_mask.reshape(-1).tolist()} " + f"restore_prefix={chunk_restore_prefix.reshape(-1).tolist()} " + f"commit_slots={chunk_commit_slots.reshape(-1).tolist()} " + f"commit_mask={chunk_commit_mask.reshape(-1).tolist()}", + flush=True, + ) + + cte_prefix_args = [ + chunk_input_ids, + chunk_attn_mask, + chunk_pos_ids, + chunk_seq_ids, + chunk_sampling, + chunk_prev_hidden, + chunk_adapter_ids, + _empty(), + _empty(), + _empty(), + _empty(), + chunk_slot_mapping, + chunk_block_table, + chunk_num_queries, + chunk_computed_context_lens, + _empty(), + _empty(), + _empty(), + _empty(), + _empty(), + _empty(), + ] + cte_args = build_cte_args( + self.config, + cte_prefix_args, + chunk_mrope, + chunk_vis_emb, + chunk_vis_mask, + hybrid_args=( + chunk_restore_slots, + chunk_restore_mask, + chunk_restore_prefix, + chunk_commit_slots, + chunk_commit_mask, + ), + ) + _debug_qwen36_arg_contract( + "runtime", + CONTEXT_ENCODING_MODEL_TAG, + self.config, + cte_args, + ) + _qwen36_prefill_timing = os.environ.get("QWEN36_PREFILL_TIMING") == "1" + _qwen36_cte_start = time.perf_counter() if _qwen36_prefill_timing else None + try: + chunk_out = self.context_encoding_model(*cte_args) + except Exception: + if hybrid_apc_request_dict is not None: + cancel_hybrid_apc_request(hybrid_apc_request_dict) + hybrid_apc_request_dict = None + raise + if _qwen36_prefill_timing and _qwen36_cte_start is not None: + print( + "[qwen36_perf] qwen_cte_call " + f"elapsed_ms={(time.perf_counter() - _qwen36_cte_start) * 1000.0:.3f} " + f"actual_chunk={actual_chunk} ctx_bs={ctx_bs} " + f"input_shape={tuple(chunk_input_ids.shape)} " + f"num_queries={chunk_num_queries.reshape(-1).tolist() if hasattr(chunk_num_queries, 'numel') and chunk_num_queries.numel() else []} " + f"computed={chunk_computed_context_lens.reshape(-1).tolist() if hasattr(chunk_computed_context_lens, 'numel') and chunk_computed_context_lens.numel() else []} " + f"restore_mask={chunk_restore_mask.reshape(-1).tolist()} " + f"commit_mask={chunk_commit_mask.reshape(-1).tolist()}", + flush=True, + ) + if actual_chunk < ctx_bs: + chunk_out = chunk_out[:actual_chunk] + output_logits.append(chunk_out) + + outputs = ( + torch.cat(output_logits, dim=0) + if len(output_logits) > 1 + else output_logits[0] + ) + self.kv_cache_populated = True + is_run_on_neuron = self.context_encoding_model.is_neuron() + if hybrid_apc_request_dict is not None: + finish_hybrid_apc_request(hybrid_apc_request_dict) + else: + _validate_qwen36_tkg_input_ids( + input_ids, + getattr(self.config, "vocab_size", None), + ) + legacy_tkg_args = _use_legacy_tkg_args() + if ( + os.environ.get("QWEN36_TKG_INPUT_DEBUG") == "1" + or os.environ.get("QWEN36_HYBRID_APC_DEBUG") == "1" + ): + max_model_len = getattr( + self.neuron_config, + "max_length", + getattr(self.neuron_config, "seq_len", None), + ) + print( + "[hybrid_apc_debug] qwen-tkg-call " + f"arg_mode={'prefix24_legacy' if legacy_tkg_args else 'hybrid29'} " + f"input_shape={_debug_tensor_shape(input_ids)} " + f"input_values={_debug_tensor_values(input_ids)} " + f"attention_shape={_debug_tensor_shape(attention_mask)} " + f"position_shape={_debug_tensor_shape(position_ids)} " + f"position_minmax={_debug_tensor_minmax(position_ids)} " + f"slot_shape={_debug_tensor_shape(slot_mapping_arg)} " + f"slot_minmax={_debug_tensor_minmax(slot_mapping_arg)} " + f"block_shape={_debug_tensor_shape(block_table_arg)} " + f"block_minmax={_debug_tensor_minmax(block_table_arg)} " + f"num_queries={_debug_tensor_values(num_queries_arg)} " + "computed_context_lens=" + f"{_debug_tensor_values(computed_context_lens_arg)} " + f"pa_num_blocks={getattr(self.neuron_config, 'pa_num_blocks', None)} " + f"block_size={getattr(self.neuron_config, 'pa_block_size', None)} " + f"seq_len={seq_len} max_model_len={max_model_len}", + flush=True, + ) + tkg_prefix_args = [ + input_ids, + attention_mask, + position_ids, + seq_ids, + sampling_params, + prev_hidden, + adapter_ids, + _empty(), + _empty(), + _empty(), + _empty(), + slot_mapping_arg, + block_table_arg, + num_queries_arg, + computed_context_lens_arg, + _empty(), + _empty(), + _empty(), + _empty(), + _empty(), + _empty(), + ] + tkg_args = build_tkg_args( + self.config, + tkg_prefix_args, + mrope_position_ids, + vision_embeddings, + vision_mask, + hybrid_args=( + hybrid_restore_slot_ids, + hybrid_restore_mask, + hybrid_restore_prefix_lens, + hybrid_commit_slot_ids, + hybrid_commit_mask, + ), + ) + _debug_qwen36_arg_contract( + "runtime", + TOKEN_GENERATION_MODEL_TAG, + self.config, + tkg_args, + ) + outputs = self.token_generation_model(*tkg_args) + is_run_on_neuron = self.token_generation_model.is_neuron() + + return outputs, is_run_on_neuron + + def get_compiler_args(self): + if self.compile_tag == CONTEXT_ENCODING_MODEL_TAG: + optimization_level = "-O1" + else: + optimization_level = "-O1" + + compiler_args = ( + "--enable-saturate-infinity " + "--enable-mixed-precision-accumulation " + f"--model-type transformer {optimization_level} " + "--auto-cast=none " + ) + return compiler_args diff --git a/contrib/models/Qwen3.5-35B-A3B/src/modeling_qwen35_vision.py b/contrib/models/Qwen3.5-35B-A3B/src/modeling_qwen35_vision.py new file mode 100644 index 00000000..0a15d2c1 --- /dev/null +++ b/contrib/models/Qwen3.5-35B-A3B/src/modeling_qwen35_vision.py @@ -0,0 +1,929 @@ +""" +Qwen3.5-27B / Qwen3.6-27B (Dense) Vision Encoder for NeuronX Distributed Inference. + +Ports the Qwen3.5/3.6 ViT encoder to run on Neuron. The vision encoder +architecture is identical across Qwen3.5-27B and Qwen3.6-27B (same patch +embed, same rotary, same merger) -- only out_hidden_size changes vs the MoE +variant (5120 vs 2048, read from config). + +The vision encoder runs as a separate compiled model from the text decoder, +compiled and loaded via NeuronBaseForImageToText. +""" + +import logging +import math +import os +from typing import Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F + +# CRITICAL: Use finite negative value instead of -inf for Neuron attention masks. +# The Neuron compiler's bfloat16 handling of -inf produces NaN that bleeds from +# padding positions into ALL positions through the transformer layers. +# -65504.0 is large enough for softmax masking but avoids NaN overflow. +_MASK_NEG_INF = -65504.0 + +logger = logging.getLogger(__name__) + +# -- NxDI imports (available on Neuron instances) -- +try: + from neuronx_distributed_inference.models.application_base import ( + NeuronApplicationBase, + ) + from neuronx_distributed_inference.models.model_wrapper import ModelWrapper + from neuronx_distributed_inference.modules.attention.attention_base import ( + NeuronAttentionBase, + ) + from neuronx_distributed_inference.modules.attention.utils import RotaryEmbedding + from neuronx_distributed.parallel_layers import layers as nxd_layers +except ImportError: + logger.warning( + "NxDI imports unavailable -- vision module can only be used on Neuron instances" + ) + +# -- HuggingFace imports for patch embed (runs on CPU) -- +try: + from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import ( + Qwen3_5MoeVisionPatchEmbed, + Qwen3_5MoeVisionPatchMerger, + Qwen3_5MoeVisionRotaryEmbedding, + ) +except ImportError: + try: + # transformers 4.57+ uses Qwen3VL* class names + from transformers.models.qwen3_vl.modeling_qwen3_vl import ( + Qwen3VLVisionPatchEmbed as Qwen3_5MoeVisionPatchEmbed, + Qwen3VLVisionPatchMerger as Qwen3_5MoeVisionPatchMerger, + Qwen3VLVisionRotaryEmbedding as Qwen3_5MoeVisionRotaryEmbedding, + ) + except ImportError: + try: + # Older transformers uses Qwen2VL* class names + from transformers.models.qwen2_vl.modeling_qwen2_vl import ( + Qwen2VLVisionPatchEmbed as Qwen3_5MoeVisionPatchEmbed, + Qwen2VLVisionPatchMerger as Qwen3_5MoeVisionPatchMerger, + Qwen2VLVisionRotaryEmbedding as Qwen3_5MoeVisionRotaryEmbedding, + ) + except ImportError: + Qwen3_5MoeVisionPatchEmbed = None + Qwen3_5MoeVisionPatchMerger = None + Qwen3_5MoeVisionRotaryEmbedding = None + + +def apply_rotary_pos_emb_vision(q, k, cos, sin): + """Apply rotary position embeddings to vision Q and K tensors. + + Uses rotate_half style (matching HF reference): + q_embed = (q * cos) + (rotate_half(q) * sin) + + Args: + q: (seq_len, num_heads, head_dim) + k: (seq_len, num_heads, head_dim) + cos: (seq_len, head_dim) + sin: (seq_len, head_dim) + """ + cos = cos.unsqueeze(-2) # (seq_len, 1, head_dim) + sin = sin.unsqueeze(-2) + + def rotate_half(x): + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed.to(q.dtype), k_embed.to(k.dtype) + + +class NeuronQwen35VisionAttention(nn.Module): + """Vision attention for Qwen3.5 MoE. + + Uses fused QKV linear (no bias in Neuron port for efficiency). + Non-causal attention with block-diagonal mask for variable-length images. + """ + + def __init__(self, config): + super().__init__() + self.hidden_size = config.hidden_size + self.num_heads = config.num_heads + self.head_dim = self.hidden_size // self.num_heads + self.scaling = self.head_dim**-0.5 + + # Fused QKV: (hidden_size -> 3 * hidden_size) with bias + self.qkv = nxd_layers.ColumnParallelLinear( + self.hidden_size, + 3 * self.hidden_size, + bias=True, + gather_output=True, + ) + self.proj = nxd_layers.RowParallelLinear( + self.hidden_size, + self.hidden_size, + bias=True, + input_is_parallel=False, + ) + + def forward(self, hidden_states, attention_mask=None, position_embeddings=None): + """ + Args: + hidden_states: (seq_len, hidden_size) + attention_mask: (1, 1, seq_len, seq_len) block-diagonal mask + position_embeddings: (cos, sin) tuple + """ + seq_len = hidden_states.shape[0] + + # QKV projection + qkv = self.qkv(hidden_states) # (seq_len, 3 * hidden_size) + qkv = qkv.reshape(seq_len, 3, self.num_heads, self.head_dim) + qkv = qkv.permute(1, 0, 2, 3) # (3, seq_len, num_heads, head_dim) + q, k, v = qkv.unbind(0) # each (seq_len, num_heads, head_dim) + + # Apply rotary embeddings + if position_embeddings is not None: + cos, sin = position_embeddings + q, k = apply_rotary_pos_emb_vision(q, k, cos, sin) + + # Reshape for batched attention: (1, num_heads, seq_len, head_dim) + q = q.transpose(0, 1).unsqueeze(0) + k = k.transpose(0, 1).unsqueeze(0) + v = v.transpose(0, 1).unsqueeze(0) + + # Scaled dot-product attention + attn_weights = torch.matmul(q, k.transpose(-1, -2)) * self.scaling + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(q.dtype) + attn_output = torch.matmul(attn_weights, v) + + # Reshape back: (seq_len, hidden_size) + attn_output = attn_output.squeeze(0).transpose(0, 1).reshape(seq_len, -1) + + # Output projection + attn_output = self.proj(attn_output) + return attn_output + + +class NeuronQwen35VisionMLP(nn.Module): + """Vision MLP with GELU activation.""" + + def __init__(self, config): + super().__init__() + self.linear_fc1 = nxd_layers.ColumnParallelLinear( + config.hidden_size, + config.intermediate_size, + bias=True, + gather_output=True, + ) + self.linear_fc2 = nxd_layers.RowParallelLinear( + config.intermediate_size, + config.hidden_size, + bias=True, + input_is_parallel=False, + ) + self.act_fn = nn.GELU() + + def forward(self, hidden_states): + return self.linear_fc2(self.act_fn(self.linear_fc1(hidden_states))) + + +class NeuronQwen35VisionBlock(nn.Module): + """Single vision transformer block: LayerNorm + Attention + LayerNorm + MLP.""" + + def __init__(self, config): + super().__init__() + self.norm1 = nn.LayerNorm(config.hidden_size, eps=1e-6) + self.norm2 = nn.LayerNorm(config.hidden_size, eps=1e-6) + self.attn = NeuronQwen35VisionAttention(config) + self.mlp = NeuronQwen35VisionMLP(config) + + def forward(self, hidden_states, attention_mask=None, position_embeddings=None): + hidden_states = hidden_states + self.attn( + self.norm1(hidden_states), + attention_mask=attention_mask, + position_embeddings=position_embeddings, + ) + hidden_states = hidden_states + self.mlp(self.norm2(hidden_states)) + return hidden_states + + +class NeuronQwen35VisionModel(nn.Module): + """Qwen3.5 MoE Vision Encoder for Neuron. + + This is the nn.Module that gets compiled and traced onto Neuron. + Patch embedding, positional embedding, and rotary embedding are computed + on CPU in the ModelWrapper and passed as inputs. + """ + + def __init__(self, config): + super().__init__() + self.config = config + self.blocks = nn.ModuleList( + [NeuronQwen35VisionBlock(config) for _ in range(config.depth)] + ) + # Merger: spatial_merge_size^2 * hidden_size -> out_hidden_size + self.merger_norm = nn.LayerNorm(config.hidden_size, eps=1e-6) + merger_hidden = config.hidden_size * (config.spatial_merge_size**2) + self.merger_fc1 = nn.Linear(merger_hidden, merger_hidden) + self.merger_act = nn.GELU() + self.merger_fc2 = nn.Linear(merger_hidden, config.out_hidden_size) + + def forward(self, hidden_states, attention_mask=None, position_embeddings=None): + """ + Args: + hidden_states: (seq_len, hidden_size) -- after patch_embed + pos_embed + attention_mask: (1, 1, seq_len, seq_len) block-diagonal mask + position_embeddings: (cos, sin) tuple for rotary + + Returns: + vision_embeddings: (merged_seq_len, out_hidden_size) + """ + for block in self.blocks: + hidden_states = block( + hidden_states, + attention_mask=attention_mask, + position_embeddings=position_embeddings, + ) + + # Apply merger: norm -> spatial merge -> fc1 -> gelu -> fc2 + hidden_states = self.merger_norm(hidden_states) + merge_size = self.config.spatial_merge_size + merged_hidden = self.config.hidden_size * (merge_size**2) + hidden_states = hidden_states.view(-1, merged_hidden) + hidden_states = self.merger_fc2(self.merger_act(self.merger_fc1(hidden_states))) + + return hidden_states + + +class CPUVisionModel(nn.Module): + """CPU-only vision encoder (pure PyTorch, no Neuron dependencies). + + Used when HBM is insufficient to load the vision encoder on Neuron + alongside the text decoder (e.g., 27B dense model on trn2.3xlarge). + """ + + def __init__(self, config): + super().__init__() + self.config = config + self.blocks = nn.ModuleList( + [self._make_block(config) for _ in range(config.depth)] + ) + self.merger_norm = nn.LayerNorm(config.hidden_size, eps=1e-6) + merger_hidden = config.hidden_size * (config.spatial_merge_size**2) + self.merger_fc1 = nn.Linear(merger_hidden, merger_hidden) + self.merger_act = nn.GELU() + self.merger_fc2 = nn.Linear(merger_hidden, config.out_hidden_size) + + @staticmethod + def _make_block(config): + """Build a single vision block with standard nn.Linear (no TP).""" + block = nn.Module() + block.norm1 = nn.LayerNorm(config.hidden_size, eps=1e-6) + block.norm2 = nn.LayerNorm(config.hidden_size, eps=1e-6) + + # Attention + attn = nn.Module() + attn.hidden_size = config.hidden_size + attn.num_heads = config.num_heads + attn.head_dim = config.hidden_size // config.num_heads + attn.scaling = attn.head_dim**-0.5 + attn.qkv = nn.Linear(config.hidden_size, 3 * config.hidden_size, bias=True) + attn.proj = nn.Linear(config.hidden_size, config.hidden_size, bias=True) + block.attn = attn + + # MLP + mlp = nn.Module() + mlp.linear_fc1 = nn.Linear( + config.hidden_size, config.intermediate_size, bias=True + ) + mlp.linear_fc2 = nn.Linear( + config.intermediate_size, config.hidden_size, bias=True + ) + mlp.act_fn = nn.GELU() + block.mlp = mlp + + return block + + def _forward_attention(self, attn, hidden_states, attention_mask, cos, sin): + seq_len = hidden_states.shape[0] + qkv = attn.qkv(hidden_states).reshape(seq_len, 3, attn.num_heads, attn.head_dim) + qkv = qkv.permute(1, 0, 2, 3) + q, k, v = qkv.unbind(0) + + if cos is not None and sin is not None: + cos_u = cos.unsqueeze(-2) + sin_u = sin.unsqueeze(-2) + + def rotate_half(x): + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + q = (q * cos_u) + (rotate_half(q) * sin_u) + k = (k * cos_u) + (rotate_half(k) * sin_u) + + q = q.transpose(0, 1).unsqueeze(0) + k = k.transpose(0, 1).unsqueeze(0) + v = v.transpose(0, 1).unsqueeze(0) + + # Use PyTorch's fused SDPA (efficient CPU implementation with memory- + # efficient attention). Falls back to eager on very old torch. + # attention_mask here is additive (bf16 with -large at masked positions). + try: + out = F.scaled_dot_product_attention( + q, k, v, + attn_mask=attention_mask, + dropout_p=0.0, + scale=attn.scaling, + ) + except Exception: + attn_weights = torch.matmul(q, k.transpose(-1, -2)) * attn.scaling + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(q.dtype) + out = torch.matmul(attn_weights, v) + out = out.squeeze(0).transpose(0, 1).reshape(seq_len, -1) + return attn.proj(out) + + def forward(self, hidden_states, attention_mask, cos, sin): + for block in self.blocks: + hidden_states = hidden_states + self._forward_attention( + block.attn, block.norm1(hidden_states), attention_mask, cos, sin + ) + hidden_states = hidden_states + block.mlp.linear_fc2( + block.mlp.act_fn(block.mlp.linear_fc1(block.norm2(hidden_states))) + ) + + hidden_states = self.merger_norm(hidden_states) + merge_size = self.config.spatial_merge_size + merged_hidden = self.config.hidden_size * (merge_size**2) + hidden_states = hidden_states.view(-1, merged_hidden) + hidden_states = self.merger_fc2(self.merger_act(self.merger_fc1(hidden_states))) + return hidden_states + + +class NeuronQwen35VisionModelWrapper(ModelWrapper): + """Wraps the vision encoder for NxDI tracing. + + Handles CPU-side operations that cannot be traced: + - Patch embedding (Conv3d) + - Positional embedding (Embedding + bilinear interpolation) + - Rotary position embedding computation + - Vision attention mask construction (block-diagonal) + - Sequence length bucketing and padding/unpadding + + Supports three modes: + 1. NxDI traced model (parallel layers) -- standard NxDI compilation + 2. Pre-compiled standalone model -- loaded from torch_neuronx.trace() output + 3. CPU-only model -- for when HBM is full (e.g., 27B dense on trn2.3xlarge) + """ + + def __init__(self, config, model_cls=None, **kwargs): + if model_cls is not None: + super().__init__(config, model_cls, **kwargs) + else: + # Standalone mode: no NxDI model_cls + nn.Module.__init__(self) + self.vision_config = config + self._compiled_model = None # Set by load_compiled() -- single bucket + self._compiled_buckets = None # Set by load_compiled() -- multi-bucket dict + self._cpu_model = None # Set by load_cpu_model() + + # These HF modules run on CPU, outside the traced graph + if Qwen3_5MoeVisionPatchEmbed is not None: + self.patch_embed = Qwen3_5MoeVisionPatchEmbed(config) + self.pos_embed = nn.Embedding( + config.num_position_embeddings, config.hidden_size + ) + self.num_grid_per_side = int(config.num_position_embeddings**0.5) + head_dim = config.hidden_size // config.num_heads + self.rotary_pos_emb = Qwen3_5MoeVisionRotaryEmbedding(head_dim // 2) + else: + logger.warning("HF Qwen3.5 MoE vision classes not available") + + self.vision_seq_len_buckets = kwargs.get( + "vision_seq_len_buckets", [1024, 4096, 16384] + ) + + def load_compiled(self, compiled_model_path): + """Load pre-compiled standalone vision encoder(s). + + Supports two modes: + 1. Single .pt file: Legacy mode, loads one compiled model for one bucket size. + 2. Directory with multiple .pt files: Multi-bucket mode. Files must be named + 'vision_encoder_{bucket_size}.pt' (e.g., 'vision_encoder_256.pt'). + Falls back to single 'vision_encoder.pt' in the directory. + + Args: + compiled_model_path: Path to a .pt file or directory containing bucket .pt files. + """ + import glob as glob_module + + logger.info(f"Loading pre-compiled vision encoder from {compiled_model_path}") + + if os.path.isfile(compiled_model_path): + # Single file mode (legacy) + self._compiled_model = torch.jit.load(compiled_model_path) + self._compiled_buckets = None + logger.info("Vision encoder loaded successfully (single bucket)") + elif os.path.isdir(compiled_model_path): + # Directory mode: look for bucket-specific files + bucket_files = sorted( + glob_module.glob( + os.path.join(compiled_model_path, "vision_encoder_*.pt") + ) + ) + if bucket_files: + self._compiled_buckets = {} + for bf in bucket_files: + # Extract bucket size from filename: vision_encoder_256.pt -> 256 + basename = os.path.basename(bf) + try: + bucket_size = int( + basename.replace("vision_encoder_", "").replace(".pt", "") + ) + self._compiled_buckets[bucket_size] = torch.jit.load(bf) + logger.info(f" Loaded vision bucket {bucket_size} from {bf}") + except ValueError: + logger.warning(f" Skipping unrecognized file: {bf}") + self._compiled_model = None + # Update vision_seq_len_buckets to match compiled buckets + self.vision_seq_len_buckets = sorted(self._compiled_buckets.keys()) + logger.info( + f"Vision encoder loaded with {len(self._compiled_buckets)} buckets: " + f"{self.vision_seq_len_buckets}" + ) + else: + # Fall back to single vision_encoder.pt in directory + single_path = os.path.join(compiled_model_path, "vision_encoder.pt") + if os.path.exists(single_path): + self._compiled_model = torch.jit.load(single_path) + self._compiled_buckets = None + logger.info( + "Vision encoder loaded successfully (single file in dir)" + ) + else: + raise FileNotFoundError( + f"No vision encoder files found in {compiled_model_path}" + ) + else: + raise FileNotFoundError( + f"Vision encoder path not found: {compiled_model_path}" + ) + + def load_vision_weights_from_hf(self, model_path): + """Load patch_embed and pos_embed weights from HF safetensors. + + Args: + model_path: Path to HF model directory + """ + from pathlib import Path + from safetensors import safe_open + + st_files = sorted( + p + for p in Path(model_path).glob("*.safetensors") + if p.suffix == ".safetensors" + ) + loaded = 0 + for sf_path in st_files: + with safe_open(str(sf_path), framework="pt") as f: + for key in f.keys(): + if key == "model.visual.patch_embed.proj.weight": + self.patch_embed.proj.weight.data.copy_(f.get_tensor(key)) + loaded += 1 + elif key == "model.visual.patch_embed.proj.bias": + self.patch_embed.proj.bias.data.copy_(f.get_tensor(key)) + loaded += 1 + elif key == "model.visual.pos_embed.weight": + self.pos_embed.weight.data.copy_(f.get_tensor(key)) + loaded += 1 + logger.info(f"Loaded {loaded} CPU-side vision weight tensors from HF") + + def load_cpu_model(self, model_path): + """Load a CPU-only vision encoder from HF safetensors. + + Use this when HBM is insufficient for the Neuron-compiled vision encoder + (e.g., 27B dense model fills trn2.3xlarge HBM). + + Args: + model_path: Path to HF model directory with safetensors + """ + from pathlib import Path + from safetensors import safe_open + + config = self.vision_config + cpu_model = CPUVisionModel(config) + + # Build key mapping from HF safetensors to CPU model + key_map = {} + for i in range(config.depth): + hf_pre = f"model.visual.blocks.{i}" + loc_pre = f"blocks.{i}" + for suffix in [ + "attn.qkv.weight", + "attn.qkv.bias", + "attn.proj.weight", + "attn.proj.bias", + "mlp.linear_fc1.weight", + "mlp.linear_fc1.bias", + "mlp.linear_fc2.weight", + "mlp.linear_fc2.bias", + "norm1.weight", + "norm1.bias", + "norm2.weight", + "norm2.bias", + ]: + key_map[f"{hf_pre}.{suffix}"] = f"{loc_pre}.{suffix}" + + key_map["model.visual.merger.norm.weight"] = "merger_norm.weight" + key_map["model.visual.merger.norm.bias"] = "merger_norm.bias" + key_map["model.visual.merger.linear_fc1.weight"] = "merger_fc1.weight" + key_map["model.visual.merger.linear_fc1.bias"] = "merger_fc1.bias" + key_map["model.visual.merger.linear_fc2.weight"] = "merger_fc2.weight" + key_map["model.visual.merger.linear_fc2.bias"] = "merger_fc2.bias" + + st_files = sorted(Path(model_path).glob("model*.safetensors")) + loaded = 0 + state_dict = cpu_model.state_dict() + + for sf_path in st_files: + with safe_open(str(sf_path), framework="pt") as f: + for key in f.keys(): + if key in key_map: + local_key = key_map[key] + if local_key in state_dict: + state_dict[local_key].copy_(f.get_tensor(key)) + loaded += 1 + + cpu_model.load_state_dict(state_dict) + cpu_model = cpu_model.to(torch.bfloat16).eval() + self._cpu_model = cpu_model + logger.info( + f"Loaded CPU vision encoder: {loaded} weights, " + f"{sum(p.numel() for p in cpu_model.parameters()) / 1e6:.1f}M params" + ) + + def _get_vision_bucket(self, seq_len): + """Find the smallest bucket that fits the sequence length.""" + for bucket in sorted(self.vision_seq_len_buckets): + if seq_len <= bucket: + return bucket + return self.vision_seq_len_buckets[-1] + + def rot_pos_emb(self, grid_thw): + """Compute rotary positional embeddings for vision tokens. + + Returns: (total_tokens, head_dim) tensor of rotary frequencies. + """ + merge_size = self.vision_config.spatial_merge_size + grid_thw_list = grid_thw.tolist() + + max_hw = max(max(h, w) for _, h, w in grid_thw_list) + freq_table = self.rotary_pos_emb(max_hw) + device = freq_table.device + + total_tokens = sum(t * h * w for t, h, w in grid_thw_list) + pos_ids = torch.empty((total_tokens, 2), dtype=torch.long, device=device) + + offset = 0 + for num_frames, height, width in grid_thw_list: + merged_h, merged_w = height // merge_size, width // merge_size + + block_rows = torch.arange(merged_h, device=device) + block_cols = torch.arange(merged_w, device=device) + intra_row = torch.arange(merge_size, device=device) + intra_col = torch.arange(merge_size, device=device) + + row_idx = ( + block_rows[:, None, None, None] * merge_size + + intra_row[None, None, :, None] + ) + col_idx = ( + block_cols[None, :, None, None] * merge_size + + intra_col[None, None, None, :] + ) + + row_idx = row_idx.expand( + merged_h, merged_w, merge_size, merge_size + ).reshape(-1) + col_idx = col_idx.expand( + merged_h, merged_w, merge_size, merge_size + ).reshape(-1) + + coords = torch.stack((row_idx, col_idx), dim=-1) + if num_frames > 1: + coords = coords.repeat(num_frames, 1) + + num_tokens = coords.shape[0] + pos_ids[offset : offset + num_tokens] = coords + offset += num_tokens + + embeddings = freq_table[pos_ids] + embeddings = embeddings.flatten(1) + return embeddings + + def fast_pos_embed_interpolate(self, grid_thw): + """Bilinear interpolation of positional embeddings for variable resolution.""" + grid_thw_list = grid_thw.tolist() + grid_ts = [row[0] for row in grid_thw_list] + grid_hs = [row[1] for row in grid_thw_list] + grid_ws = [row[2] for row in grid_thw_list] + device = self.pos_embed.weight.device + + idx_list = [[] for _ in range(4)] + weight_list = [[] for _ in range(4)] + + for t, h, w in grid_thw_list: + h_idxs = torch.linspace(0, self.num_grid_per_side - 1, h) + w_idxs = torch.linspace(0, self.num_grid_per_side - 1, w) + + h_idxs_floor = h_idxs.int() + w_idxs_floor = w_idxs.int() + h_idxs_ceil = (h_idxs.int() + 1).clip(max=self.num_grid_per_side - 1) + w_idxs_ceil = (w_idxs.int() + 1).clip(max=self.num_grid_per_side - 1) + + dh = h_idxs - h_idxs_floor + dw = w_idxs - w_idxs_floor + + base_h = h_idxs_floor * self.num_grid_per_side + base_h_ceil = h_idxs_ceil * self.num_grid_per_side + + indices = [ + (base_h[None].T + w_idxs_floor[None]).flatten(), + (base_h[None].T + w_idxs_ceil[None]).flatten(), + (base_h_ceil[None].T + w_idxs_floor[None]).flatten(), + (base_h_ceil[None].T + w_idxs_ceil[None]).flatten(), + ] + weights = [ + ((1 - dh)[None].T * (1 - dw)[None]).flatten(), + ((1 - dh)[None].T * dw[None]).flatten(), + (dh[None].T * (1 - dw)[None]).flatten(), + (dh[None].T * dw[None]).flatten(), + ] + + for i in range(4): + idx_list[i].extend(indices[i].tolist()) + weight_list[i].extend(weights[i].tolist()) + + idx_tensor = torch.tensor(idx_list, dtype=torch.long, device=device) + weight_tensor = torch.tensor( + weight_list, dtype=self.pos_embed.weight.dtype, device=device + ) + pos_embeds = self.pos_embed(idx_tensor).to(device) * weight_tensor[:, :, None] + patch_pos_embeds = pos_embeds[0] + pos_embeds[1] + pos_embeds[2] + pos_embeds[3] + + patch_pos_embeds = patch_pos_embeds.split( + [h * w for h, w in zip(grid_hs, grid_ws)] + ) + + merge_size = self.vision_config.spatial_merge_size + patch_pos_embeds_permute = [] + for pos_embed, t, h, w in zip(patch_pos_embeds, grid_ts, grid_hs, grid_ws): + pos_embed = pos_embed.repeat(t, 1) + pos_embed = ( + pos_embed.view( + t, h // merge_size, merge_size, w // merge_size, merge_size, -1 + ) + .permute(0, 1, 3, 2, 4, 5) + .flatten(0, 4) + ) + patch_pos_embeds_permute.append(pos_embed) + + return torch.cat(patch_pos_embeds_permute) + + def _build_vision_attention_mask(self, grid_thw, seq_len, dtype): + """Build block-diagonal attention mask for variable-length images. + + Each image gets its own attention block (no cross-image attention). + """ + cu_seqlens = torch.repeat_interleave( + grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0] + ).cumsum(dim=0, dtype=torch.int32) + cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0) + + # Build block-diagonal mask + mask = torch.full((seq_len, seq_len), _MASK_NEG_INF, dtype=dtype) + for i in range(len(cu_seqlens) - 1): + start = cu_seqlens[i].item() + end = cu_seqlens[i + 1].item() + mask[start:end, start:end] = 0.0 + + return mask.unsqueeze(0).unsqueeze(0) # (1, 1, seq_len, seq_len) + + def forward(self, pixel_values, image_grid_thw): + """Run vision encoding (CPU preprocessing + Neuron traced model). + + Args: + pixel_values: Raw pixel values from HF processor + image_grid_thw: (num_images, 3) -- temporal, height, width in patches + + Returns: + vision_embeddings: (total_merged_tokens, out_hidden_size) + """ + # 0. Tile path: if input exceeds all compiled Neuron buckets and tiling + # is enabled, split pixel_values into 2×2 spatial tiles of the + # (H, W) patch grid, encode each tile independently (each tile has + # 1/4 the tokens → fits in a smaller bucket), and re-interleave the + # merged outputs. Trade-off: no cross-tile attention. + seq_len_probe = pixel_values.shape[0] + max_neuron_bucket = ( + max(self._compiled_buckets.keys()) + if self._compiled_buckets is not None + else 0 + ) + gthw_list = image_grid_thw.tolist() + can_tile = ( + self._compiled_buckets is not None + and seq_len_probe > max_neuron_bucket + and len(gthw_list) == 1 + and gthw_list[0][1] % 2 == 0 + and gthw_list[0][2] % 2 == 0 + and (seq_len_probe // 4) <= max_neuron_bucket + ) + if can_tile: + logger.info( + f"seq_len={seq_len_probe} exceeds max compiled bucket " + f"{max_neuron_bucket}; using 2×2 tiled vision path." + ) + return self._tiled_forward(pixel_values, image_grid_thw) + + # 1. Patch embedding (CPU, Conv3d) + hidden_states = self.patch_embed(pixel_values) + + # 2. Positional embedding (CPU, bilinear interpolation) + pos_embeds = self.fast_pos_embed_interpolate(image_grid_thw) + hidden_states = hidden_states + pos_embeds + + # 3. Rotary position embeddings (CPU) + rotary_pos_emb = self.rot_pos_emb(image_grid_thw) + emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1) + position_embeddings = (emb.cos(), emb.sin()) + + # 4. Vision attention mask (block-diagonal) + seq_len = hidden_states.shape[0] + attention_mask = self._build_vision_attention_mask( + image_grid_thw, seq_len, hidden_states.dtype + ) + + # 5. Bucket and pad for Neuron compilation + bucket_len = self._get_vision_bucket(seq_len) + cos, sin = position_embeddings + if seq_len < bucket_len: + pad_len = bucket_len - seq_len + hidden_states = F.pad(hidden_states, (0, 0, 0, pad_len)) + cos = F.pad(cos, (0, 0, 0, pad_len)) + sin = F.pad(sin, (0, 0, 0, pad_len)) + # Extend mask with _MASK_NEG_INF for padded positions (NOT -inf, which causes NaN on Neuron) + mask = torch.full( + (1, 1, bucket_len, bucket_len), _MASK_NEG_INF, dtype=hidden_states.dtype + ) + mask[:, :, :seq_len, :seq_len] = attention_mask + attention_mask = mask + + # 6. Run vision model (Neuron compiled or CPU fallback) + use_neuron_bucket = ( + self._compiled_buckets is not None + and bucket_len in self._compiled_buckets + and seq_len <= bucket_len + ) + if use_neuron_bucket: + compiled_model = self._compiled_buckets[bucket_len] + vision_output = compiled_model( + hidden_states.to(torch.bfloat16), + attention_mask.to(torch.bfloat16), + cos.to(torch.bfloat16), + sin.to(torch.bfloat16), + ) + elif self._compiled_buckets is not None and self._cpu_model is not None: + # seq_len too large for any compiled bucket — fall back to CPU + logger.warning( + f"seq_len={seq_len} exceeds all compiled buckets " + f"{sorted(self._compiled_buckets.keys())}; using CPU vision fallback." + ) + with torch.no_grad(): + vision_output = self._cpu_model( + hidden_states.to(torch.bfloat16), + attention_mask.to(torch.bfloat16), + cos.to(torch.bfloat16), + sin.to(torch.bfloat16), + ) + elif self._compiled_model is not None: + # Single compiled model (legacy) + vision_output = self._compiled_model( + hidden_states.to(torch.bfloat16), + attention_mask.to(torch.bfloat16), + cos.to(torch.bfloat16), + sin.to(torch.bfloat16), + ) + elif self._cpu_model is not None: + # CPU-only mode: run vision encoder on CPU (no bucketing/padding needed + # but we pad anyway for consistency with the same merger math) + with torch.no_grad(): + vision_output = self._cpu_model( + hidden_states.to(torch.bfloat16), + attention_mask.to(torch.bfloat16), + cos.to(torch.bfloat16), + sin.to(torch.bfloat16), + ) + else: + # NxDI traced model: takes (hidden_states, attention_mask, position_embeddings) + vision_output = self.model(hidden_states, attention_mask, (cos, sin)) + + # 7. Unpad: only keep valid merged tokens + merge_area = self.vision_config.spatial_merge_size**2 + total_merged_tokens = sum( + t + * (h // self.vision_config.spatial_merge_size) + * (w // self.vision_config.spatial_merge_size) + for t, h, w in image_grid_thw.tolist() + ) + vision_output = vision_output[:total_merged_tokens] + + return vision_output + + def _tiled_forward(self, pixel_values, image_grid_thw): + """Split a single-image (T=1, H, W) input into 2×2 spatial tiles of + (H/2, W/2) patches each. Encode each tile independently through + `forward`, then re-interleave the merged outputs so the returned + sequence matches the full-image row-major merged ordering that the + text model expects at image-token positions. + + Requires image_grid_thw shape (1, 3) with H and W both divisible by + 2 * spatial_merge_size. Each tile has (H/2)*(W/2) patch tokens and + must fit in a compiled bucket. + """ + assert image_grid_thw.shape[0] == 1, "tiled path supports single-image inputs" + T, H, W = image_grid_thw[0].tolist() + assert T == 1, "tiled path supports T=1 (still images)" + merge = self.vision_config.spatial_merge_size + assert H % (2 * merge) == 0 and W % (2 * merge) == 0, ( + f"H={H}, W={W} must be divisible by 2*spatial_merge_size={2 * merge}" + ) + Ht, Wt = H // 2, W // 2 # per-tile patch dims + + # Reshape pixel_values (H*W, patch_dim) into (H, W, patch_dim), split + # into 4 spatial quadrants, and flatten each back to (Ht*Wt, patch_dim). + patch_dim = pixel_values.shape[1] + pv_hw = pixel_values.reshape(H, W, patch_dim) + tiles = [] + for h_off in (0, Ht): + for w_off in (0, Wt): + tile_hw = pv_hw[h_off:h_off + Ht, w_off:w_off + Wt, :] + tiles.append(tile_hw.reshape(Ht * Wt, patch_dim)) + + # Encode each tile with tile-local grid_thw. Vision-internal RoPE runs + # per-tile with positions restarting at 0; the text model's global + # mRoPE (computed from the full image_grid_thw upstream) handles + # cross-tile positioning at scatter time. + tile_gthw = torch.tensor([[1, Ht, Wt]], dtype=image_grid_thw.dtype, + device=image_grid_thw.device) + tile_merged = [] + for tile_pv in tiles: + merged = self.forward(tile_pv, tile_gthw) + tile_merged.append(merged) # each (Ht/merge * Wt/merge, out_hidden) + + # Re-interleave merged tokens back into row-major order over the full + # (H/merge, W/merge) merged grid. Merged tiles produced by tile_i + # correspond to merged-block quadrants of the full grid. + mH_full, mW_full = H // merge, W // merge + mHt, mWt = Ht // merge, Wt // merge # per-tile merged dims + out_hidden = tile_merged[0].shape[-1] + + full = torch.zeros( + (mH_full, mW_full, out_hidden), + dtype=tile_merged[0].dtype, + device=tile_merged[0].device, + ) + for idx, (h_off, w_off) in enumerate( + [(0, 0), (0, mWt), (mHt, 0), (mHt, mWt)] + ): + full[h_off:h_off + mHt, w_off:w_off + mWt] = \ + tile_merged[idx].reshape(mHt, mWt, out_hidden) + + return full.reshape(mH_full * mW_full, out_hidden) + + +class NeuronQwen35VisionForImageEncoding(NeuronApplicationBase): + """Standalone application class for vision encoding (for testing).""" + + model_cls = NeuronQwen35VisionModel + model_wrapper_cls = NeuronQwen35VisionModelWrapper + + @staticmethod + def prepare_input_args(image_path, processor): + """Prepare vision inputs from an image path. + + Args: + image_path: Path to image file + processor: HF AutoProcessor + + Returns: + pixel_values, image_grid_thw + """ + from PIL import Image + + image = Image.open(image_path).convert("RGB") + inputs = processor(images=image, return_tensors="pt") + return inputs["pixel_values"], inputs["image_grid_thw"] diff --git a/contrib/models/Qwen3.5-35B-A3B/src/modeling_qwen35_vl.py b/contrib/models/Qwen3.5-35B-A3B/src/modeling_qwen35_vl.py new file mode 100644 index 00000000..cc2c7cbd --- /dev/null +++ b/contrib/models/Qwen3.5-35B-A3B/src/modeling_qwen35_vl.py @@ -0,0 +1,701 @@ +""" +Qwen3.5-27B / Qwen3.6-27B Vision-Language Model Orchestrator for NeuronX Distributed Inference. + +This is the top-level VL model that wires together: +- The vision encoder (modeling_qwen35_vision.py) +- The text decoder (modeling_qwen35.py, dense model with vision injection) + +It handles: +- Multimodal RoPE (mRoPE) with interleaved layout +- Vision embedding injection via scatter_by_index_put +- Separate compilation and loading of vision and text models +- The CTE+TKG generation loop with vision inputs + +Architecture follows the NxDI NeuronBaseForImageToText pattern established +by Qwen3-VL in SDK 2.28, adapted for Qwen3.5/3.6 dense model's unique features: +- No deepstack (Qwen3.5/3.6 does not use intermediate vision feature injection) +- DeltaNet linear attention layers in the text decoder +- Dense SwiGLU MLP layers in the text decoder +- Interleaved mRoPE (THWTHW... layout) instead of Qwen3-VL's section-based layout +""" + +import logging +import os +from typing import Optional + +import torch +import torch.nn.functional as F + +logger = logging.getLogger(__name__) + +# NxDI imports +try: + from neuronx_distributed_inference.models.image_to_text_model_base import ( + ImageToTextInferenceConfig, + NeuronBaseForImageToText, + ) + from neuronx_distributed_inference.models.config import NeuronConfig + + HAS_NXDI_VL = True +except ImportError: + HAS_NXDI_VL = False + logger.warning("NxDI VL base classes not available -- VL model requires SDK 2.28+") + +# Local imports +try: + from src.modeling_qwen35 import ( + NeuronQwen35ForCausalLM, + NeuronQwen35Model, + Qwen35InferenceConfig, + Qwen35ModelWrapper, + ) + from src.modeling_qwen35_vision import ( + NeuronQwen35VisionModel, + NeuronQwen35VisionModelWrapper, + ) +except ImportError: + from modeling_qwen35 import ( + NeuronQwen35ForCausalLM, + NeuronQwen35Model, + Qwen35InferenceConfig, + Qwen35ModelWrapper, + ) + from modeling_qwen35_vision import ( + NeuronQwen35VisionModel, + NeuronQwen35VisionModelWrapper, + ) + + +def get_rope_index( + input_ids, + image_grid_thw=None, + video_grid_thw=None, + attention_mask=None, + image_token_id=248056, + video_token_id=248057, + vision_start_token_id=248053, + spatial_merge_size=2, +): + """Compute 3D multimodal RoPE position IDs for Qwen3.5. + + Returns position_ids of shape (3, batch_size, seq_len) where: + - Axis 0: temporal position + - Axis 1: height position + - Axis 2: width position + + For text tokens, all 3 axes have the same sequential position. + For vision tokens, each axis encodes the spatial/temporal grid position. + + Also returns rope_deltas for use during TKG decoding. + + Adapted from HuggingFace Qwen3_5Model.get_rope_index(). + """ + if video_grid_thw is not None: + video_grid_thw = torch.repeat_interleave( + video_grid_thw, video_grid_thw[:, 0], dim=0 + ) + video_grid_thw[:, 0] = 1 + + image_grid_thw_list = ( + image_grid_thw.tolist() if image_grid_thw is not None else None + ) + video_grid_thw_list = ( + video_grid_thw.tolist() if video_grid_thw is not None else None + ) + + mrope_position_deltas = [] + total_input_ids = input_ids + + if attention_mask is None: + attention_mask = torch.ones_like(total_input_ids) + + position_ids = torch.zeros( + 3, + input_ids.shape[0], + input_ids.shape[1], + dtype=input_ids.dtype, + device=input_ids.device, + ) + + image_index, video_index = 0, 0 + attention_mask = attention_mask.to(total_input_ids.device) + + for i, ids in enumerate(total_input_ids): + ids = ids[attention_mask[i] == 1] + image_nums, video_nums = 0, 0 + + vision_start_indices = torch.argwhere(ids == vision_start_token_id).squeeze(1) + if len(vision_start_indices) > 0: + vision_tokens = ids[vision_start_indices + 1] + image_nums = (vision_tokens == image_token_id).sum() + video_nums = (vision_tokens == video_token_id).sum() + + input_tokens = ids.tolist() + llm_pos_ids_list = [] + st = 0 + remain_images, remain_videos = image_nums, video_nums + + for _ in range(image_nums + video_nums): + if image_token_id in input_tokens and remain_images > 0: + ed_image = input_tokens.index(image_token_id, st) + else: + ed_image = len(input_tokens) + 1 + if video_token_id in input_tokens and remain_videos > 0: + ed_video = input_tokens.index(video_token_id, st) + else: + ed_video = len(input_tokens) + 1 + + if ed_image < ed_video: + t, h, w = image_grid_thw_list[image_index] + image_index += 1 + remain_images -= 1 + ed = ed_image + else: + t, h, w = video_grid_thw_list[video_index] + video_index += 1 + remain_videos -= 1 + ed = ed_video + + llm_grid_t = t + llm_grid_h = h // spatial_merge_size + llm_grid_w = w // spatial_merge_size + + text_len = ed - st + st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 + llm_pos_ids_list.append( + torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx + ) + + t_index = ( + torch.arange(llm_grid_t) + .view(-1, 1) + .expand(-1, llm_grid_h * llm_grid_w) + .flatten() + ) + h_index = ( + torch.arange(llm_grid_h) + .view(1, -1, 1) + .expand(llm_grid_t, -1, llm_grid_w) + .flatten() + ) + w_index = ( + torch.arange(llm_grid_w) + .view(1, 1, -1) + .expand(llm_grid_t, llm_grid_h, -1) + .flatten() + ) + llm_pos_ids_list.append( + torch.stack([t_index, h_index, w_index]) + text_len + st_idx + ) + st = ed + llm_grid_t * llm_grid_h * llm_grid_w + + if st < len(input_tokens): + st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 + text_len = len(input_tokens) - st + llm_pos_ids_list.append( + torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx + ) + + llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) + position_ids[..., i, attention_mask[i] == 1] = llm_positions.to( + position_ids.device + ) + mrope_position_deltas.append(llm_positions.max() + 1 - len(total_input_ids[i])) + + mrope_position_deltas = torch.tensor( + mrope_position_deltas, device=input_ids.device + ).unsqueeze(1) + return position_ids, mrope_position_deltas + + +class Qwen35VLInferenceConfig: + """Configuration for the full VL model (text + vision). + + Wraps the existing Qwen35InferenceConfig for text and adds + vision-specific settings. + """ + + def __init__( + self, + text_config, + vision_config, + image_token_id=248056, + video_token_id=248057, + vision_start_token_id=248053, + vision_end_token_id=248054, + spatial_merge_size=2, + vision_seq_len_buckets=None, + **kwargs, + ): + """ + Args: + text_config: Qwen35InferenceConfig instance for the text decoder + vision_config: dict with vision encoder hyperparams (depth, hidden_size, etc.) + image_token_id: Token ID for image placeholder tokens + video_token_id: Token ID for video placeholder tokens + vision_start_token_id: Token ID for <|vision_start|> + vision_end_token_id: Token ID for <|vision_end|> + spatial_merge_size: How many patches are merged (2 = 2x2 = 4 patches merged) + vision_seq_len_buckets: List of vision sequence length buckets for compilation + """ + self.text_config = text_config + self.vision_config = vision_config + self.image_token_id = image_token_id + self.video_token_id = video_token_id + self.vision_start_token_id = vision_start_token_id + self.vision_end_token_id = vision_end_token_id + self.spatial_merge_size = spatial_merge_size + self.vision_seq_len_buckets = vision_seq_len_buckets or [1024, 4096, 16384] + + +class NeuronQwen35VLForCausalLM: + """Top-level VL model for Qwen3.5/3.6-27B on Neuron. + + This class manages: + - Separate compilation/loading of vision encoder and text decoder + - CPU-side mRoPE computation + - Vision embedding injection into text decoder + - The CTE+TKG generation loop + + Note: This is NOT an NeuronBaseForImageToText subclass because the + text decoder (NeuronQwen35ForCausalLM) has extensive custom overrides + (DeltaNet state management, custom forward, custom ModelWrapper) that + don't fit the base class pattern. Instead, this class composes the two + models and handles the VL orchestration directly. + """ + + def __init__(self, model_path, text_config, vision_config=None, processor=None): + """ + Args: + model_path: Path to HF model directory + text_config: Qwen35InferenceConfig for text decoder + vision_config: Qwen35VLInferenceConfig (or None for text-only) + processor: HF AutoProcessor for image preprocessing + """ + self.model_path = model_path + self.text_config = text_config + self.vl_config = vision_config + self.processor = processor + + # Text decoder (existing implementation) + self.text_model = NeuronQwen35ForCausalLM( + model_path=model_path, config=text_config + ) + + # Vision encoder (lazy init -- only built if vl_config provided) + self.vision_model_wrapper = None + if vision_config is not None: + self._init_vision_model(vision_config) + + # mRoPE state + self.rope_deltas = None + + def _init_vision_model(self, vl_config): + """Initialize the vision encoder wrapper.""" + from types import SimpleNamespace + + vision_cfg = SimpleNamespace(**vl_config.vision_config) + self.vision_model_wrapper = NeuronQwen35VisionModelWrapper( + config=vision_cfg, + model_cls=None, # Standalone mode (no NxDI parallel layers) + vision_seq_len_buckets=vl_config.vision_seq_len_buckets, + ) + self._vl_config = vl_config + + def compile(self, compiled_model_path): + """Compile both text and vision models. + + For the vision encoder, use compile_vision_encoder.py separately + (standalone torch_neuronx.trace compilation). Then use load() to + load the pre-compiled vision encoder. + """ + # Compile text decoder + text_path = os.path.join(compiled_model_path, "text_model") + os.makedirs(text_path, exist_ok=True) + self.text_model.compile(text_path) + + # Vision encoder is compiled separately via compile_vision_encoder.py + if self.vision_model_wrapper is not None: + logger.info( + "Vision encoder must be compiled separately using " + "compile_vision_encoder.py. Use load() to load the " + "pre-compiled vision encoder." + ) + + def load(self, compiled_model_path, vision_compiled_path=None): + """Load both compiled models. + + Args: + compiled_model_path: Path to compiled text model (or parent dir) + vision_compiled_path: Path to compiled vision encoder .pt file. + If None, looks for 'vision_encoder.pt' in compiled_model_path. + """ + text_path = os.path.join(compiled_model_path, "text_model") + if os.path.exists(text_path): + self.text_model.load(text_path) + else: + # Backward compatibility: text model compiled at root + self.text_model.load(compiled_model_path) + + # Load vision encoder + if self.vision_model_wrapper is not None: + if vision_compiled_path is None: + vision_compiled_path = os.path.join( + compiled_model_path, "vision_encoder.pt" + ) + if os.path.exists(vision_compiled_path): + self.vision_model_wrapper.load_compiled(vision_compiled_path) + # Also load CPU-side weights (patch_embed, pos_embed) + self.vision_model_wrapper.load_vision_weights_from_hf(self.model_path) + logger.info("Vision encoder loaded from pre-compiled model") + else: + logger.warning( + f"No compiled vision encoder found at {vision_compiled_path}. " + "Vision encoding will not be available." + ) + + # Qwen3.5 stop token IDs (loaded from config/tokenizer) + _DEFAULT_EOS_TOKEN_IDS = { + 248044, # <|endoftext|> -- text config eos_token_id + 248046, # <|im_end|> -- tokenizer eos_token / end of assistant turn + } + + def generate( + self, + input_ids, + attention_mask=None, + pixel_values=None, + image_grid_thw=None, + video_grid_thw=None, + max_new_tokens=32, + temperature=0.0, + top_p=1.0, + top_k=0, + eos_token_ids=None, + **kwargs, + ): + """Generate text from text and/or vision inputs. + + Args: + input_ids: (batch_size, seq_len) token IDs + attention_mask: (batch_size, seq_len) attention mask + pixel_values: Vision pixel values from HF processor (or None for text-only) + image_grid_thw: (num_images, 3) grid dimensions + video_grid_thw: (num_videos, 3) grid dimensions + max_new_tokens: Maximum new tokens to generate + temperature: Sampling temperature (0.0 = greedy/argmax) + top_p: Nucleus sampling threshold (1.0 = disabled) + top_k: Top-k sampling (0 = disabled) + eos_token_ids: Set of token IDs to stop generation on + (default: {248044, 248046}) + + Returns: + generated_ids: (batch_size, seq_len + max_new_tokens) token IDs + """ + if eos_token_ids is None: + eos_token_ids = self._DEFAULT_EOS_TOKEN_IDS + + # Reset text model state for a fresh generation. + # This ensures CTE runs (not TKG) even if a prior generate() was called. + # DeltaNet recurrent states don't need explicit zeroing because the CTE + # NKI kernel always starts from zero state. + self.text_model.reset() + + has_vision = pixel_values is not None and pixel_values.numel() > 0 + + # Step 1: Compute 3D mRoPE position IDs + if has_vision and self._vl_config is not None: + position_ids, self.rope_deltas = get_rope_index( + input_ids, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + attention_mask=attention_mask, + image_token_id=self._vl_config.image_token_id, + video_token_id=self._vl_config.video_token_id, + vision_start_token_id=self._vl_config.vision_start_token_id, + spatial_merge_size=self._vl_config.spatial_merge_size, + ) + else: + # Text-only: use standard sequential position IDs + seq_len = input_ids.shape[1] + position_ids = torch.arange(seq_len).unsqueeze(0) + self.rope_deltas = None + + # Step 2: Run vision encoder and prepare injection args + llava_args = [] + batch_size = input_ids.shape[0] + if has_vision and self.vision_model_wrapper is not None: + # The vision encoder processes both image and video frames identically + # (they share the same ViT architecture). The HF processor outputs a + # single pixel_values tensor for images, and video frames are treated + # as multiple images with temporal grid > 1. + vision_embeddings = self.vision_model_wrapper(pixel_values, image_grid_thw) + # vision_embeddings: (total_merged_tokens, out_hidden_size) + + # Build vision_mask: boolean mask of ALL vision token positions + # (both image_token_id and video_token_id placeholders) + image_token_id = self._vl_config.image_token_id + video_token_id = self._vl_config.video_token_id + vision_bool_mask = (input_ids == image_token_id) | ( + input_ids == video_token_id + ) # (BS, seq_len) + + # For batch_size=1 (primary path): extract positions from batch element 0. + # For batch_size>1: each element may have different image token positions; + # we'd need per-element scatter. Currently only batch_size=1 is supported + # for VL (the compiled model uses batch_size=1 for CTE). + if batch_size > 1: + logger.warning( + "VL generation with batch_size > 1 is not fully supported. " + "Using batch element 0 for vision scatter positions." + ) + + positions = ( + vision_bool_mask[0].nonzero(as_tuple=False).squeeze(-1) + ) # (n_vision_tokens,) + + # Reshape vision_embeddings to (1, n_vision_tokens, hidden_size) + n_vis = positions.shape[0] + hidden_size = vision_embeddings.shape[-1] + vis_emb = vision_embeddings[:n_vis].unsqueeze(0) # (1, n_vis, hidden) + + # Pad to match the compiled CTE bucket size (not input length). + # The traced graph expects vision_embeddings and vision_mask with + # sequence dim == CTE bucket, so query the text model's neuron_config. + cte_bucket = getattr( + self.text_config.neuron_config, "seq_len", input_ids.shape[1] + ) + seq_len = input_ids.shape[1] + pad_limit = cte_bucket + + # Qwen3-VL pad convention: + # - vision_embeddings pad rows: zeros + # - vision_mask pad slots: point at pad_limit-1 (guaranteed to be + # a padded input slot, attention_mask==0 there so scatter is a no-op) + if n_vis < pad_limit: + pad_emb = torch.zeros( + (1, pad_limit - n_vis, hidden_size), dtype=vis_emb.dtype, + ) + vis_emb_padded = torch.cat([vis_emb, pad_emb], dim=1) + else: + vis_emb_padded = vis_emb[:, :pad_limit] + + positions_padded = torch.full( + (1, pad_limit, 1), + fill_value=pad_limit - 1, + dtype=torch.int32, + ) + positions_padded[0, :n_vis, 0] = positions[:pad_limit].to(torch.int32) + + llava_args = [vis_emb_padded, positions_padded] + + # Append 3D mRoPE position IDs for the text model. + # position_ids shape: (3, batch_size, seq_len) from get_rope_index. + # Pad to CTE bucket size on the seq dim. + if position_ids.ndim == 3: + pids = position_ids.to(torch.int32).contiguous() + if pids.shape[2] < pad_limit: + pad_amount = pad_limit - pids.shape[2] + pad_tensor = torch.zeros( + (pids.shape[0], pids.shape[1], pad_amount), dtype=torch.int32 + ) + pids = torch.cat([pids, pad_tensor], dim=2) + else: + pids = pids[:, :, :pad_limit].contiguous() + llava_args.append(pids) + else: + vision_embeddings = None + + # Step 3: Context encoding (prefill) + generated_ids = input_ids.clone() + + # CRITICAL: Always pass an explicit attention_mask for CTE. + # The base class _infer_attention_mask() assumes sequential position_ids + # (position_ids[i] >= i). When position_ids come from mRoPE temporal + # axis (non-sequential, e.g., all vision tokens share position 4), + # the inferred mask incorrectly masks out most of the sequence. + # Fix: provide a real all-ones mask for the actual token positions. + if attention_mask is None: + attention_mask = torch.ones_like(input_ids) + + # For slot 2 (position_ids): use SEQUENTIAL positions regardless of mRoPE. + # Slot 2 is only used for: (1) logit position selection via torch.max(), + # (2) attention mask inference (which we bypass with explicit mask above). + # The actual RoPE computation uses slot 21 (rotary_position_ids) from + # _get_model_outputs, NOT slot 2. Using sequential slot 2 ensures + # correct logit selection and avoids any position_ids-related issues. + seq_len = input_ids.shape[1] + cte_position_ids = torch.arange(seq_len, dtype=torch.long).unsqueeze(0) + + def _next_token_from_output(output, temperature, top_p, top_k): + # NxDI with on_device_sampling puts the sampled token into + # `output.tokens` and leaves `output.logits` as None. + # Without on_device_sampling it fills `output.logits` [B, T, V]. + head = None + if hasattr(output, "tokens") and output.tokens is not None: + head = output.tokens + elif hasattr(output, "logits") and output.logits is not None: + head = output.logits + elif isinstance(output, (tuple, list)): + head = output[0] + if head is None: + print(f"[VL] output type: {type(output).__name__}") + if hasattr(output, "keys"): + for k in list(output.keys())[:16]: + v = output[k] + print(f"[VL] {k} -> type={type(v).__name__} " + f"shape={getattr(v, 'shape', None)}") + raise RuntimeError("model output has no tokens/logits") + # Sampled token from on_device_sampling: shape [B], [B, 1], or [1, B] + if head.ndim == 3: # logits [B, T, V] + return self._sample_token(head[:, -1, :], temperature, top_p, top_k) + if head.ndim == 2: # [B, 1] or [1, B] + if head.shape[-1] == 1: + return head.squeeze(-1) + # [1, B] shape from on_device_sampling + return head.reshape(-1)[:1] if head.shape[0] == 1 else head[:, 0] + if head.ndim == 1: # [B] sampled token + return head + if head.ndim == 3: # logits [B, T, V] + return self._sample_token(head[:, -1, :], temperature, top_p, top_k) + raise RuntimeError(f"Unexpected output head shape {tuple(head.shape)}") + + with torch.no_grad(): + output = self.text_model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=cte_position_ids, + output_attentions=False, + output_hidden_states=False, + return_dict=False, + llava_args=llava_args, + ) + + next_token = _next_token_from_output(output, temperature, top_p, top_k) + generated_ids = torch.cat([generated_ids, next_token.unsqueeze(-1)], dim=-1) + + # Check EOS after first token + if next_token.item() in eos_token_ids: + return generated_ids + + # Step 4: Token generation (TKG) loop + for _ in range(max_new_tokens - 1): + pos_ids = torch.tensor([[generated_ids.shape[1] - 1]]) + if self.rope_deltas is not None: + pos_ids = pos_ids + self.rope_deltas + + last_token = generated_ids[:, -1:] + with torch.no_grad(): + output = self.text_model( + input_ids=last_token, + position_ids=pos_ids, + output_attentions=False, + output_hidden_states=False, + return_dict=False, + ) + next_token = _next_token_from_output(output, temperature, top_p, top_k) + generated_ids = torch.cat([generated_ids, next_token.unsqueeze(-1)], dim=-1) + + # Stop on EOS + if next_token.item() in eos_token_ids: + break + + return generated_ids + + @staticmethod + def _sample_token(logits, temperature=0.0, top_p=1.0, top_k=0): + """Sample a token from logits with optional temperature/top-p/top-k. + + Args: + logits: (batch_size, vocab_size) unnormalized logits + temperature: Sampling temperature. 0.0 = greedy (argmax). + top_p: Nucleus sampling threshold. 1.0 = disabled. + top_k: Top-k filtering. 0 = disabled. + + Returns: + token_id: (batch_size,) sampled token IDs + """ + if temperature <= 0.0: + return torch.argmax(logits, dim=-1) + + # Apply temperature + logits = logits / temperature + + # Top-k filtering + if top_k > 0: + top_k = min(top_k, logits.shape[-1]) + indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None] + logits[indices_to_remove] = float("-inf") + + # Top-p (nucleus) filtering + if top_p < 1.0: + sorted_logits, sorted_indices = torch.sort(logits, descending=True) + cumulative_probs = torch.cumsum( + torch.softmax(sorted_logits, dim=-1), dim=-1 + ) + # Remove tokens with cumulative probability above the threshold + sorted_indices_to_remove = cumulative_probs > top_p + # Shift right so the first token above threshold is kept + sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[ + ..., :-1 + ].clone() + sorted_indices_to_remove[..., 0] = False + # Scatter back to original indexing + indices_to_remove = sorted_indices_to_remove.scatter( + -1, sorted_indices, sorted_indices_to_remove + ) + logits[indices_to_remove] = float("-inf") + + # Sample from the filtered distribution + probs = torch.softmax(logits, dim=-1) + return torch.multinomial(probs, num_samples=1).squeeze(-1) + + @staticmethod + def prepare_input_args(text_prompt, image_path, processor, role="user"): + """Prepare inputs for vision+text generation. + + Args: + text_prompt: Text prompt string + image_path: Path to image file (or None for text-only) + processor: HF AutoProcessor + role: Message role (default "user") + + Returns: + input_ids, attention_mask, vision_inputs dict + """ + content = [] + if image_path is not None: + import base64 + from pathlib import Path + + image_data = Path(image_path).read_bytes() + b64 = base64.b64encode(image_data).decode("utf-8") + content.append( + { + "type": "image", + "url": f"data:image/jpeg;base64,{b64}", + } + ) + content.append({"type": "text", "text": text_prompt}) + + messages = [{"role": role, "content": content}] + inputs = processor.apply_chat_template( + messages, + tokenize=True, + add_generation_prompt=True, + return_tensors="pt", + return_dict=True, + ) + + input_ids = inputs["input_ids"] + attention_mask = inputs.get("attention_mask", torch.ones_like(input_ids)) + + vision_inputs = {} + if "pixel_values" in inputs: + vision_inputs["pixel_values"] = inputs["pixel_values"] + if "image_grid_thw" in inputs: + vision_inputs["image_grid_thw"] = inputs["image_grid_thw"] + if "video_grid_thw" in inputs: + vision_inputs["video_grid_thw"] = inputs["video_grid_thw"] + + return input_ids, attention_mask, vision_inputs diff --git a/contrib/models/Qwen3.5-35B-A3B/src/nki_kernels/__init__.py b/contrib/models/Qwen3.5-35B-A3B/src/nki_kernels/__init__.py new file mode 100644 index 00000000..7e78cdb9 --- /dev/null +++ b/contrib/models/Qwen3.5-35B-A3B/src/nki_kernels/__init__.py @@ -0,0 +1,10 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Custom NKI kernels for Qwen3.5-27B / Qwen3.6-27B DeltaNet layers. + +Contains three kernel implementations: +- nki_deltanet: Per-token recurrent kernel (used for token generation) +- nki_deltanet_chunked: Per-chunk kernel (legacy, superseded by fused) +- nki_deltanet_fused: Fused single-kernel chunked forward (used for context encoding) +""" diff --git a/contrib/models/Qwen3.5-35B-A3B/src/nki_kernels/nki_deltanet.py b/contrib/models/Qwen3.5-35B-A3B/src/nki_kernels/nki_deltanet.py new file mode 100644 index 00000000..5aaf172f --- /dev/null +++ b/contrib/models/Qwen3.5-35B-A3B/src/nki_kernels/nki_deltanet.py @@ -0,0 +1,607 @@ +"""NKI kernels for DeltaNet gated delta rule recurrent forward. + +NKI v3 (neuronx-cc 2.26.6360 / nki 0.5.0). Processes a SINGLE (batch, head) pair per kernel call. +The caller loops over (B, H) in PyTorch and calls this kernel for each pair. + +Input layout: All inputs are 2D contiguous tensors (S, 128). +Each call processes one (batch, head) element's full sequence. + +k_dim = v_dim = 128, which matches SBUF tile partition dimension exactly. +g and beta are scalars per token, expanded to (S, 128) by the caller. + +Two kernel variants: + deltanet_recurrent_fwd -- returns output only (original) + deltanet_recurrent_fwd_state -- returns (output, final_state) for CTE->TKG carry-over + deltanet_recurrent_step -- one-token state-in/state-out decode step + deltanet_recurrent_step_batched -- one-token batched-head decode step +""" + +import nki +import nki.isa as nisa +import nki.language as nl + +# Partition dimension max (NeuronCore SBUF tile width) +P_MAX = 128 + +# Shuffle mask: broadcast partition 0 to all partitions in a 32-wide group +_BROADCAST_MASK = [0] * 32 + + +@nki.jit +def _deltanet_recurrent_step_batched_kernel( + query: nl.ndarray, # (BH, 128) float32 + key: nl.ndarray, # (BH, 128) float32 + value: nl.ndarray, # (BH, 128) float32 + g_in: nl.ndarray, # (BH, 1) float32, log-decay scalar per head + beta_in: nl.ndarray, # (BH, 1) float32, write-gate scalar per head + state_in: nl.ndarray, # (BH * 128, 128) float32/bfloat16 +): + """Single-launch batched-head one-token DeltaNet decode step. + + The installed NKI framework on the compile hosts uses ``kernel[...]`` for + LNC selection, not custom-op SPMD grids. Keep one framework custom call by + looping over flattened ``(batch, value_head)`` rows inside the kernel. + """ + batch_heads, dim = query.shape + + output = nl.ndarray(query.shape, dtype=query.dtype, buffer=nl.shared_hbm) + state_out = nl.ndarray(state_in.shape, dtype=state_in.dtype, buffer=nl.shared_hbm) + + for bh in nl.sequential_range(batch_heads): + head_offset = bh * dim + state_offset = bh * P_MAX + + q_t = nl.ndarray((P_MAX, 1), dtype=query.dtype, buffer=nl.sbuf) + nisa.dma_copy( + dst=q_t, + src=query.ap(pattern=[[1, P_MAX]], offset=head_offset), + dge_mode=nisa.dge_mode.hwdge, + ) + + k_t = nl.ndarray((P_MAX, 1), dtype=key.dtype, buffer=nl.sbuf) + nisa.dma_copy( + dst=k_t, + src=key.ap(pattern=[[1, P_MAX]], offset=head_offset), + dge_mode=nisa.dge_mode.hwdge, + ) + + v_t = nl.ndarray((P_MAX, 1), dtype=value.dtype, buffer=nl.sbuf) + nisa.dma_copy( + dst=v_t, + src=value.ap(pattern=[[1, P_MAX]], offset=head_offset), + dge_mode=nisa.dge_mode.hwdge, + ) + + g_t = nl.ndarray((P_MAX, 1), dtype=g_in.dtype, buffer=nl.sbuf) + g_scalar = nl.ndarray((1, 1), dtype=g_in.dtype, buffer=nl.sbuf) + nisa.dma_copy(dst=g_scalar, src=g_in.ap(pattern=[[1, 1]], offset=bh)) + for i_shuf in nl.static_range(P_MAX // 32): + nisa.nc_stream_shuffle( + src=g_scalar[0:1, 0:1], + dst=g_t[i_shuf * 32 : i_shuf * 32 + 32, 0:1], + shuffle_mask=_BROADCAST_MASK, + ) + + beta_t = nl.ndarray((P_MAX, 1), dtype=beta_in.dtype, buffer=nl.sbuf) + beta_scalar = nl.ndarray((1, 1), dtype=beta_in.dtype, buffer=nl.sbuf) + nisa.dma_copy(dst=beta_scalar, src=beta_in.ap(pattern=[[1, 1]], offset=bh)) + for i_shuf in nl.static_range(P_MAX // 32): + nisa.nc_stream_shuffle( + src=beta_scalar[0:1, 0:1], + dst=beta_t[i_shuf * 32 : i_shuf * 32 + 32, 0:1], + shuffle_mask=_BROADCAST_MASK, + ) + + state = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy( + dst=state, + src=state_in[state_offset : state_offset + P_MAX, 0:dim], + ) + + exp_g = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.activation(dst=exp_g, op=nl.exp, data=g_t, bias=None, scale=1.0) + + state_decayed = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=state_decayed, + data=state, + op0=nl.multiply, + operand0=exp_g, + engine=nisa.vector_engine, + ) + + kv_mem_psum = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=kv_mem_psum, stationary=state_decayed, moving=k_t) + kv_mem = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=kv_mem, src=kv_mem_psum) + + v_sub = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=v_sub, data1=v_t, data2=kv_mem, op=nl.subtract) + + delta = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=delta, + data=v_sub, + op0=nl.multiply, + operand0=beta_t, + engine=nisa.vector_engine, + ) + + delta_row_psum = nl.ndarray((1, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=delta_row_psum, data=delta) + + delta_row_sb = nl.ndarray((1, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=delta_row_sb, src=delta_row_psum) + + delta_broadcast = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + for i_shuf in nl.static_range(P_MAX // 32): + nisa.nc_stream_shuffle( + src=delta_row_sb[0:1, 0:P_MAX], + dst=delta_broadcast[i_shuf * 32 : i_shuf * 32 + 32, 0:P_MAX], + shuffle_mask=_BROADCAST_MASK, + ) + + state_new = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.scalar_tensor_tensor( + dst=state_new, + data=delta_broadcast, + op0=nl.multiply, + operand0=k_t, + op1=nl.add, + operand1=state_decayed, + ) + + o_t_psum = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=o_t_psum, stationary=state_new, moving=q_t) + o_t = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=o_t, src=o_t_psum) + + nisa.dma_copy( + dst=output.ap(pattern=[[1, dim]], offset=head_offset), + src=o_t, + dge_mode=nisa.dge_mode.hwdge, + ) + nisa.dma_copy( + dst=state_out[state_offset : state_offset + P_MAX, 0:dim], + src=state_new, + ) + + return output, state_out + + +def deltanet_recurrent_step_batched( + query, + key, + value, + g_in, + beta_in, + state_in, +): + """Launch the one-token DeltaNet decode step across flattened BH heads.""" + return _deltanet_recurrent_step_batched_kernel( + query, + key, + value, + g_in, + beta_in, + state_in, + ) + + +@nki.jit +def deltanet_recurrent_step( + query: nl.ndarray, # (1, 128) float32 + key: nl.ndarray, # (1, 128) float32 + value: nl.ndarray, # (1, 128) float32 + g_in: nl.ndarray, # (1, 128) float32, log-decay broadcast to 128 + beta_in: nl.ndarray, # (1, 128) float32, write gate broadcast to 128 + state_in: nl.ndarray, # (128, 128) float32 +): + """Stateful one-token DeltaNet decode step. + + This is the token-generation equivalent of one iteration from + ``deltanet_recurrent_fwd_state``. The caller supplies the recurrent + state restored from the decode cache and receives the updated state. + + Returns: + output: (1, 128) float32 + state_out: (128, 128) float32 + """ + _, dim = query.shape + + output = nl.ndarray((1, dim), dtype=query.dtype, buffer=nl.shared_hbm) + state_out = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.shared_hbm) + + q_t = nl.ndarray((P_MAX, 1), dtype=query.dtype, buffer=nl.sbuf) + nisa.dma_copy(dst=q_t, src=query.ap(pattern=[[1, P_MAX]], offset=0)) + + k_t = nl.ndarray((P_MAX, 1), dtype=key.dtype, buffer=nl.sbuf) + nisa.dma_copy(dst=k_t, src=key.ap(pattern=[[1, P_MAX]], offset=0)) + + v_t = nl.ndarray((P_MAX, 1), dtype=value.dtype, buffer=nl.sbuf) + nisa.dma_copy(dst=v_t, src=value.ap(pattern=[[1, P_MAX]], offset=0)) + + g_t = nl.ndarray((P_MAX, 1), dtype=g_in.dtype, buffer=nl.sbuf) + nisa.dma_copy(dst=g_t, src=g_in.ap(pattern=[[1, P_MAX]], offset=0)) + + beta_t = nl.ndarray((P_MAX, 1), dtype=beta_in.dtype, buffer=nl.sbuf) + nisa.dma_copy(dst=beta_t, src=beta_in.ap(pattern=[[1, P_MAX]], offset=0)) + + state = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=state, src=state_in) + + exp_g = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.activation(dst=exp_g, op=nl.exp, data=g_t, bias=None, scale=1.0) + + state_decayed = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=state_decayed, + data=state, + op0=nl.multiply, + operand0=exp_g, + engine=nisa.vector_engine, + ) + nisa.tensor_copy(dst=state, src=state_decayed) + + kv_mem_psum = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=kv_mem_psum, stationary=state, moving=k_t) + kv_mem = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=kv_mem, src=kv_mem_psum) + + v_sub = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=v_sub, data1=v_t, data2=kv_mem, op=nl.subtract) + + delta = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=delta, + data=v_sub, + op0=nl.multiply, + operand0=beta_t, + engine=nisa.vector_engine, + ) + + delta_row_psum = nl.ndarray((1, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=delta_row_psum, data=delta) + + delta_row_sb = nl.ndarray((1, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=delta_row_sb, src=delta_row_psum) + + delta_broadcast = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + for i_shuf in nl.static_range(P_MAX // 32): + nisa.nc_stream_shuffle( + src=delta_row_sb[0:1, 0:P_MAX], + dst=delta_broadcast[i_shuf * 32 : i_shuf * 32 + 32, 0:P_MAX], + shuffle_mask=_BROADCAST_MASK, + ) + + outer_prod = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=outer_prod, + data=delta_broadcast, + op0=nl.multiply, + operand0=k_t, + engine=nisa.vector_engine, + ) + + state_new = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=state_new, data1=state, data2=outer_prod, op=nl.add) + nisa.tensor_copy(dst=state, src=state_new) + + o_t_psum = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=o_t_psum, stationary=state, moving=q_t) + o_t = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=o_t, src=o_t_psum) + + nisa.dma_copy(dst=output.ap(pattern=[[1, dim]], offset=0), src=o_t) + nisa.dma_copy(dst=state_out, src=state) + + return output, state_out + + +@nki.jit +def deltanet_recurrent_fwd( + query: nl.ndarray, # (S, 128) float32 + key: nl.ndarray, # (S, 128) float32 + value: nl.ndarray, # (S, 128) float32 + g_in: nl.ndarray, # (S, 128) float32, log-decay broadcast to 128 + beta_in: nl.ndarray, # (S, 128) float32, write gate broadcast to 128 +) -> nl.ndarray: + """NKI kernel for DeltaNet recurrent forward -- single (batch, head). + + Iterates over sequence tokens with sequential_range. + State matrix (128 x 128) lives in SBUF. + + Args: + query: (S, 128) float32 + key: (S, 128) float32 + value: (S, 128) float32 + g_in: (S, 128) float32 + beta_in: (S, 128) float32 + + Returns: + output: (S, 128) float32 + """ + seq_len, dim = query.shape + + # Output tensor in HBM + output = nl.ndarray((seq_len, dim), dtype=query.dtype, buffer=nl.shared_hbm) + + # Stride: for 2D (S, D), dim0 stride = D=128, dim1 stride = 1 + seq_stride = dim + + # Initialize recurrent state in SBUF: (128, 128) + state = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.memset(dst=state, value=0.0) + + # Sequential loop over tokens (state-dependent) + for t in nl.sequential_range(seq_len): + tok_offset = t * seq_stride + + # ---- Load inputs for token t ---- + q_t = nl.ndarray((P_MAX, 1), dtype=query.dtype, buffer=nl.sbuf) + nisa.dma_copy( + dst=q_t, + src=query.ap(pattern=[[1, P_MAX]], offset=tok_offset), + ) + + k_t = nl.ndarray((P_MAX, 1), dtype=key.dtype, buffer=nl.sbuf) + nisa.dma_copy( + dst=k_t, + src=key.ap(pattern=[[1, P_MAX]], offset=tok_offset), + ) + + v_t = nl.ndarray((P_MAX, 1), dtype=value.dtype, buffer=nl.sbuf) + nisa.dma_copy( + dst=v_t, + src=value.ap(pattern=[[1, P_MAX]], offset=tok_offset), + ) + + g_t = nl.ndarray((P_MAX, 1), dtype=g_in.dtype, buffer=nl.sbuf) + nisa.dma_copy( + dst=g_t, + src=g_in.ap(pattern=[[1, P_MAX]], offset=tok_offset), + ) + + beta_t = nl.ndarray((P_MAX, 1), dtype=beta_in.dtype, buffer=nl.sbuf) + nisa.dma_copy( + dst=beta_t, + src=beta_in.ap(pattern=[[1, P_MAX]], offset=tok_offset), + ) + + # ---- Step 1: Decay state -- state = state * exp(g_t) ---- + exp_g = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.activation(dst=exp_g, op=nl.exp, data=g_t, bias=None, scale=1.0) + + state_decayed = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=state_decayed, + data=state, + op0=nl.multiply, + operand0=exp_g, + engine=nisa.vector_engine, + ) + nisa.tensor_copy(dst=state, src=state_decayed) + + # ---- Step 2: Read memory -- kv_mem = state^T @ k_t ---- + kv_mem_psum = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=kv_mem_psum, stationary=state, moving=k_t) + kv_mem = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=kv_mem, src=kv_mem_psum) + + # ---- Step 3: delta = (v_t - kv_mem) * beta_t ---- + v_sub = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=v_sub, data1=v_t, data2=kv_mem, op=nl.subtract) + + delta = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=delta, + data=v_sub, + op0=nl.multiply, + operand0=beta_t, + engine=nisa.vector_engine, + ) + + # ---- Step 4: state += outer(k_t, delta) ---- + # Broadcast multiply: outer[i,j] = k_t[i] * delta[j] + # 1) Transpose delta (128,1) -> (1,128) in PSUM + # 2) Copy PSUM (1,128) -> SBUF (128,128) -- partition broadcast + # 3) Multiply by k_t (128,1) which broadcasts across free dim + # This avoids the nc_matmul P=1 outer product (wastes 127/128 TE lanes). + + # Transpose delta to get values along free dimension + delta_row_psum = nl.ndarray((1, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=delta_row_psum, data=delta) + + # Copy PSUM (1, 128) -> SBUF (1, 128) first (nki requires matching P dims) + delta_row_sb = nl.ndarray((1, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=delta_row_sb, src=delta_row_psum) + + # Broadcast (1, 128) SBUF -> (128, 128) SBUF via nc_stream_shuffle + # Each partition row gets the same delta values + delta_broadcast = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + for i_shuf in nl.static_range(P_MAX // 32): + nisa.nc_stream_shuffle( + src=delta_row_sb[0:1, 0:P_MAX], + dst=delta_broadcast[i_shuf * 32 : i_shuf * 32 + 32, 0:P_MAX], + shuffle_mask=_BROADCAST_MASK, + ) + + # Element-wise multiply: outer[i,j] = delta_broadcast[i,j] * k_t[i,0] + # tensor_scalar broadcasts (P,1) k_t across all F columns + outer_prod = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=outer_prod, + data=delta_broadcast, + op0=nl.multiply, + operand0=k_t, + engine=nisa.vector_engine, + ) + + # Accumulate into state + state_new = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=state_new, data1=state, data2=outer_prod, op=nl.add) + nisa.tensor_copy(dst=state, src=state_new) + + # ---- Step 5: o_t = state^T @ q_t ---- + o_t_psum = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=o_t_psum, stationary=state, moving=q_t) + o_t = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=o_t, src=o_t_psum) + + # ---- Store output for token t ---- + nisa.dma_copy( + dst=output.ap(pattern=[[1, dim]], offset=tok_offset), + src=o_t, + ) + + return output + + +@nki.jit +def deltanet_recurrent_fwd_state( + query: nl.ndarray, # (S, 128) float32 + key: nl.ndarray, # (S, 128) float32 + value: nl.ndarray, # (S, 128) float32 + g_in: nl.ndarray, # (S, 128) float32, log-decay broadcast to 128 + beta_in: nl.ndarray, # (S, 128) float32, write gate broadcast to 128 +): + """NKI kernel for DeltaNet recurrent forward with final state output. + + Same recurrence as deltanet_recurrent_fwd, but ALSO writes the final + recurrent state (128, 128) to an output HBM buffer. This enables + CTE -> TKG state carry-over. + + Returns: + output: (S, 128) float32 -- per-token output + final_state: (128, 128) float32 -- recurrent state after last token + """ + seq_len, dim = query.shape + + # Output tensors in HBM + output = nl.ndarray((seq_len, dim), dtype=query.dtype, buffer=nl.shared_hbm) + final_state = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.shared_hbm) + + # Stride: for 2D (S, D), dim0 stride = D=128, dim1 stride = 1 + seq_stride = dim + + # Initialize recurrent state in SBUF: (128, 128) + state = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.memset(dst=state, value=0.0) + + # Sequential loop over tokens (state-dependent) + for t in nl.sequential_range(seq_len): + tok_offset = t * seq_stride + + # ---- Load inputs for token t ---- + q_t = nl.ndarray((P_MAX, 1), dtype=query.dtype, buffer=nl.sbuf) + nisa.dma_copy( + dst=q_t, + src=query.ap(pattern=[[1, P_MAX]], offset=tok_offset), + ) + + k_t = nl.ndarray((P_MAX, 1), dtype=key.dtype, buffer=nl.sbuf) + nisa.dma_copy( + dst=k_t, + src=key.ap(pattern=[[1, P_MAX]], offset=tok_offset), + ) + + v_t = nl.ndarray((P_MAX, 1), dtype=value.dtype, buffer=nl.sbuf) + nisa.dma_copy( + dst=v_t, + src=value.ap(pattern=[[1, P_MAX]], offset=tok_offset), + ) + + g_t = nl.ndarray((P_MAX, 1), dtype=g_in.dtype, buffer=nl.sbuf) + nisa.dma_copy( + dst=g_t, + src=g_in.ap(pattern=[[1, P_MAX]], offset=tok_offset), + ) + + beta_t = nl.ndarray((P_MAX, 1), dtype=beta_in.dtype, buffer=nl.sbuf) + nisa.dma_copy( + dst=beta_t, + src=beta_in.ap(pattern=[[1, P_MAX]], offset=tok_offset), + ) + + # ---- Step 1: Decay state -- state = state * exp(g_t) ---- + exp_g = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.activation(dst=exp_g, op=nl.exp, data=g_t, bias=None, scale=1.0) + + state_decayed = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=state_decayed, + data=state, + op0=nl.multiply, + operand0=exp_g, + engine=nisa.vector_engine, + ) + nisa.tensor_copy(dst=state, src=state_decayed) + + # ---- Step 2: Read memory -- kv_mem = state^T @ k_t ---- + kv_mem_psum = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=kv_mem_psum, stationary=state, moving=k_t) + kv_mem = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=kv_mem, src=kv_mem_psum) + + # ---- Step 3: delta = (v_t - kv_mem) * beta_t ---- + v_sub = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=v_sub, data1=v_t, data2=kv_mem, op=nl.subtract) + + delta = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=delta, + data=v_sub, + op0=nl.multiply, + operand0=beta_t, + engine=nisa.vector_engine, + ) + + # ---- Step 4: state += outer(k_t, delta) ---- + # Broadcast multiply: outer[i,j] = k_t[i] * delta[j] + delta_row_psum = nl.ndarray((1, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=delta_row_psum, data=delta) + + # Copy PSUM (1, 128) -> SBUF (1, 128) first (nki requires matching P dims) + delta_row_sb = nl.ndarray((1, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=delta_row_sb, src=delta_row_psum) + + # Broadcast (1, 128) SBUF -> (128, 128) SBUF via nc_stream_shuffle + delta_broadcast = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + for i_shuf in nl.static_range(P_MAX // 32): + nisa.nc_stream_shuffle( + src=delta_row_sb[0:1, 0:P_MAX], + dst=delta_broadcast[i_shuf * 32 : i_shuf * 32 + 32, 0:P_MAX], + shuffle_mask=_BROADCAST_MASK, + ) + + outer_prod = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=outer_prod, + data=delta_broadcast, + op0=nl.multiply, + operand0=k_t, + engine=nisa.vector_engine, + ) + + state_new = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=state_new, data1=state, data2=outer_prod, op=nl.add) + nisa.tensor_copy(dst=state, src=state_new) + + # ---- Step 5: o_t = state^T @ q_t ---- + o_t_psum = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=o_t_psum, stationary=state, moving=q_t) + o_t = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=o_t, src=o_t_psum) + + # ---- Store output for token t ---- + nisa.dma_copy( + dst=output.ap(pattern=[[1, dim]], offset=tok_offset), + src=o_t, + ) + + # ---- Write final state to HBM ---- + # state is (128, 128) in SBUF, copy to final_state in HBM + # Use dma_copy with full tile: P_MAX rows, dim cols + nisa.dma_copy(dst=final_state, src=state) + + return output, final_state diff --git a/contrib/models/Qwen3.5-35B-A3B/src/nki_kernels/nki_deltanet_chunked.py b/contrib/models/Qwen3.5-35B-A3B/src/nki_kernels/nki_deltanet_chunked.py new file mode 100644 index 00000000..26b8d0ac --- /dev/null +++ b/contrib/models/Qwen3.5-35B-A3B/src/nki_kernels/nki_deltanet_chunked.py @@ -0,0 +1,431 @@ +"""NKI per-chunk DeltaNet kernel for CTE (context encoding / prefill). + +Single-chunk kernel: processes one chunk (128 tokens) with masked Neumann +power-doubling for intra-chunk correction. The caller loops over chunks in +PyTorch, passing state between calls. + +Each kernel call: + - Takes one chunk of data: q, k, v, beta, g_cumsum, g_last (all 128x128) + - Takes recurrent state_in (128x128) + - Returns chunk output (128x128) and state_out (128x128) + +No sequence-indexed DMA inside the kernel -- all inputs/outputs are full tiles. +This avoids the DMA OOB issue seen with nl.sequential_range + slice indexing +in the NxDI model compilation context. + +NKI v3 (neuronx-cc 2.26.6360 / nki 0.5.0). Uses nki.* namespace. +""" + +import nki +import nki.isa as nisa +import nki.language as nl + +P_MAX = 128 + +# Broadcast partition 0 to all partitions in a 32-wide group. +_BROADCAST_MASK = [0] * 32 + + +@nki.jit +def deltanet_chunk_step( + query, # (128, 128) float32 -- one chunk, l2-normed+scaled + key, # (128, 128) float32 -- one chunk, l2-normed + value, # (128, 128) float32 -- one chunk + beta_broadcast, # (128, 128) float32 -- write gate broadcast to 128 + g_cumsum, # (128, 128) float32 -- cumsum of g within chunk, broadcast + g_last, # (128, 128) float32 -- g_cumsum[-1], constant in chunk, broadcast + state_in, # (128, 128) float32 -- recurrent state from previous chunk + lower_mask, # (128, 128) float32 -- strict lower triangular + identity, # (128, 128) float32 -- identity matrix + lower_mask_diag, # (128, 128) float32 -- lower tri with diagonal +): + """Process one chunk of DeltaNet. + + Returns: + output: (128, 128) float32 -- chunk output + state_out: (128, 128) float32 -- updated recurrent state + """ + C, dim = query.shape # C = 128, dim = 128 + + # Output tensors in HBM + output = nl.ndarray((P_MAX, dim), dtype=query.dtype, buffer=nl.shared_hbm) + state_out = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.shared_hbm) + + # Load all inputs into SBUF + q_c = nl.ndarray((P_MAX, dim), dtype=query.dtype, buffer=nl.sbuf) + nisa.dma_copy(dst=q_c, src=query) + + k_c = nl.ndarray((P_MAX, dim), dtype=key.dtype, buffer=nl.sbuf) + nisa.dma_copy(dst=k_c, src=key) + + v_c = nl.ndarray((P_MAX, dim), dtype=value.dtype, buffer=nl.sbuf) + nisa.dma_copy(dst=v_c, src=value) + + beta_c = nl.ndarray((P_MAX, dim), dtype=beta_broadcast.dtype, buffer=nl.sbuf) + nisa.dma_copy(dst=beta_c, src=beta_broadcast) + + gc_c = nl.ndarray((P_MAX, dim), dtype=g_cumsum.dtype, buffer=nl.sbuf) + nisa.dma_copy(dst=gc_c, src=g_cumsum) + + gl_c = nl.ndarray((P_MAX, dim), dtype=g_last.dtype, buffer=nl.sbuf) + nisa.dma_copy(dst=gl_c, src=g_last) + + state = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=state, src=state_in) + + # Load masks + eye = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=eye, src=identity) + + Lmask = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=Lmask, src=lower_mask) + + Lmask_d = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=Lmask_d, src=lower_mask_diag) + + # ============================================================ + # k_beta = K * beta, v_beta = V * beta + # ============================================================ + k_beta = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=k_beta, data1=k_c, data2=beta_c, op=nl.multiply) + + v_beta = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=v_beta, data1=v_c, data2=beta_c, op=nl.multiply) + + # ============================================================ + # Stable decay factors from cumulative log-decay + # + # The caller passes g_cumsum and g_last broadcast to (128, 128). Extract + # one column and build pairwise decays as exp(gc[i] - gc[j]) so no + # individual exp(-gc[j]) term can overflow. + # ============================================================ + gc_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=gc_p[0:P_MAX, 0:1], src=gc_c[0:P_MAX, 0:1]) + + gl_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=gl_p[0:P_MAX, 0:1], src=gl_c[0:P_MAX, 0:1]) + + exp_gc_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.activation( + dst=exp_gc_p[0:P_MAX, 0:1], + op=nl.exp, + data=gc_p[0:P_MAX, 0:1], + bias=None, + scale=1.0, + ) + + exp_gl_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.activation( + dst=exp_gl_p[0:P_MAX, 0:1], + op=nl.exp, + data=gl_p[0:P_MAX, 0:1], + bias=None, + scale=1.0, + ) + + gc_padded = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.memset(dst=gc_padded, value=0.0) + nisa.tensor_copy(dst=gc_padded[0:P_MAX, 0:1], src=gc_p[0:P_MAX, 0:1]) + + gc_row_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=gc_row_psum, data=gc_padded) + + gc_row = nl.ndarray((1, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=gc_row[0:1, 0:P_MAX], src=gc_row_psum[0:1, 0:P_MAX]) + + gc_row_broadcast = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + for i_shuf in nl.static_range(P_MAX // 32): + nisa.nc_stream_shuffle( + src=gc_row[0:1, 0:P_MAX], + dst=gc_row_broadcast[i_shuf * 32 : i_shuf * 32 + 32, 0:P_MAX], + shuffle_mask=_BROADCAST_MASK, + ) + + gc_col_strict = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=gc_col_strict, + data=Lmask, + op0=nl.multiply, + operand0=gc_p, + engine=nisa.vector_engine, + ) + gc_row_strict = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=gc_row_strict, data1=gc_row_broadcast, data2=Lmask, op=nl.multiply + ) + g_diff_strict = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=g_diff_strict, + data1=gc_col_strict, + data2=gc_row_strict, + op=nl.subtract, + ) + decay_strict_raw = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.activation( + dst=decay_strict_raw, + op=nl.exp, + data=g_diff_strict, + bias=None, + scale=1.0, + ) + decay_strict = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=decay_strict, data1=decay_strict_raw, data2=Lmask, op=nl.multiply + ) + + gc_col_diag = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=gc_col_diag, + data=Lmask_d, + op0=nl.multiply, + operand0=gc_p, + engine=nisa.vector_engine, + ) + gc_row_diag = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=gc_row_diag, data1=gc_row_broadcast, data2=Lmask_d, op=nl.multiply + ) + g_diff_diag = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=g_diff_diag, + data1=gc_col_diag, + data2=gc_row_diag, + op=nl.subtract, + ) + decay_diag_raw = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.activation( + dst=decay_diag_raw, + op=nl.exp, + data=g_diff_diag, + bias=None, + scale=1.0, + ) + decay_diag = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=decay_diag, data1=decay_diag_raw, data2=Lmask_d, op=nl.multiply + ) + + # ============================================================ + # Phase 1: Build A matrix (intra-chunk correction) + # QK = k_beta @ k^T -- contract over features + # ============================================================ + kb_T_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=kb_T_psum, stationary=k_beta, moving=eye) + kb_T = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=kb_T, src=kb_T_psum) + + k_T_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=k_T_psum, stationary=k_c, moving=eye) + k_T = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=k_T, src=k_T_psum) + + QK_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=QK_psum, stationary=kb_T, moving=k_T) + QK = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=QK, src=QK_psum) + + # QK_decay[i,j] = QK[i,j] * exp(gc[i] - gc[j]) for i > j. + QK_decay = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=QK_decay, data1=QK, data2=decay_strict, op=nl.multiply) + + # A = -QK_decay * lower_mask + neg_QK_decay = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=neg_QK_decay, + data=QK_decay, + op0=nl.multiply, + operand0=-1.0, + engine=nisa.vector_engine, + ) + A_mat = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=A_mat, data1=neg_QK_decay, data2=Lmask, op=nl.multiply) + + # ============================================================ + # Masked Neumann power-doubling: + # N = (I + A)(I + A^2)(I + A^4)...(I + A^64) + # + # A_mat is strictly lower triangular, so A^128 = 0. Re-mask after every + # square/multiply so numerical residue cannot leak above the diagonal. + # ============================================================ + P_acc = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=P_acc, data1=eye, data2=A_mat, op=nl.add) + + A_pow = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=A_pow, src=A_mat) + + for _round in nl.sequential_range(6): + Ap_T_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=Ap_T_psum, data=A_pow) + Ap_T = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=Ap_T, src=Ap_T_psum) + + Ap_sq_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=Ap_sq_psum, stationary=Ap_T, moving=A_pow) + nisa.tensor_copy(dst=A_pow, src=Ap_sq_psum) + nisa.tensor_tensor(dst=A_pow, data1=A_pow, data2=Lmask, op=nl.multiply) + + IpA = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=IpA, data1=eye, data2=A_pow, op=nl.add) + + IpA_T_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=IpA_T_psum, data=IpA) + IpA_T = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=IpA_T, src=IpA_T_psum) + + Pacc_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=Pacc_psum, stationary=IpA_T, moving=P_acc) + nisa.tensor_copy(dst=P_acc, src=Pacc_psum) + nisa.tensor_tensor(dst=P_acc, data1=P_acc, data2=Lmask_d, op=nl.multiply) + + # ============================================================ + # Apply N: value_corr = N @ v_beta, k_cumdecay = N @ (k_beta * exp_gc) + # ============================================================ + N_T_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=N_T_psum, data=P_acc) + N_T = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=N_T, src=N_T_psum) + + vc_psum = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=vc_psum, stationary=N_T, moving=v_beta) + value_corr = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=value_corr, src=vc_psum) + + kb_exp_gc = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=kb_exp_gc, + data=k_beta, + op0=nl.multiply, + operand0=exp_gc_p, + engine=nisa.vector_engine, + ) + + kcd_psum = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=kcd_psum, stationary=N_T, moving=kb_exp_gc) + k_cumdecay = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=k_cumdecay, src=kcd_psum) + + # ============================================================ + # Phase 2: Inter-chunk state propagation + # attn_intra = (q @ k^T) * decay_mask * lower_mask_diag + # ============================================================ + q_T_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=q_T_psum, stationary=q_c, moving=eye) + q_T = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=q_T, src=q_T_psum) + + qk_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=qk_psum, stationary=q_T, moving=k_T) + qk_raw = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=qk_raw, src=qk_psum) + + attn_intra = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=attn_intra, data1=qk_raw, data2=decay_diag, op=nl.multiply) + + # ============================================================ + # v_prime = k_cumdecay @ state + # ============================================================ + kcd_T_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=kcd_T_psum, stationary=k_cumdecay, moving=eye) + kcd_T = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=kcd_T, src=kcd_T_psum) + + vp_psum = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=vp_psum, stationary=kcd_T, moving=state) + v_prime = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=v_prime, src=vp_psum) + + v_new = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=v_new, data1=value_corr, data2=v_prime, op=nl.subtract) + + # ============================================================ + # attn_inter = (q * exp(g_cumsum)) @ state + # ============================================================ + q_exp = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=q_exp, + data=q_c, + op0=nl.multiply, + operand0=exp_gc_p, + engine=nisa.vector_engine, + ) + + qe_T_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=qe_T_psum, stationary=q_exp, moving=eye) + qe_T = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=qe_T, src=qe_T_psum) + + ai_psum = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=ai_psum, stationary=qe_T, moving=state) + attn_inter = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=attn_inter, src=ai_psum) + + # ============================================================ + # attn_intra @ v_new + # ============================================================ + ai_T_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=ai_T_psum, stationary=attn_intra, moving=eye) + ai_T = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=ai_T, src=ai_T_psum) + + intra_psum = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=intra_psum, stationary=ai_T, moving=v_new) + intra_out = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=intra_out, src=intra_psum) + + # ============================================================ + # chunk_output = attn_inter + intra_out + # ============================================================ + chunk_out = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=chunk_out, data1=attn_inter, data2=intra_out, op=nl.add) + + nisa.dma_copy(dst=output, src=chunk_out) + + # ============================================================ + # State update: state_new = state * exp(g_last) + # + (k * exp(g_last - gc))^T @ v_new + # ============================================================ + gl_minus_gc_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=gl_minus_gc_p, + data1=gl_p, + data2=gc_p, + op=nl.subtract, + ) + exp_gl_minus_gc_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.activation( + dst=exp_gl_minus_gc_p, + op=nl.exp, + data=gl_minus_gc_p, + bias=None, + scale=1.0, + ) + + k_raw_decay = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=k_raw_decay, + data=k_c, + op0=nl.multiply, + operand0=exp_gl_minus_gc_p, + engine=nisa.vector_engine, + ) + + kv_psum = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=kv_psum, stationary=k_raw_decay, moving=v_new) + kv_outer = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=kv_outer, src=kv_psum) + + state_decayed = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=state_decayed, + data=state, + op0=nl.multiply, + operand0=exp_gl_p, + engine=nisa.vector_engine, + ) + + state_new = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=state_new, data1=state_decayed, data2=kv_outer, op=nl.add) + + nisa.dma_copy(dst=state_out, src=state_new) + + return output, state_out diff --git a/contrib/models/Qwen3.5-35B-A3B/src/nki_kernels/nki_deltanet_fused.py b/contrib/models/Qwen3.5-35B-A3B/src/nki_kernels/nki_deltanet_fused.py new file mode 100644 index 00000000..62e7add7 --- /dev/null +++ b/contrib/models/Qwen3.5-35B-A3B/src/nki_kernels/nki_deltanet_fused.py @@ -0,0 +1,2991 @@ +"""Fused single-kernel DeltaNet chunked forward for CTE (context encoding). + +SSD-style architecture: processes ALL chunks for one (batch, head) pair in +a single NKI kernel call. State (128x128) persists in SBUF across chunks — +no HBM round-trips for inter-chunk state propagation. + +Key optimizations over nki_deltanet_chunked.py: + 1. Single kernel call per (B,H) instead of B*H*num_chunks calls + 2. State in SBUF across all chunks (no HBM state read/write per chunk) + 3. In-kernel cumsum via tensor_tensor_scan (no PyTorch cumsum) + 4. Masks and constants loaded once, reused across chunks + 5. Uses tensor_scalar for partition-broadcast (no explicit broadcast loops) + 6. nc_transpose (Vector Engine) for all 128x128 transposes instead of + nc_matmul(moving=eye) (Tensor Engine) — frees TE for actual math + +neuronx-cc 2.26.6360 / nki 0.5.0. k_dim = v_dim = 128 = P_MAX exactly. +Chunk size = 128 = P_MAX (one tile per chunk). + +Mathematical framework: + Per-chunk direct triangular solve for intra-chunk correction: + QK_decay[i,j] = QK[i,j] * exp(gc[i] - gc[j]) for i > j + A = -QK_decay * lower_mask + v_new = solve((I - A), v_beta - (k_beta * exp(gc)) @ state) + + Inter-chunk state propagation: + attn_inter = (q * exp(gc)) @ state + attn_intra = (q @ k^T) * (strict_decay + I) + output = attn_inter + attn_intra @ v_new + state = exp(g_last) * (state + k_raw_decay^T @ v_new) +""" + +import os + +import numpy as np + +import nki +import nki.isa as nisa +import nki.language as nl + +P_MAX = 128 # Partition dim = k_dim = v_dim +CHUNK_SIZE = int(os.environ.get("QWEN36_DELTANET_CHUNK_SIZE", "128")) +L2_EPS_SQUARED = 1.0e-12 +QUERY_SCALE = P_MAX ** -0.5 +SOLVE_BLOCK_SIZE = int(os.environ.get("QWEN36_DELTANET_SOLVE_BLOCK_SIZE", "32")) +if ( + CHUNK_SIZE <= 0 + or P_MAX % CHUNK_SIZE != 0 + or CHUNK_SIZE % 32 != 0 + or SOLVE_BLOCK_SIZE <= 0 + or CHUNK_SIZE % SOLVE_BLOCK_SIZE != 0 + or SOLVE_BLOCK_SIZE % 32 != 0 + or SOLVE_BLOCK_SIZE & (SOLVE_BLOCK_SIZE - 1) != 0 +): + raise ValueError( + "QWEN36_DELTANET_CHUNK_SIZE must be a positive divisor of P_MAX " + "and a multiple of the 32-partition broadcast group, while " + "QWEN36_DELTANET_SOLVE_BLOCK_SIZE must be positive, divide " + "CHUNK_SIZE, be a power of two, and be a multiple of 32; " + f"P_MAX={P_MAX}, CHUNK_SIZE={CHUNK_SIZE}, got {SOLVE_BLOCK_SIZE}" + ) +MAX_SOLVE_SCAN_STEPS = SOLVE_BLOCK_SIZE.bit_length() - 1 +SOLVE_SCAN_STEPS = int( + os.environ.get("QWEN36_DELTANET_SOLVE_SCAN_STEPS", str(MAX_SOLVE_SCAN_STEPS)) +) +if SOLVE_SCAN_STEPS <= 0 or SOLVE_SCAN_STEPS > MAX_SOLVE_SCAN_STEPS: + raise ValueError( + "QWEN36_DELTANET_SOLVE_SCAN_STEPS must be in " + f"[1, {MAX_SOLVE_SCAN_STEPS}] for SOLVE_BLOCK_SIZE={SOLVE_BLOCK_SIZE}; " + f"got {SOLVE_SCAN_STEPS}" + ) +SOLVE_ACTIVE_PREFIX_K = os.environ.get( + "QWEN36_DELTANET_SOLVE_ACTIVE_PREFIX_K", + "0", +).lower() not in ("0", "false", "no", "off") +SOLVE_MODE = os.environ.get("QWEN36_DELTANET_SOLVE_MODE", "doubling").lower() +AUTOCP_CP_CHUNKS = int(os.environ.get("QWEN36_DELTANET_AUTOCP_CP_CHUNKS", "4")) +if SOLVE_MODE not in ("doubling", "kkt_hier"): + raise ValueError( + "QWEN36_DELTANET_SOLVE_MODE must be one of " + "('doubling', 'kkt_hier'); " + f"got {SOLVE_MODE!r}" + ) +SOLVE_KKT_HIER = SOLVE_MODE == "kkt_hier" +if SOLVE_KKT_HIER and (SOLVE_BLOCK_SIZE != P_MAX or CHUNK_SIZE != P_MAX): + raise ValueError( + "QWEN36_DELTANET_SOLVE_MODE=kkt_hier currently expects " + f"QWEN36_DELTANET_CHUNK_SIZE={P_MAX} and " + f"QWEN36_DELTANET_SOLVE_BLOCK_SIZE={P_MAX}; " + f"got CHUNK_SIZE={CHUNK_SIZE}, SOLVE_BLOCK_SIZE={SOLVE_BLOCK_SIZE}" + ) +if AUTOCP_CP_CHUNKS <= 0: + raise ValueError( + "QWEN36_DELTANET_AUTOCP_CP_CHUNKS must be positive; " + f"got {AUTOCP_CP_CHUNKS}" + ) + +# Broadcast partition 0 to all partitions in a 32-wide group +_BROADCAST_MASK = [0] * 32 + + +def _make_lower_mask(): + """Strict lower triangular active chunk block in a 128x128 constant.""" + mask = np.zeros((P_MAX, P_MAX), dtype=np.float32) + mask[:CHUNK_SIZE, :CHUNK_SIZE] = np.tril( + np.ones((CHUNK_SIZE, CHUNK_SIZE), dtype=np.float32), k=-1 + ) + return mask + + +def _make_lower_mask_diag(): + """Lower triangular active chunk block with diagonal in a 128x128 constant.""" + mask = np.zeros((P_MAX, P_MAX), dtype=np.float32) + mask[:CHUNK_SIZE, :CHUNK_SIZE] = np.tril( + np.ones((CHUNK_SIZE, CHUNK_SIZE), dtype=np.float32), k=0 + ) + return mask + + +def _make_identity(): + """Identity active chunk block in a 128x128 constant.""" + identity = np.zeros((P_MAX, P_MAX), dtype=np.float32) + identity[:CHUNK_SIZE, :CHUNK_SIZE] = np.eye(CHUNK_SIZE, dtype=np.float32) + return identity + + +def _matmul_square(dst, left, right, size): + left_trans_psum = nl.ndarray((size, size), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=left_trans_psum, data=left) + left_trans = nl.ndarray((size, size), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=left_trans, src=left_trans_psum) + + out_psum = nl.ndarray((size, size), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=out_psum, stationary=left_trans, moving=right) + nisa.tensor_copy(dst=dst, src=out_psum) + + +def _offdiag_combine_t(dst, left_t, cross_t, right_t, size): + tmp = nl.ndarray((size, size), dtype=nl.float32, buffer=nl.sbuf) + _matmul_square(tmp, left_t, cross_t, size) + _matmul_square(dst, tmp, right_t, size) + + +def _leaf_inverse32_t(dst, A_T, Imat, start): + nisa.tensor_copy(dst=dst, src=Imat[0:32, 0:32]) + + power_t = nl.ndarray((32, 32), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=power_t, src=A_T[start : start + 32, start : start + 32]) + + power_psum = nl.ndarray((32, 32), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=power_psum, data=power_t) + power = nl.ndarray((32, 32), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=power, src=power_psum) + + for _scan_i in nl.static_range(5): + correction = nl.ndarray((32, 32), dtype=nl.float32, buffer=nl.sbuf) + _matmul_square(correction, dst, power_t, 32) + + next_inv_t = nl.ndarray((32, 32), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=next_inv_t, data1=dst, data2=correction, op=nl.add) + nisa.tensor_copy(dst=dst, src=next_inv_t) + + if _scan_i != 4: + power_next = nl.ndarray((32, 32), dtype=nl.float32, buffer=nl.sbuf) + _matmul_square(power_next, power, power, 32) + + power_next_t_psum = nl.ndarray( + (32, 32), dtype=nl.float32, buffer=nl.psum + ) + nisa.nc_transpose(dst=power_next_t_psum, data=power_next) + power_next_t = nl.ndarray((32, 32), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=power_next_t, src=power_next_t_psum) + + nisa.tensor_copy(dst=power, src=power_next) + nisa.tensor_copy(dst=power_t, src=power_next_t) + + +def _inverse64_t(dst, A_T, Imat, start): + nisa.memset(dst=dst, value=0.0) + + for leaf_idx in nl.static_range(2): + leaf_offset = leaf_idx * 32 + leaf_start = start + leaf_offset + leaf_t = nl.ndarray((32, 32), dtype=nl.float32, buffer=nl.sbuf) + _leaf_inverse32_t(leaf_t, A_T, Imat, leaf_start) + nisa.tensor_copy( + dst=dst[leaf_offset : leaf_offset + 32, leaf_offset : leaf_offset + 32], + src=leaf_t, + ) + + left32_t = nl.ndarray((32, 32), dtype=nl.float32, buffer=nl.sbuf) + right32_t = nl.ndarray((32, 32), dtype=nl.float32, buffer=nl.sbuf) + cross32_t = nl.ndarray((32, 32), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=left32_t, src=dst[0:32, 0:32]) + nisa.tensor_copy(dst=right32_t, src=dst[32:64, 32:64]) + nisa.tensor_copy(dst=cross32_t, src=A_T[start : start + 32, start + 32 : start + 64]) + + off32_t = nl.ndarray((32, 32), dtype=nl.float32, buffer=nl.sbuf) + _offdiag_combine_t(off32_t, left32_t, cross32_t, right32_t, 32) + nisa.tensor_copy(dst=dst[0:32, 32:64], src=off32_t) + + +def _hierarchical_kkt_solve128(v_new, A_T, Imat, solve_rhs, dim): + n_lo_t = nl.ndarray((64, 64), dtype=nl.float32, buffer=nl.sbuf) + n_hi_t = nl.ndarray((64, 64), dtype=nl.float32, buffer=nl.sbuf) + _inverse64_t(n_lo_t, A_T, Imat, 0) + _inverse64_t(n_hi_t, A_T, Imat, 64) + + n128_t = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.memset(dst=n128_t, value=0.0) + nisa.tensor_copy(dst=n128_t[0:64, 0:64], src=n_lo_t) + nisa.tensor_copy(dst=n128_t[64:128, 64:128], src=n_hi_t) + + cross64_t = nl.ndarray((64, 64), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=cross64_t, src=A_T[0:64, 64:128]) + + off64_t = nl.ndarray((64, 64), dtype=nl.float32, buffer=nl.sbuf) + _offdiag_combine_t(off64_t, n_lo_t, cross64_t, n_hi_t, 64) + nisa.tensor_copy(dst=n128_t[0:64, 64:128], src=off64_t) + + solved_psum = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=solved_psum, stationary=n128_t, moving=solve_rhs) + nisa.tensor_copy(dst=v_new, src=solved_psum) + + +def _blocked_doubling_solve(v_new, A_T, solve_rhs, dim): + for solve_block in nl.static_range(CHUNK_SIZE // SOLVE_BLOCK_SIZE): + block_start = solve_block * SOLVE_BLOCK_SIZE + block_end = block_start + SOLVE_BLOCK_SIZE + + prev_contrib = nl.ndarray( + (SOLVE_BLOCK_SIZE, dim), dtype=nl.float32, buffer=nl.sbuf + ) + if solve_block == 0: + nisa.memset(dst=prev_contrib, value=0.0) + else: + prev_psum = nl.ndarray( + (SOLVE_BLOCK_SIZE, dim), dtype=nl.float32, buffer=nl.psum + ) + if SOLVE_ACTIVE_PREFIX_K: + nisa.nc_matmul( + dst=prev_psum, + stationary=A_T[0:block_start, block_start:block_end], + moving=v_new[0:block_start, 0:dim], + ) + else: + nisa.nc_matmul( + dst=prev_psum, + stationary=A_T[0:CHUNK_SIZE, block_start:block_end], + moving=v_new[0:CHUNK_SIZE, 0:dim], + ) + nisa.tensor_copy(dst=prev_contrib, src=prev_psum) + + solve_rhs_block = nl.ndarray( + (SOLVE_BLOCK_SIZE, dim), dtype=nl.float32, buffer=nl.sbuf + ) + nisa.tensor_copy( + dst=solve_rhs_block, + src=solve_rhs[block_start:block_end, 0:dim], + ) + + residual_block = nl.ndarray( + (SOLVE_BLOCK_SIZE, dim), dtype=nl.float32, buffer=nl.sbuf + ) + nisa.tensor_tensor( + dst=residual_block, + data1=solve_rhs_block, + data2=prev_contrib, + op=nl.add, + ) + + A_diag_T = nl.ndarray( + (SOLVE_BLOCK_SIZE, SOLVE_BLOCK_SIZE), + dtype=nl.float32, + buffer=nl.sbuf, + ) + nisa.tensor_copy( + dst=A_diag_T, + src=A_T[block_start:block_end, block_start:block_end], + ) + + A_power_T = nl.ndarray( + (SOLVE_BLOCK_SIZE, SOLVE_BLOCK_SIZE), + dtype=nl.float32, + buffer=nl.sbuf, + ) + nisa.tensor_copy(dst=A_power_T, src=A_diag_T) + + A_power_psum = nl.ndarray( + (SOLVE_BLOCK_SIZE, SOLVE_BLOCK_SIZE), + dtype=nl.float32, + buffer=nl.psum, + ) + nisa.nc_transpose(dst=A_power_psum, data=A_power_T) + A_power = nl.ndarray( + (SOLVE_BLOCK_SIZE, SOLVE_BLOCK_SIZE), + dtype=nl.float32, + buffer=nl.sbuf, + ) + nisa.tensor_copy(dst=A_power, src=A_power_psum) + + local_v = nl.ndarray( + (SOLVE_BLOCK_SIZE, dim), dtype=nl.float32, buffer=nl.sbuf + ) + nisa.tensor_copy(dst=local_v, src=residual_block) + + for _scan_i in nl.static_range(SOLVE_SCAN_STEPS): + correction_psum = nl.ndarray( + (SOLVE_BLOCK_SIZE, dim), dtype=nl.float32, buffer=nl.psum + ) + nisa.nc_matmul( + dst=correction_psum, + stationary=A_power_T, + moving=local_v, + ) + correction = nl.ndarray( + (SOLVE_BLOCK_SIZE, dim), dtype=nl.float32, buffer=nl.sbuf + ) + nisa.tensor_copy(dst=correction, src=correction_psum) + + local_next = nl.ndarray( + (SOLVE_BLOCK_SIZE, dim), dtype=nl.float32, buffer=nl.sbuf + ) + nisa.tensor_tensor( + dst=local_next, data1=local_v, data2=correction, op=nl.add + ) + + nisa.tensor_copy(dst=local_v, src=local_next) + + if _scan_i == SOLVE_SCAN_STEPS - 2: + A_power_next_T_psum = nl.ndarray( + (SOLVE_BLOCK_SIZE, SOLVE_BLOCK_SIZE), + dtype=nl.float32, + buffer=nl.psum, + ) + nisa.nc_matmul( + dst=A_power_next_T_psum, + stationary=A_power, + moving=A_power_T, + ) + nisa.tensor_copy(dst=A_power_T, src=A_power_next_T_psum) + elif _scan_i != SOLVE_SCAN_STEPS - 1: + A_power_next_psum = nl.ndarray( + (SOLVE_BLOCK_SIZE, SOLVE_BLOCK_SIZE), + dtype=nl.float32, + buffer=nl.psum, + ) + nisa.nc_matmul( + dst=A_power_next_psum, + stationary=A_power_T, + moving=A_power, + ) + A_power_next = nl.ndarray( + (SOLVE_BLOCK_SIZE, SOLVE_BLOCK_SIZE), + dtype=nl.float32, + buffer=nl.sbuf, + ) + nisa.tensor_copy(dst=A_power_next, src=A_power_next_psum) + + A_power_next_T_psum = nl.ndarray( + (SOLVE_BLOCK_SIZE, SOLVE_BLOCK_SIZE), + dtype=nl.float32, + buffer=nl.psum, + ) + nisa.nc_transpose(dst=A_power_next_T_psum, data=A_power_next) + A_power_next_T = nl.ndarray( + (SOLVE_BLOCK_SIZE, SOLVE_BLOCK_SIZE), + dtype=nl.float32, + buffer=nl.sbuf, + ) + nisa.tensor_copy(dst=A_power_next_T, src=A_power_next_T_psum) + + nisa.tensor_copy(dst=A_power, src=A_power_next) + nisa.tensor_copy(dst=A_power_T, src=A_power_next_T) + + nisa.tensor_copy( + dst=v_new[block_start:block_end, 0:dim], + src=local_v[0:SOLVE_BLOCK_SIZE, 0:dim], + ) + + +@nki.jit +def deltanet_fused_chunked_fwd( + query: nl.ndarray, # (S, 128) float32 — raw Q; normalized in-kernel + key: nl.ndarray, # (S, 128) float32 — raw K; normalized in-kernel + value: nl.ndarray, # (S, 128) float32 + g_in: nl.ndarray, # (S, 1) float32 — per-token log-decay (NOT cumsum) + beta_in: nl.ndarray, # (S, 1) float32 — per-token write gate + initial_state: nl.ndarray, # (128, 128) float32 — recurrent checkpoint or zeros + lower_mask: nl.ndarray, # (128, 128) float32 — strict lower tri + identity: nl.ndarray, # (128, 128) float32 — identity + lower_mask_diag: nl.ndarray, # (128, 128) float32 — lower tri with diag +): + """Fused chunked DeltaNet forward — single kernel call per (batch, head). + + Processes all chunks sequentially within the kernel, keeping the recurrent + state (128x128) in SBUF across chunks. Returns per-token output and + final state. + + Input requirements: + - S must be divisible by 128 (pad before calling) + - query/key are raw projected chunks; l2-norm and Q scale are in-kernel + - g_in is RAW log-decay (cumsum computed in-kernel via tensor_tensor_scan) + - beta_in is sigmoid(b) (write gate) + - initial_state is zero for cold prefill, or the restored GDN checkpoint + + Returns: + output: (S, 128) float32 + final_state: (128, 128) float32 + """ + seq_len = query.shape[0] + dim = query.shape[1] # 128 + num_chunks = seq_len // CHUNK_SIZE + + # Output tensors in HBM + output = nl.ndarray((seq_len, dim), dtype=query.dtype, buffer=nl.shared_hbm) + final_state_out = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.shared_hbm) + + # ================================================================ + # Load constant masks into SBUF once (reused across all chunks) + # ================================================================ + Lmask = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=Lmask, src=lower_mask) + + UMask_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=UMask_psum, data=Lmask) + UMask = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=UMask, src=UMask_psum) + + Imat = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=Imat, src=identity) + + # Ones vector for cumsum scan: (1, CHUNK_SIZE) + ones_1xC = nl.ndarray((1, CHUNK_SIZE), dtype=nl.float32, buffer=nl.sbuf) + nisa.memset(dst=ones_1xC, value=1.0) + + # Zero initial for cumsum scan + zero_11 = nl.ndarray((1, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.memset(dst=zero_11, value=0.0) + + # ================================================================ + # Initialize recurrent state in SBUF — persists across ALL chunks + # ================================================================ + state = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=state, src=initial_state) + + # ================================================================ + # Sequential chunk processing + # ================================================================ + for i_chunk in nl.sequential_range(num_chunks): + chunk_start = i_chunk * CHUNK_SIZE + + # ---- Load chunk data from HBM ---- + q_c = nl.ndarray((P_MAX, dim), dtype=query.dtype, buffer=nl.sbuf) + if CHUNK_SIZE == P_MAX: + nisa.dma_copy( + dst=q_c, + src=query[chunk_start : chunk_start + CHUNK_SIZE, 0:dim], + ) + else: + nisa.memset(dst=q_c, value=0.0) + nisa.dma_copy( + dst=q_c[0:CHUNK_SIZE, 0:dim], + src=query[chunk_start : chunk_start + CHUNK_SIZE, 0:dim], + ) + + k_c = nl.ndarray((P_MAX, dim), dtype=key.dtype, buffer=nl.sbuf) + if CHUNK_SIZE == P_MAX: + nisa.dma_copy( + dst=k_c, + src=key[chunk_start : chunk_start + CHUNK_SIZE, 0:dim], + ) + else: + nisa.memset(dst=k_c, value=0.0) + nisa.dma_copy( + dst=k_c[0:CHUNK_SIZE, 0:dim], + src=key[chunk_start : chunk_start + CHUNK_SIZE, 0:dim], + ) + + q_square = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=q_square, data1=q_c, data2=q_c, op=nl.multiply) + q_norm_sq = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_reduce(dst=q_norm_sq, data=q_square, op=nl.add, axis=1) + q_norm_sq_clamped = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=q_norm_sq_clamped, + data=q_norm_sq, + op0=nl.maximum, + operand0=L2_EPS_SQUARED, + engine=nisa.vector_engine, + ) + q_inv_norm = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=q_inv_norm, + data=q_norm_sq_clamped, + op0=nl.rsqrt, + operand0=0.0, + engine=nisa.gpsimd_engine, + ) + q_norm = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=q_norm, + data=q_c, + op0=nl.multiply, + operand0=q_inv_norm, + op1=nl.multiply, + operand1=QUERY_SCALE, + engine=nisa.vector_engine, + ) + + k_square = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=k_square, data1=k_c, data2=k_c, op=nl.multiply) + k_norm_sq = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_reduce(dst=k_norm_sq, data=k_square, op=nl.add, axis=1) + k_norm_sq_clamped = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=k_norm_sq_clamped, + data=k_norm_sq, + op0=nl.maximum, + operand0=L2_EPS_SQUARED, + engine=nisa.vector_engine, + ) + k_inv_norm = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=k_inv_norm, + data=k_norm_sq_clamped, + op0=nl.rsqrt, + operand0=0.0, + engine=nisa.gpsimd_engine, + ) + k_norm = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=k_norm, + data=k_c, + op0=nl.multiply, + operand0=k_inv_norm, + engine=nisa.vector_engine, + ) + + v_c = nl.ndarray((P_MAX, dim), dtype=value.dtype, buffer=nl.sbuf) + if CHUNK_SIZE == P_MAX: + nisa.dma_copy( + dst=v_c, + src=value[chunk_start : chunk_start + CHUNK_SIZE, 0:dim], + ) + else: + nisa.memset(dst=v_c, value=0.0) + nisa.dma_copy( + dst=v_c[0:CHUNK_SIZE, 0:dim], + src=value[chunk_start : chunk_start + CHUNK_SIZE, 0:dim], + ) + + # g: (CHUNK_SIZE, 1) — raw log-decay per token + g_chunk_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + if CHUNK_SIZE != P_MAX: + nisa.memset(dst=g_chunk_p, value=0.0) + nisa.dma_copy( + dst=g_chunk_p[0:CHUNK_SIZE, 0:1], + src=g_in[chunk_start : chunk_start + CHUNK_SIZE, 0:1], + ) + + # beta: (CHUNK_SIZE, 1) — write gate scalar per token + beta_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + if CHUNK_SIZE != P_MAX: + nisa.memset(dst=beta_p, value=0.0) + nisa.dma_copy( + dst=beta_p[0:CHUNK_SIZE, 0:1], + src=beta_in[chunk_start : chunk_start + CHUNK_SIZE, 0:1], + ) + + # ---- In-kernel cumsum of g via tensor_tensor_scan ---- + # Need g as (1, CHUNK_SIZE) for scan along free dim. Use direct + # vector transpose instead of padding through a full 128x128 tile. + g_tp_psum = nl.ndarray((1, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=g_tp_psum, data=g_chunk_p) + + g_row = nl.ndarray((1, CHUNK_SIZE), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy( + dst=g_row[0:1, 0:CHUNK_SIZE], + src=g_tp_psum[0:1, 0:CHUNK_SIZE], + ) + + # cumsum: gc_row[t] = 1.0 * gc_row[t-1] + g_row[t] + gc_row = nl.ndarray((1, CHUNK_SIZE), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor_scan( + dst=gc_row[0:1, 0:CHUNK_SIZE], + data0=ones_1xC[0:1, 0:CHUNK_SIZE], + data1=g_row[0:1, 0:CHUNK_SIZE], + initial=zero_11[0:1, 0:1], + op0=nl.multiply, + op1=nl.add, + ) + + # Transpose gc back to (CHUNK_SIZE, 1) partition layout. + gc_tp_psum = nl.ndarray((CHUNK_SIZE, 1), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=gc_tp_psum, data=gc_row) + + # gc_p: (P_MAX, 1) — cumulative sum of g per token in this chunk + gc_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + if CHUNK_SIZE != P_MAX: + nisa.memset(dst=gc_p, value=0.0) + nisa.tensor_copy( + dst=gc_p[0:CHUNK_SIZE, 0:1], + src=gc_tp_psum[0:CHUNK_SIZE, 0:1], + ) + + # g_last = gc[-1] (scalar) — needed for state decay + gl_11 = nl.ndarray((1, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy( + dst=gl_11[0:1, 0:1], + src=gc_row[0:1, CHUNK_SIZE - 1 : CHUNK_SIZE], + ) + + # ---- Compute exp(gc) and exp(g_last) as (P_MAX, 1) scalars ---- + # These (P_MAX, 1) tensors are used with tensor_scalar to broadcast + # across the free dimension without explicit (P_MAX, dim) copies. + + exp_gc_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.activation( + dst=exp_gc_p[0:P_MAX, 0:1], + op=nl.exp, + data=gc_p[0:P_MAX, 0:1], + bias=None, + scale=1.0, + ) + + exp_gl_11 = nl.ndarray((1, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.activation( + dst=exp_gl_11, + op=nl.exp, + data=gl_11, + bias=None, + scale=1.0, + ) + + gl_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + for i_shuf in nl.static_range(P_MAX // 32): + nisa.nc_stream_shuffle( + src=gl_11[0:1, 0:1], + dst=gl_p[i_shuf * 32 : i_shuf * 32 + 32, 0:1], + shuffle_mask=_BROADCAST_MASK, + ) + + exp_gl_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + for i_shuf in nl.static_range(P_MAX // 32): + nisa.nc_stream_shuffle( + src=exp_gl_11[0:1, 0:1], + dst=exp_gl_p[i_shuf * 32 : i_shuf * 32 + 32, 0:1], + shuffle_mask=_BROADCAST_MASK, + ) + + # ============================================================ + # Stable pairwise decay factors from cumulative log-decay. + # + # The original fused path used split scaling: + # exp(gc[i]) * exp(-gc[j]) + # That can materialize huge unused intermediates. Build the same + # causal decay matrices as the per-chunk kernel using exp(gc[i]-gc[j]) + # and mask after the exp so upper-triangular values cannot leak into + # later matmuls. + # ============================================================ + gc_row_broadcast = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + if CHUNK_SIZE != P_MAX: + nisa.memset(dst=gc_row_broadcast, value=0.0) + for i_shuf in nl.static_range(P_MAX // 32): + nisa.nc_stream_shuffle( + src=gc_row[0:1, 0:CHUNK_SIZE], + dst=gc_row_broadcast[ + i_shuf * 32 : i_shuf * 32 + 32, 0:CHUNK_SIZE + ], + shuffle_mask=_BROADCAST_MASK, + ) + + gc_col_strict_t = nl.ndarray( + (P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf + ) + nisa.tensor_tensor( + dst=gc_col_strict_t, + data1=gc_row_broadcast, + data2=UMask, + op=nl.multiply, + ) + gc_row_strict_t = nl.ndarray( + (P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf + ) + nisa.tensor_scalar( + dst=gc_row_strict_t, + data=UMask, + op0=nl.multiply, + operand0=gc_p, + engine=nisa.vector_engine, + ) + g_diff_strict_t = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=g_diff_strict_t, + data1=gc_col_strict_t, + data2=gc_row_strict_t, + op=nl.subtract, + ) + decay_strict_t_raw = nl.ndarray( + (P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf + ) + nisa.activation( + dst=decay_strict_t_raw, + op=nl.exp, + data=g_diff_strict_t, + bias=None, + scale=1.0, + ) + decay_strict_t = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=decay_strict_t, + data1=decay_strict_t_raw, + data2=UMask, + op=nl.multiply, + ) + + decay_diag_t = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=decay_diag_t, data1=decay_strict_t, data2=Imat, op=nl.add + ) + + # ============================================================ + # k_beta = K * beta, v_beta = V * beta + # tensor_scalar broadcasts beta_p (P_MAX, 1) across free dim + # ============================================================ + k_beta = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=k_beta, + data=k_norm, + op0=nl.multiply, + operand0=beta_p, + engine=nisa.vector_engine, + ) + + v_beta = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=v_beta, + data=v_c, + op0=nl.multiply, + operand0=beta_p, + engine=nisa.vector_engine, + ) + + # ============================================================ + # Phase 1: Build A matrix (intra-chunk correction) + # Transpose K and K_beta for matmul + # ============================================================ + kb_T_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=kb_T_psum, data=k_beta) + kb_T = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=kb_T, src=kb_T_psum) + + k_T_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=k_T_psum, data=k_norm) + k_T = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=k_T, src=k_T_psum) + + # QK_T[j, i] = k_norm[j] @ k_beta[i]. Build the transposed solve + # matrix directly, avoiding a full A -> A_T transpose per chunk. + QK_T_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=QK_T_psum, stationary=k_T, moving=kb_T) + QK_T = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=QK_T, src=QK_T_psum) + + # A_T[j, i] = -QK[i, j] * exp(gc[i] - gc[j]) for i > j. + QK_decay_t = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=QK_decay_t, data1=QK_T, data2=decay_strict_t, op=nl.multiply + ) + + A_T = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=A_T, + data=QK_decay_t, + op0=nl.multiply, + operand0=-1.0, + engine=nisa.vector_engine, + ) + # ============================================================ + # Build the single RHS needed for v_new. + # + # Materializing N = inv(I - A) would compute: + # value_corr = N @ v_beta + # k_cumdecay = N @ (k_beta * exp(gc)) + # v_new = value_corr - k_cumdecay @ state + # + # By associativity: + # v_new = N @ (v_beta - (k_beta * exp(gc)) @ state) + # + # Solve this RHS directly. This is equivalent to the nilpotent + # Neumann series, but avoids repeated matrix squaring, which is + # numerically unstable for realistic Qwen decay gates. + # ============================================================ + kb_exp_gc = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=kb_exp_gc, + data=k_beta, + op0=nl.multiply, + operand0=exp_gc_p, + engine=nisa.vector_engine, + ) + + kbe_T_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=kbe_T_psum, data=kb_exp_gc) + kbe_T = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=kbe_T, src=kbe_T_psum) + + kbe_state_psum = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=kbe_state_psum, stationary=kbe_T, moving=state) + kbe_state = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=kbe_state, src=kbe_state_psum) + + solve_rhs = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=solve_rhs, data1=v_beta, data2=kbe_state, op=nl.subtract) + + # ============================================================ + # Blocked forward substitution for: + # v_new = solve((I - A), solve_rhs) + # + # A is strictly lower triangular. Compute each solve block's + # contribution from previously solved rows with one dense matmul, then + # solve the small diagonal block row-by-row. This keeps the algebra + # exact while moving the wide part of the triangular solve onto TE + # tiles, closer to the FlashQLA/FLA blocked chunked-prefill structure. + # ============================================================ + v_new = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.memset(dst=v_new, value=0.0) + + if SOLVE_KKT_HIER: + _hierarchical_kkt_solve128(v_new, A_T, Imat, solve_rhs, dim) + else: + _blocked_doubling_solve(v_new, A_T, solve_rhs, dim) + + # ============================================================ + # Phase 2: Inter-chunk state propagation + # attn_intra = (q @ k^T) * (strict_decay + identity) + # ============================================================ + q_T_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=q_T_psum, data=q_norm) + q_T = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=q_T, src=q_T_psum) + + # ai_T[j, i] = (q[i] @ k[j]) * transpose(decay_diag)[j, i]. + qk_T_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=qk_T_psum, stationary=k_T, moving=q_T) + qk_raw_t = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=qk_raw_t, src=qk_T_psum) + + ai_T = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=ai_T, data1=qk_raw_t, data2=decay_diag_t, op=nl.multiply + ) + + # ============================================================ + # attn_inter = (q * exp(gc)) @ state (state is in SBUF!) + # ============================================================ + q_exp = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=q_exp, + data=q_norm, + op0=nl.multiply, + operand0=exp_gc_p, + engine=nisa.vector_engine, + ) + + qe_T_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=qe_T_psum, data=q_exp) + qe_T = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=qe_T, src=qe_T_psum) + + ai_psum = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=ai_psum, stationary=qe_T, moving=state) + attn_inter = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=attn_inter, src=ai_psum) + + # ============================================================ + # attn_intra @ v_new + # ============================================================ + intra_psum = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=intra_psum, stationary=ai_T, moving=v_new) + intra_out = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=intra_out, src=intra_psum) + + # ============================================================ + # chunk_output = attn_inter + intra_out + # ============================================================ + chunk_out = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=chunk_out, data1=attn_inter, data2=intra_out, op=nl.add) + + # Store output chunk to HBM + if CHUNK_SIZE == P_MAX: + nisa.dma_copy( + dst=output[chunk_start : chunk_start + CHUNK_SIZE, 0:dim], + src=chunk_out, + ) + else: + nisa.dma_copy( + dst=output[chunk_start : chunk_start + CHUNK_SIZE, 0:dim], + src=chunk_out[0:CHUNK_SIZE, 0:dim], + ) + + # ============================================================ + # State update: state = exp(g_last) * (state + k_raw_decay^T @ v_new) + # state is updated IN-PLACE in SBUF — no HBM round-trip! + # ============================================================ + + # k_raw_decay contributes as exp(g_last) * (k * exp(-gc))^T @ v_new. + # Compute the equivalent stable form k * exp(g_last - gc) directly so + # no exp(-gc) intermediate can overflow. + exp_gl_minus_gc_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.activation( + dst=exp_gl_minus_gc_p[0:P_MAX, 0:1], + op=nl.exp, + data=gc_p[0:P_MAX, 0:1], + bias=gl_p[0:P_MAX, 0:1], + scale=-1.0, + ) + + # k_raw_decay = k * exp(g_last - gc) + k_raw_decay = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=k_raw_decay, + data=k_norm, + op0=nl.multiply, + operand0=exp_gl_minus_gc_p, + engine=nisa.vector_engine, + ) + + # k_raw_decay^T @ v_new → (dim, dim) outer product sum + # nc_matmul: result[M,N] = sum_K stationary[K,M] * moving[K,N] + # stationary=k_raw_decay (P_MAX, dim), moving=v_new (P_MAX, dim) + # Result: sum over tokens -> (dim, dim) + kv_psum = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=kv_psum, stationary=k_raw_decay, moving=v_new) + kv_outer = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=kv_outer, src=kv_psum) + + # state = state * exp(g_last) + kv_outer + state_decayed = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=state_decayed, + data=state, + op0=nl.multiply, + operand0=exp_gl_p, + engine=nisa.vector_engine, + ) + nisa.tensor_tensor(dst=state, data1=state_decayed, data2=kv_outer, op=nl.add) + + # ---- Write final state to HBM ---- + nisa.dma_copy(dst=final_state_out, src=state) + + return output, final_state_out + + +@nki.jit +def deltanet_autocp_affine_chunk( + query: nl.ndarray, # (128, 128) float32 - raw Q; normalized in-kernel + key: nl.ndarray, # (128, 128) float32 - raw K; normalized in-kernel + value: nl.ndarray, # (128, 128) float32 + g_in: nl.ndarray, # (128, 1) float32 - per-token log-decay + beta_in: nl.ndarray, # (128, 1) float32 - per-token write gate + lower_mask: nl.ndarray, # (128, 128) float32 - strict lower tri + identity: nl.ndarray, # (128, 128) float32 - identity + lower_mask_diag: nl.ndarray, # (128, 128) float32 - lower tri with diag +): + """Build one chunk's state-independent AutoCP affine pieces. + + For one 128-token DeltaNet chunk: + output = output_base + output_state @ state + next_state = state_matrix @ state + state_bias + + This probe deliberately mirrors the fused CTE chunk math and returns the + four intermediate tensors to HBM for isolated correctness validation before + wiring an AutoCP prepass into serving. + """ + dim = query.shape[1] + + output_base_out = nl.ndarray( + (P_MAX, dim), dtype=nl.float32, buffer=nl.shared_hbm + ) + output_state_out = nl.ndarray( + (P_MAX, dim), dtype=nl.float32, buffer=nl.shared_hbm + ) + state_matrix_out = nl.ndarray( + (P_MAX, dim), dtype=nl.float32, buffer=nl.shared_hbm + ) + state_bias_out = nl.ndarray( + (P_MAX, dim), dtype=nl.float32, buffer=nl.shared_hbm + ) + + Lmask = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=Lmask, src=lower_mask) + + Imat = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=Imat, src=identity) + + Lmask_diag = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=Lmask_diag, src=lower_mask_diag) + + ones_1xC = nl.ndarray((1, CHUNK_SIZE), dtype=nl.float32, buffer=nl.sbuf) + nisa.memset(dst=ones_1xC, value=1.0) + + zero_11 = nl.ndarray((1, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.memset(dst=zero_11, value=0.0) + + q_c = nl.ndarray((P_MAX, dim), dtype=query.dtype, buffer=nl.sbuf) + nisa.dma_copy(dst=q_c, src=query[0:CHUNK_SIZE, 0:dim]) + + k_c = nl.ndarray((P_MAX, dim), dtype=key.dtype, buffer=nl.sbuf) + nisa.dma_copy(dst=k_c, src=key[0:CHUNK_SIZE, 0:dim]) + + q_square = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=q_square, data1=q_c, data2=q_c, op=nl.multiply) + q_norm_sq = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_reduce(dst=q_norm_sq, data=q_square, op=nl.add, axis=1) + q_norm_sq_clamped = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=q_norm_sq_clamped, + data=q_norm_sq, + op0=nl.maximum, + operand0=L2_EPS_SQUARED, + engine=nisa.vector_engine, + ) + q_inv_norm = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=q_inv_norm, + data=q_norm_sq_clamped, + op0=nl.rsqrt, + operand0=0.0, + engine=nisa.gpsimd_engine, + ) + q_norm = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=q_norm, + data=q_c, + op0=nl.multiply, + operand0=q_inv_norm, + op1=nl.multiply, + operand1=QUERY_SCALE, + engine=nisa.vector_engine, + ) + + k_square = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=k_square, data1=k_c, data2=k_c, op=nl.multiply) + k_norm_sq = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_reduce(dst=k_norm_sq, data=k_square, op=nl.add, axis=1) + k_norm_sq_clamped = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=k_norm_sq_clamped, + data=k_norm_sq, + op0=nl.maximum, + operand0=L2_EPS_SQUARED, + engine=nisa.vector_engine, + ) + k_inv_norm = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=k_inv_norm, + data=k_norm_sq_clamped, + op0=nl.rsqrt, + operand0=0.0, + engine=nisa.gpsimd_engine, + ) + k_norm = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=k_norm, + data=k_c, + op0=nl.multiply, + operand0=k_inv_norm, + engine=nisa.vector_engine, + ) + + v_c = nl.ndarray((P_MAX, dim), dtype=value.dtype, buffer=nl.sbuf) + nisa.dma_copy(dst=v_c, src=value[0:CHUNK_SIZE, 0:dim]) + + g_chunk_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=g_chunk_p[0:CHUNK_SIZE, 0:1], src=g_in[0:CHUNK_SIZE, 0:1]) + + beta_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=beta_p[0:CHUNK_SIZE, 0:1], src=beta_in[0:CHUNK_SIZE, 0:1]) + + g_padded = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.memset(dst=g_padded, value=0.0) + nisa.tensor_copy(dst=g_padded[0:CHUNK_SIZE, 0:1], src=g_chunk_p) + + g_tp_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=g_tp_psum, data=g_padded) + + g_row = nl.ndarray((1, CHUNK_SIZE), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=g_row[0:1, 0:CHUNK_SIZE], src=g_tp_psum[0:1, 0:CHUNK_SIZE]) + + gc_row = nl.ndarray((1, CHUNK_SIZE), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor_scan( + dst=gc_row[0:1, 0:CHUNK_SIZE], + data0=ones_1xC[0:1, 0:CHUNK_SIZE], + data1=g_row[0:1, 0:CHUNK_SIZE], + initial=zero_11[0:1, 0:1], + op0=nl.multiply, + op1=nl.add, + ) + + gc_padded = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.memset(dst=gc_padded, value=0.0) + nisa.tensor_copy(dst=gc_padded[0:1, 0:CHUNK_SIZE], src=gc_row) + + gc_tp_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=gc_tp_psum, data=gc_padded) + + gc_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=gc_p[0:CHUNK_SIZE, 0:1], src=gc_tp_psum[0:CHUNK_SIZE, 0:1]) + + gl_11 = nl.ndarray((1, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy( + dst=gl_11[0:1, 0:1], + src=gc_row[0:1, CHUNK_SIZE - 1 : CHUNK_SIZE], + ) + + exp_gc_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.activation( + dst=exp_gc_p[0:P_MAX, 0:1], + op=nl.exp, + data=gc_p[0:P_MAX, 0:1], + bias=None, + scale=1.0, + ) + + gl_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + for i_shuf in nl.static_range(P_MAX // 32): + nisa.nc_stream_shuffle( + src=gl_11[0:1, 0:1], + dst=gl_p[i_shuf * 32 : i_shuf * 32 + 32, 0:1], + shuffle_mask=_BROADCAST_MASK, + ) + + exp_gl_11 = nl.ndarray((1, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.activation( + dst=exp_gl_11, + op=nl.exp, + data=gl_11, + bias=None, + scale=1.0, + ) + + exp_gl_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + for i_shuf in nl.static_range(P_MAX // 32): + nisa.nc_stream_shuffle( + src=exp_gl_11[0:1, 0:1], + dst=exp_gl_p[i_shuf * 32 : i_shuf * 32 + 32, 0:1], + shuffle_mask=_BROADCAST_MASK, + ) + + gc_row_broadcast = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + for i_shuf in nl.static_range(P_MAX // 32): + nisa.nc_stream_shuffle( + src=gc_row[0:1, 0:P_MAX], + dst=gc_row_broadcast[i_shuf * 32 : i_shuf * 32 + 32, 0:P_MAX], + shuffle_mask=_BROADCAST_MASK, + ) + + gc_col_strict = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=gc_col_strict, + data=Lmask, + op0=nl.multiply, + operand0=gc_p, + engine=nisa.vector_engine, + ) + gc_row_strict = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=gc_row_strict, data1=gc_row_broadcast, data2=Lmask, op=nl.multiply + ) + g_diff_strict = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=g_diff_strict, + data1=gc_col_strict, + data2=gc_row_strict, + op=nl.subtract, + ) + decay_strict_raw = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.activation( + dst=decay_strict_raw, + op=nl.exp, + data=g_diff_strict, + bias=None, + scale=1.0, + ) + decay_strict = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=decay_strict, data1=decay_strict_raw, data2=Lmask, op=nl.multiply + ) + + gc_col_diag = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=gc_col_diag, + data=Lmask_diag, + op0=nl.multiply, + operand0=gc_p, + engine=nisa.vector_engine, + ) + gc_row_diag = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=gc_row_diag, + data1=gc_row_broadcast, + data2=Lmask_diag, + op=nl.multiply, + ) + g_diff_diag = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=g_diff_diag, + data1=gc_col_diag, + data2=gc_row_diag, + op=nl.subtract, + ) + decay_diag_raw = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.activation( + dst=decay_diag_raw, + op=nl.exp, + data=g_diff_diag, + bias=None, + scale=1.0, + ) + decay_diag = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=decay_diag, data1=decay_diag_raw, data2=Lmask_diag, op=nl.multiply + ) + + k_beta = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=k_beta, + data=k_norm, + op0=nl.multiply, + operand0=beta_p, + engine=nisa.vector_engine, + ) + + v_beta = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=v_beta, + data=v_c, + op0=nl.multiply, + operand0=beta_p, + engine=nisa.vector_engine, + ) + + kb_T_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=kb_T_psum, data=k_beta) + kb_T = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=kb_T, src=kb_T_psum) + + k_T_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=k_T_psum, data=k_norm) + k_T = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=k_T, src=k_T_psum) + + QK_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=QK_psum, stationary=kb_T, moving=k_T) + QK = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=QK, src=QK_psum) + + QK_decay = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=QK_decay, data1=QK, data2=decay_strict, op=nl.multiply) + neg_QK_decay = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=neg_QK_decay, + data=QK_decay, + op0=nl.multiply, + operand0=-1.0, + engine=nisa.vector_engine, + ) + + A_T_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=A_T_psum, data=neg_QK_decay) + A_T = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=A_T, src=A_T_psum) + + value_u = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.memset(dst=value_u, value=0.0) + if SOLVE_KKT_HIER: + _hierarchical_kkt_solve128(value_u, A_T, Imat, v_beta, dim) + else: + _blocked_doubling_solve(value_u, A_T, v_beta, dim) + + kb_exp_gc = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=kb_exp_gc, + data=k_beta, + op0=nl.multiply, + operand0=exp_gc_p, + engine=nisa.vector_engine, + ) + state_w = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.memset(dst=state_w, value=0.0) + if SOLVE_KKT_HIER: + _hierarchical_kkt_solve128(state_w, A_T, Imat, kb_exp_gc, dim) + else: + _blocked_doubling_solve(state_w, A_T, kb_exp_gc, dim) + + q_T_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=q_T_psum, data=q_norm) + q_T = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=q_T, src=q_T_psum) + + qk_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=qk_psum, stationary=q_T, moving=k_T) + qk_raw = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=qk_raw, src=qk_psum) + + attn_intra = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=attn_intra, data1=qk_raw, data2=decay_diag, op=nl.multiply + ) + + ai_T_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=ai_T_psum, data=attn_intra) + ai_T = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=ai_T, src=ai_T_psum) + + output_base_psum = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=output_base_psum, stationary=ai_T, moving=value_u) + output_base = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=output_base, src=output_base_psum) + + q_exp = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=q_exp, + data=q_norm, + op0=nl.multiply, + operand0=exp_gc_p, + engine=nisa.vector_engine, + ) + + output_state_corr_psum = nl.ndarray( + (P_MAX, dim), dtype=nl.float32, buffer=nl.psum + ) + nisa.nc_matmul(dst=output_state_corr_psum, stationary=ai_T, moving=state_w) + output_state_corr = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=output_state_corr, src=output_state_corr_psum) + + output_state = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=output_state, + data1=q_exp, + data2=output_state_corr, + op=nl.subtract, + ) + + gl_minus_gc_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=gl_minus_gc_p, + data1=gl_p, + data2=gc_p, + op=nl.subtract, + ) + exp_gl_minus_gc_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.activation( + dst=exp_gl_minus_gc_p[0:P_MAX, 0:1], + op=nl.exp, + data=gl_minus_gc_p[0:P_MAX, 0:1], + bias=None, + scale=1.0, + ) + + k_raw_decay = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=k_raw_decay, + data=k_norm, + op0=nl.multiply, + operand0=exp_gl_minus_gc_p, + engine=nisa.vector_engine, + ) + + state_bias_psum = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=state_bias_psum, stationary=k_raw_decay, moving=value_u) + state_bias = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=state_bias, src=state_bias_psum) + + state_corr_psum = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=state_corr_psum, stationary=k_raw_decay, moving=state_w) + state_corr = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=state_corr, src=state_corr_psum) + + exp_gl_identity = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=exp_gl_identity, + data=Imat, + op0=nl.multiply, + operand0=exp_gl_p, + engine=nisa.vector_engine, + ) + state_matrix = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=state_matrix, + data1=exp_gl_identity, + data2=state_corr, + op=nl.subtract, + ) + + nisa.dma_copy(dst=output_base_out, src=output_base) + nisa.dma_copy(dst=output_state_out, src=output_state) + nisa.dma_copy(dst=state_matrix_out, src=state_matrix) + nisa.dma_copy(dst=state_bias_out, src=state_bias) + + return output_base_out, output_state_out, state_matrix_out, state_bias_out + + +@nki.jit +def deltanet_autocp_affine_sequence( + query: nl.ndarray, # (S, 128) float32 - raw Q; normalized in-kernel + key: nl.ndarray, # (S, 128) float32 - raw K; normalized in-kernel + value: nl.ndarray, # (S, 128) float32 + g_in: nl.ndarray, # (S, 1) float32 - per-token log-decay + beta_in: nl.ndarray, # (S, 1) float32 - per-token write gate + lower_mask: nl.ndarray, # (128, 128) float32 - strict lower tri + identity: nl.ndarray, # (128, 128) float32 - identity + lower_mask_diag: nl.ndarray, # (128, 128) float32 - kept for call compatibility +): + """Build AutoCP affine pieces for one sequence with LNC-striped chunks.""" + seq_len = query.shape[0] + dim = query.shape[1] + num_chunks = seq_len // CHUNK_SIZE + program_idx = nl.program_id(axis=0) + num_programs = nl.num_programs(axes=0) + chunks_per_program = num_chunks // num_programs + + output_base_out = nl.ndarray( + (num_chunks, P_MAX, dim), dtype=nl.float32, buffer=nl.shared_hbm + ) + output_state_out = nl.ndarray( + (num_chunks, P_MAX, dim), dtype=nl.float32, buffer=nl.shared_hbm + ) + state_matrix_out = nl.ndarray( + (num_chunks, P_MAX, dim), dtype=nl.float32, buffer=nl.shared_hbm + ) + state_bias_out = nl.ndarray( + (num_chunks, P_MAX, dim), dtype=nl.float32, buffer=nl.shared_hbm + ) + + Lmask = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=Lmask, src=lower_mask) + + Imat = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=Imat, src=identity) + + ones_1xC = nl.ndarray((1, CHUNK_SIZE), dtype=nl.float32, buffer=nl.sbuf) + nisa.memset(dst=ones_1xC, value=1.0) + + zero_11 = nl.ndarray((1, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.memset(dst=zero_11, value=0.0) + + for chunk_loop in nl.sequential_range(chunks_per_program): + chunk_idx = program_idx * chunks_per_program + chunk_loop + chunk_start = chunk_idx * CHUNK_SIZE + + q_c = nl.ndarray((P_MAX, dim), dtype=query.dtype, buffer=nl.sbuf) + nisa.dma_copy( + dst=q_c, + src=query[chunk_start : chunk_start + CHUNK_SIZE, 0:dim], + ) + + k_c = nl.ndarray((P_MAX, dim), dtype=key.dtype, buffer=nl.sbuf) + nisa.dma_copy( + dst=k_c, + src=key[chunk_start : chunk_start + CHUNK_SIZE, 0:dim], + ) + + q_square = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=q_square, data1=q_c, data2=q_c, op=nl.multiply) + q_norm_sq = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_reduce(dst=q_norm_sq, data=q_square, op=nl.add, axis=1) + q_norm_sq_clamped = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=q_norm_sq_clamped, + data=q_norm_sq, + op0=nl.maximum, + operand0=L2_EPS_SQUARED, + engine=nisa.vector_engine, + ) + q_inv_norm = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=q_inv_norm, + data=q_norm_sq_clamped, + op0=nl.rsqrt, + operand0=0.0, + engine=nisa.gpsimd_engine, + ) + q_norm = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=q_norm, + data=q_c, + op0=nl.multiply, + operand0=q_inv_norm, + op1=nl.multiply, + operand1=QUERY_SCALE, + engine=nisa.vector_engine, + ) + + k_square = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=k_square, data1=k_c, data2=k_c, op=nl.multiply) + k_norm_sq = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_reduce(dst=k_norm_sq, data=k_square, op=nl.add, axis=1) + k_norm_sq_clamped = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=k_norm_sq_clamped, + data=k_norm_sq, + op0=nl.maximum, + operand0=L2_EPS_SQUARED, + engine=nisa.vector_engine, + ) + k_inv_norm = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=k_inv_norm, + data=k_norm_sq_clamped, + op0=nl.rsqrt, + operand0=0.0, + engine=nisa.gpsimd_engine, + ) + k_norm = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=k_norm, + data=k_c, + op0=nl.multiply, + operand0=k_inv_norm, + engine=nisa.vector_engine, + ) + + v_c = nl.ndarray((P_MAX, dim), dtype=value.dtype, buffer=nl.sbuf) + nisa.dma_copy( + dst=v_c, + src=value[chunk_start : chunk_start + CHUNK_SIZE, 0:dim], + ) + + g_chunk_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy( + dst=g_chunk_p[0:CHUNK_SIZE, 0:1], + src=g_in[chunk_start : chunk_start + CHUNK_SIZE, 0:1], + ) + + beta_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy( + dst=beta_p[0:CHUNK_SIZE, 0:1], + src=beta_in[chunk_start : chunk_start + CHUNK_SIZE, 0:1], + ) + + g_padded = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.memset(dst=g_padded, value=0.0) + nisa.tensor_copy(dst=g_padded[0:CHUNK_SIZE, 0:1], src=g_chunk_p) + + g_tp_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=g_tp_psum, data=g_padded) + + g_row = nl.ndarray((1, CHUNK_SIZE), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=g_row[0:1, 0:CHUNK_SIZE], src=g_tp_psum[0:1, 0:CHUNK_SIZE]) + + gc_row = nl.ndarray((1, CHUNK_SIZE), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor_scan( + dst=gc_row[0:1, 0:CHUNK_SIZE], + data0=ones_1xC[0:1, 0:CHUNK_SIZE], + data1=g_row[0:1, 0:CHUNK_SIZE], + initial=zero_11[0:1, 0:1], + op0=nl.multiply, + op1=nl.add, + ) + + gc_padded = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.memset(dst=gc_padded, value=0.0) + nisa.tensor_copy(dst=gc_padded[0:1, 0:CHUNK_SIZE], src=gc_row) + + gc_tp_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=gc_tp_psum, data=gc_padded) + + gc_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=gc_p[0:CHUNK_SIZE, 0:1], src=gc_tp_psum[0:CHUNK_SIZE, 0:1]) + + gl_11 = nl.ndarray((1, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy( + dst=gl_11[0:1, 0:1], + src=gc_row[0:1, CHUNK_SIZE - 1 : CHUNK_SIZE], + ) + + exp_gc_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.activation( + dst=exp_gc_p[0:P_MAX, 0:1], + op=nl.exp, + data=gc_p[0:P_MAX, 0:1], + bias=None, + scale=1.0, + ) + + exp_gl_11 = nl.ndarray((1, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.activation( + dst=exp_gl_11, + op=nl.exp, + data=gl_11, + bias=None, + scale=1.0, + ) + + gl_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + for i_shuf in nl.static_range(P_MAX // 32): + nisa.nc_stream_shuffle( + src=gl_11[0:1, 0:1], + dst=gl_p[i_shuf * 32 : i_shuf * 32 + 32, 0:1], + shuffle_mask=_BROADCAST_MASK, + ) + + exp_gl_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + for i_shuf in nl.static_range(P_MAX // 32): + nisa.nc_stream_shuffle( + src=exp_gl_11[0:1, 0:1], + dst=exp_gl_p[i_shuf * 32 : i_shuf * 32 + 32, 0:1], + shuffle_mask=_BROADCAST_MASK, + ) + + gc_row_broadcast = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + for i_shuf in nl.static_range(P_MAX // 32): + nisa.nc_stream_shuffle( + src=gc_row[0:1, 0:P_MAX], + dst=gc_row_broadcast[i_shuf * 32 : i_shuf * 32 + 32, 0:P_MAX], + shuffle_mask=_BROADCAST_MASK, + ) + + gc_col_strict = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=gc_col_strict, + data=Lmask, + op0=nl.multiply, + operand0=gc_p, + engine=nisa.vector_engine, + ) + gc_row_strict = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=gc_row_strict, data1=gc_row_broadcast, data2=Lmask, op=nl.multiply + ) + g_diff_strict = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=g_diff_strict, + data1=gc_col_strict, + data2=gc_row_strict, + op=nl.subtract, + ) + decay_strict_raw = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.activation( + dst=decay_strict_raw, + op=nl.exp, + data=g_diff_strict, + bias=None, + scale=1.0, + ) + decay_strict = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=decay_strict, data1=decay_strict_raw, data2=Lmask, op=nl.multiply + ) + + decay_diag = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=decay_diag, data1=decay_strict, data2=Imat, op=nl.add) + + k_beta = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=k_beta, + data=k_norm, + op0=nl.multiply, + operand0=beta_p, + engine=nisa.vector_engine, + ) + + v_beta = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=v_beta, + data=v_c, + op0=nl.multiply, + operand0=beta_p, + engine=nisa.vector_engine, + ) + + kb_T_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=kb_T_psum, data=k_beta) + kb_T = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=kb_T, src=kb_T_psum) + + k_T_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=k_T_psum, data=k_norm) + k_T = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=k_T, src=k_T_psum) + + QK_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=QK_psum, stationary=kb_T, moving=k_T) + QK = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=QK, src=QK_psum) + + QK_decay = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=QK_decay, data1=QK, data2=decay_strict, op=nl.multiply) + neg_QK_decay = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=neg_QK_decay, + data=QK_decay, + op0=nl.multiply, + operand0=-1.0, + engine=nisa.vector_engine, + ) + + A_T_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=A_T_psum, data=neg_QK_decay) + A_T = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=A_T, src=A_T_psum) + + value_u = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.memset(dst=value_u, value=0.0) + if SOLVE_KKT_HIER: + _hierarchical_kkt_solve128(value_u, A_T, Imat, v_beta, dim) + else: + _blocked_doubling_solve(value_u, A_T, v_beta, dim) + + kb_exp_gc = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=kb_exp_gc, + data=k_beta, + op0=nl.multiply, + operand0=exp_gc_p, + engine=nisa.vector_engine, + ) + state_w = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.memset(dst=state_w, value=0.0) + if SOLVE_KKT_HIER: + _hierarchical_kkt_solve128(state_w, A_T, Imat, kb_exp_gc, dim) + else: + _blocked_doubling_solve(state_w, A_T, kb_exp_gc, dim) + + q_T_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=q_T_psum, data=q_norm) + q_T = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=q_T, src=q_T_psum) + + qk_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=qk_psum, stationary=q_T, moving=k_T) + qk_raw = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=qk_raw, src=qk_psum) + + attn_intra = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=attn_intra, data1=qk_raw, data2=decay_diag, op=nl.multiply + ) + + ai_T_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=ai_T_psum, data=attn_intra) + ai_T = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=ai_T, src=ai_T_psum) + + output_base_psum = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=output_base_psum, stationary=ai_T, moving=value_u) + output_base = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=output_base, src=output_base_psum) + + q_exp = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=q_exp, + data=q_norm, + op0=nl.multiply, + operand0=exp_gc_p, + engine=nisa.vector_engine, + ) + + output_state_corr_psum = nl.ndarray( + (P_MAX, dim), dtype=nl.float32, buffer=nl.psum + ) + nisa.nc_matmul(dst=output_state_corr_psum, stationary=ai_T, moving=state_w) + output_state_corr = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=output_state_corr, src=output_state_corr_psum) + + output_state = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=output_state, + data1=q_exp, + data2=output_state_corr, + op=nl.subtract, + ) + + gl_minus_gc_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=gl_minus_gc_p, + data1=gl_p, + data2=gc_p, + op=nl.subtract, + ) + exp_gl_minus_gc_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.activation( + dst=exp_gl_minus_gc_p[0:P_MAX, 0:1], + op=nl.exp, + data=gl_minus_gc_p[0:P_MAX, 0:1], + bias=None, + scale=1.0, + ) + + k_raw_decay = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=k_raw_decay, + data=k_norm, + op0=nl.multiply, + operand0=exp_gl_minus_gc_p, + engine=nisa.vector_engine, + ) + + state_bias_psum = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=state_bias_psum, stationary=k_raw_decay, moving=value_u) + state_bias = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=state_bias, src=state_bias_psum) + + state_corr_psum = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=state_corr_psum, stationary=k_raw_decay, moving=state_w) + state_corr = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=state_corr, src=state_corr_psum) + + exp_gl_identity = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=exp_gl_identity, + data=Imat, + op0=nl.multiply, + operand0=exp_gl_p, + engine=nisa.vector_engine, + ) + state_matrix = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=state_matrix, + data1=exp_gl_identity, + data2=state_corr, + op=nl.subtract, + ) + + nisa.dma_copy(dst=output_base_out[chunk_idx, 0:P_MAX, 0:dim], src=output_base) + nisa.dma_copy(dst=output_state_out[chunk_idx, 0:P_MAX, 0:dim], src=output_state) + nisa.dma_copy(dst=state_matrix_out[chunk_idx, 0:P_MAX, 0:dim], src=state_matrix) + nisa.dma_copy(dst=state_bias_out[chunk_idx, 0:P_MAX, 0:dim], src=state_bias) + + return output_base_out, output_state_out, state_matrix_out, state_bias_out + + +@nki.jit +def deltanet_autocp_state_summary_sequence( + key: nl.ndarray, # (S, 128) float32 - raw K; normalized in-kernel + value: nl.ndarray, # (S, 128) float32 + g_in: nl.ndarray, # (S, 1) float32 - per-token log-decay + beta_in: nl.ndarray, # (S, 1) float32 - per-token write gate + lower_mask: nl.ndarray, # (128, 128) float32 - strict lower tri + identity: nl.ndarray, # (128, 128) float32 - identity +): + """Build compact AutoCP segment state summaries. + + This is the first production-shaped AutoCP prepass: it skips query/output + affine pieces and emits only per-segment state transforms: + + state_{seg+1} = segment_matrix_seg @ state_seg + segment_bias_seg + + Segment replay can then use the existing recurrent fused kernel from the + corrected segment initial states. + """ + seq_len = key.shape[0] + dim = key.shape[1] + num_chunks = seq_len // CHUNK_SIZE + num_segments = num_chunks // AUTOCP_CP_CHUNKS + program_idx = nl.program_id(axis=0) + num_programs = nl.num_programs(axes=0) + segments_per_program = num_segments // num_programs + + segment_matrix_out = nl.ndarray( + (num_segments, P_MAX, dim), dtype=nl.float32, buffer=nl.shared_hbm + ) + segment_bias_out = nl.ndarray( + (num_segments, P_MAX, dim), dtype=nl.float32, buffer=nl.shared_hbm + ) + + Lmask = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=Lmask, src=lower_mask) + + Imat = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=Imat, src=identity) + + ones_1xC = nl.ndarray((1, CHUNK_SIZE), dtype=nl.float32, buffer=nl.sbuf) + nisa.memset(dst=ones_1xC, value=1.0) + + zero_11 = nl.ndarray((1, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.memset(dst=zero_11, value=0.0) + + for segment_loop in nl.sequential_range(segments_per_program): + segment_idx = program_idx * segments_per_program + segment_loop + first_chunk = segment_idx * AUTOCP_CP_CHUNKS + + segment_matrix = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=segment_matrix, src=Imat) + + segment_bias = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.memset(dst=segment_bias, value=0.0) + + for local_chunk in nl.sequential_range(AUTOCP_CP_CHUNKS): + chunk_idx = first_chunk + local_chunk + chunk_start = chunk_idx * CHUNK_SIZE + + k_c = nl.ndarray((P_MAX, dim), dtype=key.dtype, buffer=nl.sbuf) + nisa.dma_copy( + dst=k_c, + src=key[chunk_start : chunk_start + CHUNK_SIZE, 0:dim], + ) + + k_square = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=k_square, data1=k_c, data2=k_c, op=nl.multiply) + k_norm_sq = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_reduce(dst=k_norm_sq, data=k_square, op=nl.add, axis=1) + k_norm_sq_clamped = nl.ndarray( + (P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf + ) + nisa.tensor_scalar( + dst=k_norm_sq_clamped, + data=k_norm_sq, + op0=nl.maximum, + operand0=L2_EPS_SQUARED, + engine=nisa.vector_engine, + ) + k_inv_norm = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=k_inv_norm, + data=k_norm_sq_clamped, + op0=nl.rsqrt, + operand0=0.0, + engine=nisa.gpsimd_engine, + ) + k_norm = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=k_norm, + data=k_c, + op0=nl.multiply, + operand0=k_inv_norm, + engine=nisa.vector_engine, + ) + + v_c = nl.ndarray((P_MAX, dim), dtype=value.dtype, buffer=nl.sbuf) + nisa.dma_copy( + dst=v_c, + src=value[chunk_start : chunk_start + CHUNK_SIZE, 0:dim], + ) + + g_chunk_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy( + dst=g_chunk_p[0:CHUNK_SIZE, 0:1], + src=g_in[chunk_start : chunk_start + CHUNK_SIZE, 0:1], + ) + + beta_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy( + dst=beta_p[0:CHUNK_SIZE, 0:1], + src=beta_in[chunk_start : chunk_start + CHUNK_SIZE, 0:1], + ) + + g_padded = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.memset(dst=g_padded, value=0.0) + nisa.tensor_copy(dst=g_padded[0:CHUNK_SIZE, 0:1], src=g_chunk_p) + + g_tp_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=g_tp_psum, data=g_padded) + + g_row = nl.ndarray((1, CHUNK_SIZE), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy( + dst=g_row[0:1, 0:CHUNK_SIZE], + src=g_tp_psum[0:1, 0:CHUNK_SIZE], + ) + + gc_row = nl.ndarray((1, CHUNK_SIZE), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor_scan( + dst=gc_row[0:1, 0:CHUNK_SIZE], + data0=ones_1xC[0:1, 0:CHUNK_SIZE], + data1=g_row[0:1, 0:CHUNK_SIZE], + initial=zero_11[0:1, 0:1], + op0=nl.multiply, + op1=nl.add, + ) + + gc_padded = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.memset(dst=gc_padded, value=0.0) + nisa.tensor_copy(dst=gc_padded[0:1, 0:CHUNK_SIZE], src=gc_row) + + gc_tp_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=gc_tp_psum, data=gc_padded) + + gc_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy( + dst=gc_p[0:CHUNK_SIZE, 0:1], + src=gc_tp_psum[0:CHUNK_SIZE, 0:1], + ) + + gl_11 = nl.ndarray((1, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy( + dst=gl_11[0:1, 0:1], + src=gc_row[0:1, CHUNK_SIZE - 1 : CHUNK_SIZE], + ) + + exp_gc_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.activation( + dst=exp_gc_p[0:P_MAX, 0:1], + op=nl.exp, + data=gc_p[0:P_MAX, 0:1], + bias=None, + scale=1.0, + ) + + exp_gl_11 = nl.ndarray((1, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.activation( + dst=exp_gl_11, + op=nl.exp, + data=gl_11, + bias=None, + scale=1.0, + ) + + gl_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + for i_shuf in nl.static_range(P_MAX // 32): + nisa.nc_stream_shuffle( + src=gl_11[0:1, 0:1], + dst=gl_p[i_shuf * 32 : i_shuf * 32 + 32, 0:1], + shuffle_mask=_BROADCAST_MASK, + ) + + exp_gl_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + for i_shuf in nl.static_range(P_MAX // 32): + nisa.nc_stream_shuffle( + src=exp_gl_11[0:1, 0:1], + dst=exp_gl_p[i_shuf * 32 : i_shuf * 32 + 32, 0:1], + shuffle_mask=_BROADCAST_MASK, + ) + + gc_row_broadcast = nl.ndarray( + (P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf + ) + for i_shuf in nl.static_range(P_MAX // 32): + nisa.nc_stream_shuffle( + src=gc_row[0:1, 0:P_MAX], + dst=gc_row_broadcast[i_shuf * 32 : i_shuf * 32 + 32, 0:P_MAX], + shuffle_mask=_BROADCAST_MASK, + ) + + gc_col_strict = nl.ndarray( + (P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf + ) + nisa.tensor_scalar( + dst=gc_col_strict, + data=Lmask, + op0=nl.multiply, + operand0=gc_p, + engine=nisa.vector_engine, + ) + gc_row_strict = nl.ndarray( + (P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf + ) + nisa.tensor_tensor( + dst=gc_row_strict, + data1=gc_row_broadcast, + data2=Lmask, + op=nl.multiply, + ) + g_diff_strict = nl.ndarray( + (P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf + ) + nisa.tensor_tensor( + dst=g_diff_strict, + data1=gc_col_strict, + data2=gc_row_strict, + op=nl.subtract, + ) + decay_strict_raw = nl.ndarray( + (P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf + ) + nisa.activation( + dst=decay_strict_raw, + op=nl.exp, + data=g_diff_strict, + bias=None, + scale=1.0, + ) + decay_strict = nl.ndarray( + (P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf + ) + nisa.tensor_tensor( + dst=decay_strict, + data1=decay_strict_raw, + data2=Lmask, + op=nl.multiply, + ) + + k_beta = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=k_beta, + data=k_norm, + op0=nl.multiply, + operand0=beta_p, + engine=nisa.vector_engine, + ) + + v_beta = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=v_beta, + data=v_c, + op0=nl.multiply, + operand0=beta_p, + engine=nisa.vector_engine, + ) + + kb_T_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=kb_T_psum, data=k_beta) + kb_T = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=kb_T, src=kb_T_psum) + + k_T_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=k_T_psum, data=k_norm) + k_T = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=k_T, src=k_T_psum) + + QK_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=QK_psum, stationary=kb_T, moving=k_T) + QK = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=QK, src=QK_psum) + + QK_decay = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=QK_decay, + data1=QK, + data2=decay_strict, + op=nl.multiply, + ) + neg_QK_decay = nl.ndarray( + (P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf + ) + nisa.tensor_scalar( + dst=neg_QK_decay, + data=QK_decay, + op0=nl.multiply, + operand0=-1.0, + engine=nisa.vector_engine, + ) + + A_T_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=A_T_psum, data=neg_QK_decay) + A_T = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=A_T, src=A_T_psum) + + value_u = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.memset(dst=value_u, value=0.0) + if SOLVE_KKT_HIER: + _hierarchical_kkt_solve128(value_u, A_T, Imat, v_beta, dim) + else: + _blocked_doubling_solve(value_u, A_T, v_beta, dim) + + kb_exp_gc = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=kb_exp_gc, + data=k_beta, + op0=nl.multiply, + operand0=exp_gc_p, + engine=nisa.vector_engine, + ) + state_w = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.memset(dst=state_w, value=0.0) + if SOLVE_KKT_HIER: + _hierarchical_kkt_solve128(state_w, A_T, Imat, kb_exp_gc, dim) + else: + _blocked_doubling_solve(state_w, A_T, kb_exp_gc, dim) + + gl_minus_gc_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=gl_minus_gc_p, + data1=gl_p, + data2=gc_p, + op=nl.subtract, + ) + exp_gl_minus_gc_p = nl.ndarray( + (P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf + ) + nisa.activation( + dst=exp_gl_minus_gc_p[0:P_MAX, 0:1], + op=nl.exp, + data=gl_minus_gc_p[0:P_MAX, 0:1], + bias=None, + scale=1.0, + ) + + k_raw_decay = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=k_raw_decay, + data=k_norm, + op0=nl.multiply, + operand0=exp_gl_minus_gc_p, + engine=nisa.vector_engine, + ) + + state_bias = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + state_bias_psum = nl.ndarray( + (P_MAX, dim), dtype=nl.float32, buffer=nl.psum + ) + nisa.nc_matmul(dst=state_bias_psum, stationary=k_raw_decay, moving=value_u) + nisa.tensor_copy(dst=state_bias, src=state_bias_psum) + + state_corr_psum = nl.ndarray( + (P_MAX, dim), dtype=nl.float32, buffer=nl.psum + ) + nisa.nc_matmul(dst=state_corr_psum, stationary=k_raw_decay, moving=state_w) + state_corr = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=state_corr, src=state_corr_psum) + + exp_gl_identity = nl.ndarray( + (P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf + ) + nisa.tensor_scalar( + dst=exp_gl_identity, + data=Imat, + op0=nl.multiply, + operand0=exp_gl_p, + engine=nisa.vector_engine, + ) + state_matrix = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=state_matrix, + data1=exp_gl_identity, + data2=state_corr, + op=nl.subtract, + ) + + state_matrix_t_psum = nl.ndarray( + (P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum + ) + nisa.nc_transpose(dst=state_matrix_t_psum, data=state_matrix) + state_matrix_t = nl.ndarray( + (P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf + ) + nisa.tensor_copy(dst=state_matrix_t, src=state_matrix_t_psum) + + composed_matrix_psum = nl.ndarray( + (P_MAX, dim), dtype=nl.float32, buffer=nl.psum + ) + nisa.nc_matmul( + dst=composed_matrix_psum, + stationary=state_matrix_t, + moving=segment_matrix, + ) + composed_matrix = nl.ndarray( + (P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf + ) + nisa.tensor_copy(dst=composed_matrix, src=composed_matrix_psum) + + propagated_bias_psum = nl.ndarray( + (P_MAX, dim), dtype=nl.float32, buffer=nl.psum + ) + nisa.nc_matmul( + dst=propagated_bias_psum, + stationary=state_matrix_t, + moving=segment_bias, + ) + propagated_bias = nl.ndarray( + (P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf + ) + nisa.tensor_copy(dst=propagated_bias, src=propagated_bias_psum) + + composed_bias = nl.ndarray( + (P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf + ) + nisa.tensor_tensor( + dst=composed_bias, + data1=propagated_bias, + data2=state_bias, + op=nl.add, + ) + + nisa.tensor_copy(dst=segment_matrix, src=composed_matrix) + nisa.tensor_copy(dst=segment_bias, src=composed_bias) + + nisa.dma_copy( + dst=segment_matrix_out[segment_idx, 0:P_MAX, 0:dim], + src=segment_matrix, + ) + nisa.dma_copy( + dst=segment_bias_out[segment_idx, 0:P_MAX, 0:dim], + src=segment_bias, + ) + + return segment_matrix_out, segment_bias_out + + +@nki.jit +def deltanet_autocp_state_prefix( + state_matrix: nl.ndarray, # (N, 128, 128) float32 + state_bias: nl.ndarray, # (N, 128, 128) float32 + initial_state: nl.ndarray, # (128, 128) float32 +): + """Apply per-chunk AutoCP state transforms and emit chunk initial states. + + Given per-chunk transforms: + state_{i+1} = state_matrix_i @ state_i + state_bias_i + + returns: + chunk_states[i] = state_i + final_state = state_N + + This is the isolated state-prefix correctness probe. A later production + path can replace the loop body with a tree/parallel prefix over the same + HBM interface. + """ + num_chunks = state_matrix.shape[0] + dim = initial_state.shape[1] + + chunk_states_out = nl.ndarray( + (num_chunks, P_MAX, dim), dtype=nl.float32, buffer=nl.shared_hbm + ) + final_state_out = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.shared_hbm) + + state = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=state, src=initial_state[0:P_MAX, 0:dim]) + + for i_chunk in nl.sequential_range(num_chunks): + nisa.dma_copy( + dst=chunk_states_out[i_chunk, 0:P_MAX, 0:dim], + src=state, + ) + + matrix = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy( + dst=matrix, + src=state_matrix[i_chunk, 0:P_MAX, 0:dim], + ) + + bias = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy( + dst=bias, + src=state_bias[i_chunk, 0:P_MAX, 0:dim], + ) + + matrix_t_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=matrix_t_psum, data=matrix) + matrix_t = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=matrix_t, src=matrix_t_psum) + + propagated_psum = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=propagated_psum, stationary=matrix_t, moving=state) + propagated = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=propagated, src=propagated_psum) + + next_state = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=next_state, + data1=propagated, + data2=bias, + op=nl.add, + ) + nisa.tensor_copy(dst=state, src=next_state) + + nisa.dma_copy(dst=final_state_out, src=state) + + return chunk_states_out, final_state_out + + +@nki.jit +def deltanet_autocp_apply_output( + output_base: nl.ndarray, # (N, 128, 128) float32 + output_state: nl.ndarray, # (N, 128, 128) float32 + chunk_states: nl.ndarray, # (N, 128, 128) float32 +): + """Apply AutoCP chunk initial states to state-dependent output terms.""" + num_chunks = output_base.shape[0] + dim = output_base.shape[2] + + output = nl.ndarray( + (num_chunks * CHUNK_SIZE, dim), dtype=nl.float32, buffer=nl.shared_hbm + ) + + for i_chunk in nl.sequential_range(num_chunks): + base = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy( + dst=base, + src=output_base[i_chunk, 0:P_MAX, 0:dim], + ) + + state_coeff = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy( + dst=state_coeff, + src=output_state[i_chunk, 0:P_MAX, 0:dim], + ) + + state = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy( + dst=state, + src=chunk_states[i_chunk, 0:P_MAX, 0:dim], + ) + + coeff_t_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=coeff_t_psum, data=state_coeff) + coeff_t = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=coeff_t, src=coeff_t_psum) + + state_out_psum = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=state_out_psum, stationary=coeff_t, moving=state) + state_out = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=state_out, src=state_out_psum) + + chunk_output = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=chunk_output, + data1=base, + data2=state_out, + op=nl.add, + ) + + nisa.dma_copy( + dst=output[i_chunk * CHUNK_SIZE : i_chunk * CHUNK_SIZE + CHUNK_SIZE, 0:dim], + src=chunk_output, + ) + + return output + + +@nki.jit +def deltanet_autocp_prefix_apply_output( + output_base: nl.ndarray, # (N, 128, 128) float32 + output_state: nl.ndarray, # (N, 128, 128) float32 + state_matrix: nl.ndarray, # (N, 128, 128) float32 + state_bias: nl.ndarray, # (N, 128, 128) float32 + initial_state: nl.ndarray, # (128, 128) float32 +): + """Fused AutoCP state-prefix and output-apply pass. + + This removes the intermediate chunk_states HBM tensor and one custom-call + from the AutoCP probe path. It intentionally remains an exact sequential + prefix over dense 128x128 chunk transforms; a matrix-affine prefix cannot be + represented by tensor_tensor_scan's elementwise recurrence. + """ + num_chunks = output_base.shape[0] + dim = output_base.shape[2] + + output = nl.ndarray( + (num_chunks * CHUNK_SIZE, dim), dtype=nl.float32, buffer=nl.shared_hbm + ) + final_state_out = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.shared_hbm) + + state = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=state, src=initial_state[0:P_MAX, 0:dim]) + + for i_chunk in nl.sequential_range(num_chunks): + base = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy( + dst=base, + src=output_base[i_chunk, 0:P_MAX, 0:dim], + ) + + state_coeff = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy( + dst=state_coeff, + src=output_state[i_chunk, 0:P_MAX, 0:dim], + ) + + coeff_t_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=coeff_t_psum, data=state_coeff) + coeff_t = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=coeff_t, src=coeff_t_psum) + + state_out_psum = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=state_out_psum, stationary=coeff_t, moving=state) + state_out = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=state_out, src=state_out_psum) + + chunk_output = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=chunk_output, + data1=base, + data2=state_out, + op=nl.add, + ) + nisa.dma_copy( + dst=output[i_chunk * CHUNK_SIZE : i_chunk * CHUNK_SIZE + CHUNK_SIZE, 0:dim], + src=chunk_output, + ) + + matrix = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy( + dst=matrix, + src=state_matrix[i_chunk, 0:P_MAX, 0:dim], + ) + + bias = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy( + dst=bias, + src=state_bias[i_chunk, 0:P_MAX, 0:dim], + ) + + matrix_t_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=matrix_t_psum, data=matrix) + matrix_t = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=matrix_t, src=matrix_t_psum) + + propagated_psum = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=propagated_psum, stationary=matrix_t, moving=state) + propagated = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=propagated, src=propagated_psum) + + next_state = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=next_state, + data1=propagated, + data2=bias, + op=nl.add, + ) + nisa.tensor_copy(dst=state, src=next_state) + + nisa.dma_copy(dst=final_state_out, src=state) + + return output, final_state_out + + +@nki.jit +def deltanet_fused_chunked_fwd_multihead( + query: nl.ndarray, # (BH, S, 128) float32 — raw Q; normalized in-kernel + key: nl.ndarray, # (BH, S, 128) float32 — raw K; normalized in-kernel + value: nl.ndarray, # (BH, S, 128) float32 + g_in: nl.ndarray, # (BH, S, 1) float32 + beta_in: nl.ndarray, # (BH, S, 1) float32 + initial_state: nl.ndarray, # (BH, 128, 128) float32 + lower_mask: nl.ndarray, # (128, 128) float32 + identity: nl.ndarray, # (128, 128) float32 + lower_mask_diag: nl.ndarray, # (128, 128) float32 +): + """Fused chunked DeltaNet forward for one or more heads with SPMD sharding.""" + num_heads = query.shape[0] + seq_len = query.shape[1] + dim = query.shape[2] + num_chunks = seq_len // CHUNK_SIZE + head_idx = nl.program_id(axis=0) + + output = nl.ndarray( + (num_heads, seq_len, dim), dtype=query.dtype, buffer=nl.shared_hbm + ) + final_state_out = nl.ndarray( + (num_heads, P_MAX, dim), dtype=nl.float32, buffer=nl.shared_hbm + ) + + Lmask = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=Lmask, src=lower_mask) + + UMask_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=UMask_psum, data=Lmask) + UMask = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=UMask, src=UMask_psum) + + Imat = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=Imat, src=identity) + + ones_1xC = nl.ndarray((1, CHUNK_SIZE), dtype=nl.float32, buffer=nl.sbuf) + nisa.memset(dst=ones_1xC, value=1.0) + + zero_11 = nl.ndarray((1, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.memset(dst=zero_11, value=0.0) + + state = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=state, src=initial_state[head_idx, 0:P_MAX, 0:dim]) + + for i_chunk in nl.sequential_range(num_chunks): + chunk_start = i_chunk * CHUNK_SIZE + + q_c = nl.ndarray((P_MAX, dim), dtype=query.dtype, buffer=nl.sbuf) + if CHUNK_SIZE == P_MAX: + nisa.dma_copy( + dst=q_c, + src=query[head_idx, chunk_start : chunk_start + CHUNK_SIZE, 0:dim], + ) + else: + nisa.memset(dst=q_c, value=0.0) + nisa.dma_copy( + dst=q_c[0:CHUNK_SIZE, 0:dim], + src=query[head_idx, chunk_start : chunk_start + CHUNK_SIZE, 0:dim], + ) + + k_c = nl.ndarray((P_MAX, dim), dtype=key.dtype, buffer=nl.sbuf) + if CHUNK_SIZE == P_MAX: + nisa.dma_copy( + dst=k_c, + src=key[head_idx, chunk_start : chunk_start + CHUNK_SIZE, 0:dim], + ) + else: + nisa.memset(dst=k_c, value=0.0) + nisa.dma_copy( + dst=k_c[0:CHUNK_SIZE, 0:dim], + src=key[head_idx, chunk_start : chunk_start + CHUNK_SIZE, 0:dim], + ) + + q_square = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=q_square, data1=q_c, data2=q_c, op=nl.multiply) + q_norm_sq = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_reduce(dst=q_norm_sq, data=q_square, op=nl.add, axis=1) + q_norm_sq_clamped = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=q_norm_sq_clamped, + data=q_norm_sq, + op0=nl.maximum, + operand0=L2_EPS_SQUARED, + engine=nisa.vector_engine, + ) + q_inv_norm = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=q_inv_norm, + data=q_norm_sq_clamped, + op0=nl.rsqrt, + operand0=0.0, + engine=nisa.gpsimd_engine, + ) + q_norm = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=q_norm, + data=q_c, + op0=nl.multiply, + operand0=q_inv_norm, + op1=nl.multiply, + operand1=QUERY_SCALE, + engine=nisa.vector_engine, + ) + + k_square = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=k_square, data1=k_c, data2=k_c, op=nl.multiply) + k_norm_sq = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_reduce(dst=k_norm_sq, data=k_square, op=nl.add, axis=1) + k_norm_sq_clamped = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=k_norm_sq_clamped, + data=k_norm_sq, + op0=nl.maximum, + operand0=L2_EPS_SQUARED, + engine=nisa.vector_engine, + ) + k_inv_norm = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=k_inv_norm, + data=k_norm_sq_clamped, + op0=nl.rsqrt, + operand0=0.0, + engine=nisa.gpsimd_engine, + ) + k_norm = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=k_norm, + data=k_c, + op0=nl.multiply, + operand0=k_inv_norm, + engine=nisa.vector_engine, + ) + + v_c = nl.ndarray((P_MAX, dim), dtype=value.dtype, buffer=nl.sbuf) + if CHUNK_SIZE == P_MAX: + nisa.dma_copy( + dst=v_c, + src=value[head_idx, chunk_start : chunk_start + CHUNK_SIZE, 0:dim], + ) + else: + nisa.memset(dst=v_c, value=0.0) + nisa.dma_copy( + dst=v_c[0:CHUNK_SIZE, 0:dim], + src=value[head_idx, chunk_start : chunk_start + CHUNK_SIZE, 0:dim], + ) + + g_chunk_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + if CHUNK_SIZE != P_MAX: + nisa.memset(dst=g_chunk_p, value=0.0) + nisa.dma_copy( + dst=g_chunk_p[0:CHUNK_SIZE, 0:1], + src=g_in[head_idx, chunk_start : chunk_start + CHUNK_SIZE, 0:1], + ) + + beta_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + if CHUNK_SIZE != P_MAX: + nisa.memset(dst=beta_p, value=0.0) + nisa.dma_copy( + dst=beta_p[0:CHUNK_SIZE, 0:1], + src=beta_in[head_idx, chunk_start : chunk_start + CHUNK_SIZE, 0:1], + ) + + g_tp_psum = nl.ndarray((1, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=g_tp_psum, data=g_chunk_p) + + g_row = nl.ndarray((1, CHUNK_SIZE), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy( + dst=g_row[0:1, 0:CHUNK_SIZE], + src=g_tp_psum[0:1, 0:CHUNK_SIZE], + ) + + gc_row = nl.ndarray((1, CHUNK_SIZE), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor_scan( + dst=gc_row[0:1, 0:CHUNK_SIZE], + data0=ones_1xC[0:1, 0:CHUNK_SIZE], + data1=g_row[0:1, 0:CHUNK_SIZE], + initial=zero_11[0:1, 0:1], + op0=nl.multiply, + op1=nl.add, + ) + + gc_tp_psum = nl.ndarray((CHUNK_SIZE, 1), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=gc_tp_psum, data=gc_row) + + gc_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + if CHUNK_SIZE != P_MAX: + nisa.memset(dst=gc_p, value=0.0) + nisa.tensor_copy( + dst=gc_p[0:CHUNK_SIZE, 0:1], + src=gc_tp_psum[0:CHUNK_SIZE, 0:1], + ) + + gl_11 = nl.ndarray((1, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy( + dst=gl_11[0:1, 0:1], + src=gc_row[0:1, CHUNK_SIZE - 1 : CHUNK_SIZE], + ) + + exp_gc_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.activation( + dst=exp_gc_p[0:P_MAX, 0:1], + op=nl.exp, + data=gc_p[0:P_MAX, 0:1], + bias=None, + scale=1.0, + ) + + exp_gl_11 = nl.ndarray((1, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.activation( + dst=exp_gl_11, + op=nl.exp, + data=gl_11, + bias=None, + scale=1.0, + ) + + gl_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + for i_shuf in nl.static_range(P_MAX // 32): + nisa.nc_stream_shuffle( + src=gl_11[0:1, 0:1], + dst=gl_p[i_shuf * 32 : i_shuf * 32 + 32, 0:1], + shuffle_mask=_BROADCAST_MASK, + ) + + exp_gl_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + for i_shuf in nl.static_range(P_MAX // 32): + nisa.nc_stream_shuffle( + src=exp_gl_11[0:1, 0:1], + dst=exp_gl_p[i_shuf * 32 : i_shuf * 32 + 32, 0:1], + shuffle_mask=_BROADCAST_MASK, + ) + + gc_row_broadcast = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + if CHUNK_SIZE != P_MAX: + nisa.memset(dst=gc_row_broadcast, value=0.0) + for i_shuf in nl.static_range(P_MAX // 32): + nisa.nc_stream_shuffle( + src=gc_row[0:1, 0:CHUNK_SIZE], + dst=gc_row_broadcast[ + i_shuf * 32 : i_shuf * 32 + 32, 0:CHUNK_SIZE + ], + shuffle_mask=_BROADCAST_MASK, + ) + + gc_col_strict_t = nl.ndarray( + (P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf + ) + nisa.tensor_tensor( + dst=gc_col_strict_t, + data1=gc_row_broadcast, + data2=UMask, + op=nl.multiply, + ) + gc_row_strict_t = nl.ndarray( + (P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf + ) + nisa.tensor_scalar( + dst=gc_row_strict_t, + data=UMask, + op0=nl.multiply, + operand0=gc_p, + engine=nisa.vector_engine, + ) + g_diff_strict_t = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=g_diff_strict_t, + data1=gc_col_strict_t, + data2=gc_row_strict_t, + op=nl.subtract, + ) + decay_strict_t_raw = nl.ndarray( + (P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf + ) + nisa.activation( + dst=decay_strict_t_raw, + op=nl.exp, + data=g_diff_strict_t, + bias=None, + scale=1.0, + ) + decay_strict_t = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=decay_strict_t, + data1=decay_strict_t_raw, + data2=UMask, + op=nl.multiply, + ) + + decay_diag_t = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=decay_diag_t, data1=decay_strict_t, data2=Imat, op=nl.add + ) + + k_beta = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=k_beta, + data=k_norm, + op0=nl.multiply, + operand0=beta_p, + engine=nisa.vector_engine, + ) + + v_beta = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=v_beta, + data=v_c, + op0=nl.multiply, + operand0=beta_p, + engine=nisa.vector_engine, + ) + + kb_T_psum = nl.ndarray((P_MAX, CHUNK_SIZE), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=kb_T_psum, data=k_beta[0:CHUNK_SIZE, 0:dim]) + kb_T = nl.ndarray((P_MAX, CHUNK_SIZE), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=kb_T, src=kb_T_psum) + + k_T_psum = nl.ndarray((P_MAX, CHUNK_SIZE), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=k_T_psum, data=k_norm[0:CHUNK_SIZE, 0:dim]) + k_T = nl.ndarray((P_MAX, CHUNK_SIZE), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=k_T, src=k_T_psum) + + QK_T_psum = nl.ndarray( + (CHUNK_SIZE, CHUNK_SIZE), dtype=nl.float32, buffer=nl.psum + ) + nisa.nc_matmul(dst=QK_T_psum, stationary=k_T, moving=kb_T) + QK_T = nl.ndarray((CHUNK_SIZE, CHUNK_SIZE), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=QK_T, src=QK_T_psum) + + QK_decay_t = nl.ndarray( + (CHUNK_SIZE, CHUNK_SIZE), dtype=nl.float32, buffer=nl.sbuf + ) + nisa.tensor_tensor( + dst=QK_decay_t, + data1=QK_T, + data2=decay_strict_t[0:CHUNK_SIZE, 0:CHUNK_SIZE], + op=nl.multiply, + ) + + A_T = nl.ndarray((CHUNK_SIZE, CHUNK_SIZE), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=A_T, + data=QK_decay_t, + op0=nl.multiply, + operand0=-1.0, + engine=nisa.vector_engine, + ) + kb_exp_gc = nl.ndarray((CHUNK_SIZE, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=kb_exp_gc, + data=k_beta[0:CHUNK_SIZE, 0:dim], + op0=nl.multiply, + operand0=exp_gc_p[0:CHUNK_SIZE, 0:1], + engine=nisa.vector_engine, + ) + + kbe_T_psum = nl.ndarray((P_MAX, CHUNK_SIZE), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=kbe_T_psum, data=kb_exp_gc) + kbe_T = nl.ndarray((P_MAX, CHUNK_SIZE), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=kbe_T, src=kbe_T_psum) + + kbe_state_psum = nl.ndarray((CHUNK_SIZE, dim), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=kbe_state_psum, stationary=kbe_T, moving=state) + kbe_state = nl.ndarray((CHUNK_SIZE, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=kbe_state, src=kbe_state_psum) + + solve_rhs = nl.ndarray((CHUNK_SIZE, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=solve_rhs, + data1=v_beta[0:CHUNK_SIZE, 0:dim], + data2=kbe_state, + op=nl.subtract, + ) + + v_new = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.memset(dst=v_new, value=0.0) + + if SOLVE_KKT_HIER: + _hierarchical_kkt_solve128(v_new, A_T, Imat, solve_rhs, dim) + else: + _blocked_doubling_solve(v_new, A_T, solve_rhs, dim) + + q_T_psum = nl.ndarray((P_MAX, CHUNK_SIZE), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=q_T_psum, data=q_norm[0:CHUNK_SIZE, 0:dim]) + q_T = nl.ndarray((P_MAX, CHUNK_SIZE), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=q_T, src=q_T_psum) + + qk_T_psum = nl.ndarray( + (CHUNK_SIZE, CHUNK_SIZE), dtype=nl.float32, buffer=nl.psum + ) + nisa.nc_matmul(dst=qk_T_psum, stationary=k_T, moving=q_T) + qk_raw_t = nl.ndarray( + (CHUNK_SIZE, CHUNK_SIZE), dtype=nl.float32, buffer=nl.sbuf + ) + nisa.tensor_copy(dst=qk_raw_t, src=qk_T_psum) + + ai_T = nl.ndarray((CHUNK_SIZE, CHUNK_SIZE), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=ai_T, + data1=qk_raw_t, + data2=decay_diag_t[0:CHUNK_SIZE, 0:CHUNK_SIZE], + op=nl.multiply, + ) + + q_exp = nl.ndarray((CHUNK_SIZE, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=q_exp, + data=q_norm[0:CHUNK_SIZE, 0:dim], + op0=nl.multiply, + operand0=exp_gc_p[0:CHUNK_SIZE, 0:1], + engine=nisa.vector_engine, + ) + + qe_T_psum = nl.ndarray((P_MAX, CHUNK_SIZE), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=qe_T_psum, data=q_exp) + qe_T = nl.ndarray((P_MAX, CHUNK_SIZE), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=qe_T, src=qe_T_psum) + + ai_psum = nl.ndarray((CHUNK_SIZE, dim), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=ai_psum, stationary=qe_T, moving=state) + attn_inter = nl.ndarray((CHUNK_SIZE, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=attn_inter, src=ai_psum) + + intra_psum = nl.ndarray((CHUNK_SIZE, dim), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul( + dst=intra_psum, + stationary=ai_T, + moving=v_new[0:CHUNK_SIZE, 0:dim], + ) + intra_out = nl.ndarray((CHUNK_SIZE, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=intra_out, src=intra_psum) + + chunk_out = nl.ndarray((CHUNK_SIZE, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=chunk_out, data1=attn_inter, data2=intra_out, op=nl.add) + + nisa.dma_copy( + dst=output[head_idx, chunk_start : chunk_start + CHUNK_SIZE, 0:dim], + src=chunk_out, + ) + + exp_gl_minus_gc_p = nl.ndarray( + (CHUNK_SIZE, 1), dtype=nl.float32, buffer=nl.sbuf + ) + nisa.activation( + dst=exp_gl_minus_gc_p[0:CHUNK_SIZE, 0:1], + op=nl.exp, + data=gc_p[0:CHUNK_SIZE, 0:1], + bias=gl_p[0:CHUNK_SIZE, 0:1], + scale=-1.0, + ) + + k_raw_decay = nl.ndarray((CHUNK_SIZE, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=k_raw_decay, + data=k_norm[0:CHUNK_SIZE, 0:dim], + op0=nl.multiply, + operand0=exp_gl_minus_gc_p, + engine=nisa.vector_engine, + ) + + kv_psum = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul( + dst=kv_psum, + stationary=k_raw_decay, + moving=v_new[0:CHUNK_SIZE, 0:dim], + ) + kv_outer = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=kv_outer, src=kv_psum) + + state_decayed = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=state_decayed, + data=state, + op0=nl.multiply, + operand0=exp_gl_p, + engine=nisa.vector_engine, + ) + nisa.tensor_tensor(dst=state, data1=state_decayed, data2=kv_outer, op=nl.add) + + nisa.dma_copy(dst=final_state_out[head_idx, 0:P_MAX, 0:dim], src=state) + + return output, final_state_out diff --git a/contrib/models/Qwen3.5-35B-A3B/src/nki_kernels/nki_deltanet_fused_legacy.py b/contrib/models/Qwen3.5-35B-A3B/src/nki_kernels/nki_deltanet_fused_legacy.py new file mode 100644 index 00000000..087dff81 --- /dev/null +++ b/contrib/models/Qwen3.5-35B-A3B/src/nki_kernels/nki_deltanet_fused_legacy.py @@ -0,0 +1,613 @@ +"""Fused single-kernel DeltaNet chunked forward for CTE (context encoding). + +SSD-style architecture: processes ALL chunks for one (batch, head) pair in +a single NKI kernel call. State (128x128) persists in SBUF across chunks — +no HBM round-trips for inter-chunk state propagation. + +Key optimizations over nki_deltanet_chunked.py: + 1. Single kernel call per (B,H) instead of B*H*num_chunks calls + 2. State in SBUF across all chunks (no HBM state read/write per chunk) + 3. In-kernel cumsum via tensor_tensor_scan (no PyTorch cumsum) + 4. Masks and constants loaded once, reused across chunks + 5. Uses tensor_scalar for partition-broadcast (no explicit broadcast loops) + 6. nc_transpose (Vector Engine) for all 128x128 transposes instead of + nc_matmul(moving=eye) (Tensor Engine) — frees TE for actual math + +neuronx-cc 2.26.6360 / nki 0.5.0. k_dim = v_dim = 128 = P_MAX exactly. +Chunk size = 128 = P_MAX (one tile per chunk). + +Mathematical framework: + Per-chunk direct triangular solve for intra-chunk correction: + QK_decay[i,j] = QK[i,j] * exp(gc[i] - gc[j]) for i > j + A = -QK_decay * lower_mask + v_new = solve((I - A), v_beta - (k_beta * exp(gc)) @ state) + + Inter-chunk state propagation: + attn_inter = (q * exp(gc)) @ state + attn_intra = (q @ k^T) * decay_mask * lower_mask_diag + output = attn_inter + attn_intra @ v_new + state = exp(g_last) * (state + k_raw_decay^T @ v_new) +""" + +import numpy as np + +import nki +import nki.isa as nisa +import nki.language as nl + +P_MAX = 128 # Partition dim = chunk_size = k_dim = v_dim +CHUNK_SIZE = 128 + +# Broadcast partition 0 to all partitions in a 32-wide group +_BROADCAST_MASK = [0] * 32 + + +def _make_lower_mask(): + """Strict lower triangular (128x128) as numpy constant.""" + return np.tril(np.ones((CHUNK_SIZE, CHUNK_SIZE), dtype=np.float32), k=-1) + + +def _make_lower_mask_diag(): + """Lower triangular with diagonal (128x128) as numpy constant.""" + return np.tril(np.ones((CHUNK_SIZE, CHUNK_SIZE), dtype=np.float32), k=0) + + +def _make_identity(): + """Identity matrix (128x128) as numpy constant.""" + return np.eye(CHUNK_SIZE, dtype=np.float32) + + +@nki.jit +def deltanet_fused_chunked_fwd( + query: nl.ndarray, # (S, 128) float32 — l2-normed and scaled + key: nl.ndarray, # (S, 128) float32 — l2-normed + value: nl.ndarray, # (S, 128) float32 + g_in: nl.ndarray, # (S, 1) float32 — per-token log-decay (NOT cumsum) + beta_in: nl.ndarray, # (S, 1) float32 — per-token write gate + initial_state: nl.ndarray, # (128, 128) float32 — recurrent checkpoint or zeros + lower_mask: nl.ndarray, # (128, 128) float32 — strict lower tri + identity: nl.ndarray, # (128, 128) float32 — identity + lower_mask_diag: nl.ndarray, # (128, 128) float32 — lower tri with diag +): + """Fused chunked DeltaNet forward — single kernel call per (batch, head). + + Processes all chunks sequentially within the kernel, keeping the recurrent + state (128x128) in SBUF across chunks. Returns per-token output and + final state. + + Input requirements: + - S must be divisible by 128 (pad before calling) + - query must be l2-normed and scaled by 1/sqrt(k_dim) + - key must be l2-normed + - g_in is RAW log-decay (cumsum computed in-kernel via tensor_tensor_scan) + - beta_in is sigmoid(b) (write gate) + - initial_state is zero for cold prefill, or the restored GDN checkpoint + + Returns: + output: (S, 128) float32 + final_state: (128, 128) float32 + """ + seq_len = query.shape[0] + dim = query.shape[1] # 128 + num_chunks = seq_len // CHUNK_SIZE + + # Output tensors in HBM + output = nl.ndarray((seq_len, dim), dtype=query.dtype, buffer=nl.shared_hbm) + final_state_out = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.shared_hbm) + + # ================================================================ + # Load constant masks into SBUF once (reused across all chunks) + # ================================================================ + eye = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=eye, src=identity) + + Lmask = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=Lmask, src=lower_mask) + + Lmask_d = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=Lmask_d, src=lower_mask_diag) + + # Ones vector for cumsum scan: (1, CHUNK_SIZE) + ones_1xC = nl.ndarray((1, CHUNK_SIZE), dtype=nl.float32, buffer=nl.sbuf) + nisa.memset(dst=ones_1xC, value=1.0) + + # Zero initial for cumsum scan + zero_11 = nl.ndarray((1, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.memset(dst=zero_11, value=0.0) + + # ================================================================ + # Initialize recurrent state in SBUF — persists across ALL chunks + # ================================================================ + state = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy(dst=state, src=initial_state) + + # ================================================================ + # Sequential chunk processing + # ================================================================ + for i_chunk in nl.sequential_range(num_chunks): + chunk_start = i_chunk * CHUNK_SIZE + + # ---- Load chunk data from HBM ---- + q_c = nl.ndarray((P_MAX, dim), dtype=query.dtype, buffer=nl.sbuf) + nisa.dma_copy( + dst=q_c, + src=query[chunk_start : chunk_start + CHUNK_SIZE, 0:dim], + ) + + k_c = nl.ndarray((P_MAX, dim), dtype=key.dtype, buffer=nl.sbuf) + nisa.dma_copy( + dst=k_c, + src=key[chunk_start : chunk_start + CHUNK_SIZE, 0:dim], + ) + + v_c = nl.ndarray((P_MAX, dim), dtype=value.dtype, buffer=nl.sbuf) + nisa.dma_copy( + dst=v_c, + src=value[chunk_start : chunk_start + CHUNK_SIZE, 0:dim], + ) + + # g: (CHUNK_SIZE, 1) — raw log-decay per token + g_chunk_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy( + dst=g_chunk_p[0:CHUNK_SIZE, 0:1], + src=g_in[chunk_start : chunk_start + CHUNK_SIZE, 0:1], + ) + + # beta: (CHUNK_SIZE, 1) — write gate scalar per token + beta_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy( + dst=beta_p[0:CHUNK_SIZE, 0:1], + src=beta_in[chunk_start : chunk_start + CHUNK_SIZE, 0:1], + ) + + # ---- In-kernel cumsum of g via tensor_tensor_scan ---- + # Need g as (1, CHUNK_SIZE) for scan along free dim. + # Transpose: (CHUNK_SIZE, 1) -> (1, CHUNK_SIZE) via nc_transpose + g_padded = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.memset(dst=g_padded, value=0.0) + nisa.tensor_copy( + dst=g_padded[0:CHUNK_SIZE, 0:1], + src=g_chunk_p[0:CHUNK_SIZE, 0:1], + ) + + g_tp_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=g_tp_psum, data=g_padded) + + g_row = nl.ndarray((1, CHUNK_SIZE), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy( + dst=g_row[0:1, 0:CHUNK_SIZE], + src=g_tp_psum[0:1, 0:CHUNK_SIZE], + ) + + # cumsum: gc_row[t] = 1.0 * gc_row[t-1] + g_row[t] + gc_row = nl.ndarray((1, CHUNK_SIZE), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor_scan( + dst=gc_row[0:1, 0:CHUNK_SIZE], + data0=ones_1xC[0:1, 0:CHUNK_SIZE], + data1=g_row[0:1, 0:CHUNK_SIZE], + initial=zero_11[0:1, 0:1], + op0=nl.multiply, + op1=nl.add, + ) + + # Transpose gc back to (CHUNK_SIZE, 1) partition layout + gc_padded = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.memset(dst=gc_padded, value=0.0) + nisa.tensor_copy( + dst=gc_padded[0:1, 0:CHUNK_SIZE], + src=gc_row[0:1, 0:CHUNK_SIZE], + ) + + gc_tp_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=gc_tp_psum, data=gc_padded) + + # gc_p: (P_MAX, 1) — cumulative sum of g per token in this chunk + gc_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy( + dst=gc_p[0:CHUNK_SIZE, 0:1], + src=gc_tp_psum[0:CHUNK_SIZE, 0:1], + ) + + # g_last = gc[-1] (scalar) — needed for state decay + gl_11 = nl.ndarray((1, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy( + dst=gl_11[0:1, 0:1], + src=gc_row[0:1, CHUNK_SIZE - 1 : CHUNK_SIZE], + ) + + # ---- Compute exp(gc) and exp(g_last) as (P_MAX, 1) scalars ---- + # These (P_MAX, 1) tensors are used with tensor_scalar to broadcast + # across the free dimension without explicit (P_MAX, dim) copies. + + exp_gc_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.activation( + dst=exp_gc_p[0:P_MAX, 0:1], + op=nl.exp, + data=gc_p[0:P_MAX, 0:1], + bias=None, + scale=1.0, + ) + + # g_last: scalar, then broadcast to (P_MAX, 1) for direct + # exp(g_last - gc) in the state update. + gl_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + for i_shuf in nl.static_range(P_MAX // 32): + nisa.nc_stream_shuffle( + src=gl_11[0:1, 0:1], + dst=gl_p[i_shuf * 32 : i_shuf * 32 + 32, 0:1], + shuffle_mask=_BROADCAST_MASK, + ) + + # exp(g_last): scalar, then broadcast to (P_MAX, 1) + exp_gl_11 = nl.ndarray((1, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.activation( + dst=exp_gl_11, + op=nl.exp, + data=gl_11, + bias=None, + scale=1.0, + ) + + exp_gl_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + for i_shuf in nl.static_range(P_MAX // 32): + nisa.nc_stream_shuffle( + src=exp_gl_11[0:1, 0:1], + dst=exp_gl_p[i_shuf * 32 : i_shuf * 32 + 32, 0:1], + shuffle_mask=_BROADCAST_MASK, + ) + + # ============================================================ + # Stable pairwise decay factors from cumulative log-decay. + # + # The original fused path used split scaling: + # exp(gc[i]) * exp(-gc[j]) + # That can materialize huge unused intermediates. Build the same + # causal decay matrices as the per-chunk kernel using exp(gc[i]-gc[j]) + # and mask after the exp so upper-triangular values cannot leak into + # later matmuls. + # ============================================================ + gc_row_broadcast = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + for i_shuf in nl.static_range(P_MAX // 32): + nisa.nc_stream_shuffle( + src=gc_row[0:1, 0:P_MAX], + dst=gc_row_broadcast[i_shuf * 32 : i_shuf * 32 + 32, 0:P_MAX], + shuffle_mask=_BROADCAST_MASK, + ) + + gc_col_strict = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=gc_col_strict, + data=Lmask, + op0=nl.multiply, + operand0=gc_p, + engine=nisa.vector_engine, + ) + gc_row_strict = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=gc_row_strict, data1=gc_row_broadcast, data2=Lmask, op=nl.multiply + ) + g_diff_strict = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=g_diff_strict, + data1=gc_col_strict, + data2=gc_row_strict, + op=nl.subtract, + ) + decay_strict_raw = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.activation( + dst=decay_strict_raw, + op=nl.exp, + data=g_diff_strict, + bias=None, + scale=1.0, + ) + decay_strict = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=decay_strict, data1=decay_strict_raw, data2=Lmask, op=nl.multiply + ) + + gc_col_diag = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=gc_col_diag, + data=Lmask_d, + op0=nl.multiply, + operand0=gc_p, + engine=nisa.vector_engine, + ) + gc_row_diag = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=gc_row_diag, data1=gc_row_broadcast, data2=Lmask_d, op=nl.multiply + ) + g_diff_diag = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=g_diff_diag, + data1=gc_col_diag, + data2=gc_row_diag, + op=nl.subtract, + ) + decay_diag_raw = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.activation( + dst=decay_diag_raw, + op=nl.exp, + data=g_diff_diag, + bias=None, + scale=1.0, + ) + decay_diag = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=decay_diag, data1=decay_diag_raw, data2=Lmask_d, op=nl.multiply + ) + + # ============================================================ + # k_beta = K * beta, v_beta = V * beta + # tensor_scalar broadcasts beta_p (P_MAX, 1) across free dim + # ============================================================ + k_beta = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=k_beta, + data=k_c, + op0=nl.multiply, + operand0=beta_p, + engine=nisa.vector_engine, + ) + + v_beta = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=v_beta, + data=v_c, + op0=nl.multiply, + operand0=beta_p, + engine=nisa.vector_engine, + ) + + # ============================================================ + # Phase 1: Build A matrix (intra-chunk correction) + # Transpose K and K_beta for matmul + # ============================================================ + kb_T_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=kb_T_psum, data=k_beta) + kb_T = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=kb_T, src=kb_T_psum) + + k_T_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=k_T_psum, data=k_c) + k_T = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=k_T, src=k_T_psum) + + # QK = k_beta^T @ k (contract over features) + QK_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=QK_psum, stationary=kb_T, moving=k_T) + QK = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=QK, src=QK_psum) + + # QK_decay[i,j] = QK[i,j] * exp(gc[i] - gc[j]) for i > j. + QK_decay = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=QK_decay, data1=QK, data2=decay_strict, op=nl.multiply) + + # A = -QK_decay * lower_mask + neg_QK_decay = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=neg_QK_decay, + data=QK_decay, + op0=nl.multiply, + operand0=-1.0, + engine=nisa.vector_engine, + ) + A_mat = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=A_mat, data1=neg_QK_decay, data2=Lmask, op=nl.multiply) + + # ============================================================ + # Build the single RHS needed for v_new. + # + # Materializing N = inv(I - A) would compute: + # value_corr = N @ v_beta + # k_cumdecay = N @ (k_beta * exp(gc)) + # v_new = value_corr - k_cumdecay @ state + # + # By associativity: + # v_new = N @ (v_beta - (k_beta * exp(gc)) @ state) + # + # Solve this RHS directly. This is equivalent to the nilpotent + # Neumann series, but avoids repeated matrix squaring, which is + # numerically unstable for realistic Qwen decay gates. + # ============================================================ + kb_exp_gc = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=kb_exp_gc, + data=k_beta, + op0=nl.multiply, + operand0=exp_gc_p, + engine=nisa.vector_engine, + ) + + kbe_T_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=kbe_T_psum, data=kb_exp_gc) + kbe_T = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=kbe_T, src=kbe_T_psum) + + kbe_state_psum = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=kbe_state_psum, stationary=kbe_T, moving=state) + kbe_state = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=kbe_state, src=kbe_state_psum) + + solve_rhs = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=solve_rhs, data1=v_beta, data2=kbe_state, op=nl.subtract) + + # ============================================================ + # Direct forward substitution for: + # v_new = solve((I - A_mat), solve_rhs) + # + # A_mat is strictly lower triangular, so row i only depends on rows + # < i. The full-matmul plus row-select form keeps the shape static + # and compiler-safe while updating exactly one solved row per step. + # ============================================================ + v_new = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.memset(dst=v_new, value=0.0) + + A_T_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=A_T_psum, data=A_mat) + A_T = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=A_T, src=A_T_psum) + + for solve_i in nl.static_range(P_MAX): + row_psum = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=row_psum, stationary=A_T, moving=v_new) + row_prod = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=row_prod, src=row_psum) + + row_with_rhs = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=row_with_rhs, + data1=row_prod, + data2=solve_rhs, + op=nl.add, + ) + + row_mask = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy( + dst=row_mask[0:P_MAX, 0:1], + src=eye[0:P_MAX, solve_i : solve_i + 1], + ) + + row_update = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=row_update, + data=row_with_rhs, + op0=nl.multiply, + operand0=row_mask, + engine=nisa.vector_engine, + ) + + v_next = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=v_next, data1=v_new, data2=row_update, op=nl.add) + nisa.tensor_copy(dst=v_new, src=v_next) + + # ============================================================ + # Phase 2: Inter-chunk state propagation + # attn_intra = (q @ k^T) * decay_mask * lower_mask_diag + # ============================================================ + q_T_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=q_T_psum, data=q_c) + q_T = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=q_T, src=q_T_psum) + + qk_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=qk_psum, stationary=q_T, moving=k_T) + qk_raw = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=qk_raw, src=qk_psum) + + # qk_decay[i,j] = (q @ k^T)[i,j] * exp(gc[i] - gc[j]) for i >= j. + qk_decay = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=qk_decay, data1=qk_raw, data2=decay_diag, op=nl.multiply) + + attn_intra = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=attn_intra, data1=qk_decay, data2=Lmask_d, op=nl.multiply + ) + + # ============================================================ + # attn_inter = (q * exp(gc)) @ state (state is in SBUF!) + # ============================================================ + q_exp = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=q_exp, + data=q_c, + op0=nl.multiply, + operand0=exp_gc_p, + engine=nisa.vector_engine, + ) + + qe_T_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=qe_T_psum, data=q_exp) + qe_T = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=qe_T, src=qe_T_psum) + + ai_psum = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=ai_psum, stationary=qe_T, moving=state) + attn_inter = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=attn_inter, src=ai_psum) + + # ============================================================ + # attn_intra @ v_new + # ============================================================ + ai_T_psum = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.psum) + nisa.nc_transpose(dst=ai_T_psum, data=attn_intra) + ai_T = nl.ndarray((P_MAX, P_MAX), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=ai_T, src=ai_T_psum) + + intra_psum = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=intra_psum, stationary=ai_T, moving=v_new) + intra_out = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=intra_out, src=intra_psum) + + # ============================================================ + # chunk_output = attn_inter + intra_out + # ============================================================ + chunk_out = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor(dst=chunk_out, data1=attn_inter, data2=intra_out, op=nl.add) + + # Store output chunk to HBM + nisa.dma_copy( + dst=output[chunk_start : chunk_start + CHUNK_SIZE, 0:dim], + src=chunk_out, + ) + + # ============================================================ + # State update: state = exp(g_last) * (state + k_raw_decay^T @ v_new) + # state is updated IN-PLACE in SBUF — no HBM round-trip! + # ============================================================ + + # k_raw_decay contributes as exp(g_last) * (k * exp(-gc))^T @ v_new. + # Compute the equivalent stable form k * exp(g_last - gc) directly so + # no exp(-gc) intermediate can overflow. + gl_minus_gc_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=gl_minus_gc_p, + data1=gl_p, + data2=gc_p, + op=nl.subtract, + ) + exp_gl_minus_gc_p = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.activation( + dst=exp_gl_minus_gc_p[0:P_MAX, 0:1], + op=nl.exp, + data=gl_minus_gc_p[0:P_MAX, 0:1], + bias=None, + scale=1.0, + ) + + # k_raw_decay = k * exp(g_last - gc) + k_raw_decay = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=k_raw_decay, + data=k_c, + op0=nl.multiply, + operand0=exp_gl_minus_gc_p, + engine=nisa.vector_engine, + ) + + # k_raw_decay^T @ v_new → (dim, dim) outer product sum + # nc_matmul: result[M,N] = sum_K stationary[K,M] * moving[K,N] + # stationary=k_raw_decay (P_MAX, dim), moving=v_new (P_MAX, dim) + # Result: sum over tokens -> (dim, dim) + kv_psum = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.psum) + nisa.nc_matmul(dst=kv_psum, stationary=k_raw_decay, moving=v_new) + kv_outer = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_copy(dst=kv_outer, src=kv_psum) + + # state = state * exp(g_last) + kv_outer + # tensor_scalar broadcasts exp_gl_p (P_MAX, 1) across free dim. + state_decayed = nl.ndarray((P_MAX, dim), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=state_decayed, + data=state, + op0=nl.multiply, + operand0=exp_gl_p, + engine=nisa.vector_engine, + ) + nisa.tensor_tensor(dst=state, data1=state_decayed, data2=kv_outer, op=nl.add) + + # ---- Write final state to HBM ---- + nisa.dma_copy(dst=final_state_out, src=state) + + return output, final_state_out diff --git a/contrib/models/Qwen3.5-35B-A3B/src/nki_kernels/qwen_qk_norm_rope.py b/contrib/models/Qwen3.5-35B-A3B/src/nki_kernels/qwen_qk_norm_rope.py new file mode 100644 index 00000000..e5535254 --- /dev/null +++ b/contrib/models/Qwen3.5-35B-A3B/src/nki_kernels/qwen_qk_norm_rope.py @@ -0,0 +1,230 @@ +"""Qwen3.6-specific Q/K RMSNorm + partial-RoPE NKI kernel. + +The model's full-attention layers use head_dim=256 and partial RoPE over the +first 64 dimensions. This kernel consumes projected Q/K tensors in BSD layout +and returns normalized/rotated Q/K tensors in BHSD layout, replacing the +separate move_heads_front + q/k RMSNorm + partial RoPE torch ops. +""" + +import nki +import nki.isa as nisa +import nki.language as nl + +P_MAX = 128 +D_HEAD = 256 +ROPE_DIM = 64 +ROPE_HALF = 32 +_BROADCAST_MASK = [0] * 32 + + +def _broadcast_row_to_tile(row, out): + for i_shuf in nl.static_range(P_MAX // 32): + nisa.nc_stream_shuffle( + src=row[0:1, 0:D_HEAD], + dst=out[i_shuf * 32 : i_shuf * 32 + 32, 0:D_HEAD], + shuffle_mask=_BROADCAST_MASK, + ) + + +def _normalize_rope_store( + proj, + gamma, + cos_cache, + sin_cache, + out, + eps, +): + batch_size, seq_len, width = proj.shape + num_heads = width // D_HEAD + gamma_2d = gamma.reshape((1, D_HEAD)) + + gamma_row = nl.ndarray((1, D_HEAD), dtype=gamma.dtype, buffer=nl.sbuf) + nisa.dma_copy(dst=gamma_row, src=gamma_2d[0:1, 0:D_HEAD]) + gamma_tile = nl.ndarray((P_MAX, D_HEAD), dtype=gamma.dtype, buffer=nl.sbuf) + _broadcast_row_to_tile(gamma_row, gamma_tile) + + num_seq_tiles = (seq_len + P_MAX - 1) // P_MAX + for b_idx in nl.sequential_range(batch_size): + for h_idx in nl.sequential_range(num_heads): + col_start = h_idx * D_HEAD + for tile_idx in nl.affine_range(num_seq_tiles): + seq_start = tile_idx * P_MAX + p_size = min(P_MAX, seq_len - seq_start) + + x = nl.ndarray((P_MAX, D_HEAD), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy( + dst=x[0:p_size, 0:D_HEAD], + src=proj[ + b_idx, + seq_start : seq_start + p_size, + col_start : col_start + D_HEAD, + ], + ) + + square = nl.ndarray((P_MAX, D_HEAD), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=square[0:p_size, 0:D_HEAD], + data1=x[0:p_size, 0:D_HEAD], + data2=x[0:p_size, 0:D_HEAD], + op=nl.multiply, + ) + + sumsq = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_reduce( + dst=sumsq[0:p_size, 0:1], + data=square[0:p_size, 0:D_HEAD], + op=nl.add, + axis=1, + ) + + variance = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=variance[0:p_size, 0:1], + data=sumsq[0:p_size, 0:1], + op0=nl.multiply, + operand0=(1.0 / D_HEAD), + ) + + variance_eps = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=variance_eps[0:p_size, 0:1], + data=variance[0:p_size, 0:1], + op0=nl.add, + operand0=eps, + ) + + inv_rms = nl.ndarray((P_MAX, 1), dtype=nl.float32, buffer=nl.sbuf) + nisa.activation( + dst=inv_rms[0:p_size, 0:1], + data=variance_eps[0:p_size, 0:1], + op=nl.rsqrt, + ) + + normed = nl.ndarray((P_MAX, D_HEAD), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=normed[0:p_size, 0:D_HEAD], + data=x[0:p_size, 0:D_HEAD], + op0=nl.multiply, + operand0=inv_rms[0:p_size, 0:1], + engine=nisa.vector_engine, + ) + nisa.tensor_tensor( + dst=normed[0:p_size, 0:D_HEAD], + data1=normed[0:p_size, 0:D_HEAD], + data2=gamma_tile[0:p_size, 0:D_HEAD], + op=nl.multiply, + ) + + nisa.dma_copy( + dst=out[b_idx, h_idx, seq_start : seq_start + p_size, 0:D_HEAD], + src=normed[0:p_size, 0:D_HEAD], + ) + + cos_tile = nl.ndarray((P_MAX, ROPE_DIM), dtype=nl.float32, buffer=nl.sbuf) + sin_tile = nl.ndarray((P_MAX, ROPE_DIM), dtype=nl.float32, buffer=nl.sbuf) + nisa.dma_copy( + dst=cos_tile[0:p_size, 0:ROPE_DIM], + src=cos_cache[b_idx, seq_start : seq_start + p_size, 0:ROPE_DIM], + ) + nisa.dma_copy( + dst=sin_tile[0:p_size, 0:ROPE_DIM], + src=sin_cache[b_idx, seq_start : seq_start + p_size, 0:ROPE_DIM], + ) + + neg_hi = nl.ndarray((P_MAX, ROPE_HALF), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_scalar( + dst=neg_hi[0:p_size, 0:ROPE_HALF], + data=normed[0:p_size, ROPE_HALF:ROPE_DIM], + op0=nl.multiply, + operand0=-1.0, + ) + + lo_cos = nl.ndarray((P_MAX, ROPE_HALF), dtype=nl.float32, buffer=nl.sbuf) + hi_sin_neg = nl.ndarray((P_MAX, ROPE_HALF), dtype=nl.float32, buffer=nl.sbuf) + rope_lo = nl.ndarray((P_MAX, ROPE_HALF), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=lo_cos[0:p_size, 0:ROPE_HALF], + data1=normed[0:p_size, 0:ROPE_HALF], + data2=cos_tile[0:p_size, 0:ROPE_HALF], + op=nl.multiply, + ) + nisa.tensor_tensor( + dst=hi_sin_neg[0:p_size, 0:ROPE_HALF], + data1=neg_hi[0:p_size, 0:ROPE_HALF], + data2=sin_tile[0:p_size, 0:ROPE_HALF], + op=nl.multiply, + ) + nisa.tensor_tensor( + dst=rope_lo[0:p_size, 0:ROPE_HALF], + data1=lo_cos[0:p_size, 0:ROPE_HALF], + data2=hi_sin_neg[0:p_size, 0:ROPE_HALF], + op=nl.add, + ) + + hi_cos = nl.ndarray((P_MAX, ROPE_HALF), dtype=nl.float32, buffer=nl.sbuf) + lo_sin = nl.ndarray((P_MAX, ROPE_HALF), dtype=nl.float32, buffer=nl.sbuf) + rope_hi = nl.ndarray((P_MAX, ROPE_HALF), dtype=nl.float32, buffer=nl.sbuf) + nisa.tensor_tensor( + dst=hi_cos[0:p_size, 0:ROPE_HALF], + data1=normed[0:p_size, ROPE_HALF:ROPE_DIM], + data2=cos_tile[0:p_size, ROPE_HALF:ROPE_DIM], + op=nl.multiply, + ) + nisa.tensor_tensor( + dst=lo_sin[0:p_size, 0:ROPE_HALF], + data1=normed[0:p_size, 0:ROPE_HALF], + data2=sin_tile[0:p_size, ROPE_HALF:ROPE_DIM], + op=nl.multiply, + ) + nisa.tensor_tensor( + dst=rope_hi[0:p_size, 0:ROPE_HALF], + data1=hi_cos[0:p_size, 0:ROPE_HALF], + data2=lo_sin[0:p_size, 0:ROPE_HALF], + op=nl.add, + ) + + nisa.dma_copy( + dst=out[b_idx, h_idx, seq_start : seq_start + p_size, 0:ROPE_HALF], + src=rope_lo[0:p_size, 0:ROPE_HALF], + ) + nisa.dma_copy( + dst=out[ + b_idx, + h_idx, + seq_start : seq_start + p_size, + ROPE_HALF:ROPE_DIM, + ], + src=rope_hi[0:p_size, 0:ROPE_HALF], + ) + + +@nki.jit +def qwen_qk_norm_partial_rope_kernel( + q_proj: nl.ndarray, + k_proj: nl.ndarray, + q_gamma: nl.ndarray, + k_gamma: nl.ndarray, + cos_cache: nl.ndarray, + sin_cache: nl.ndarray, + eps: float, +): + batch_size, seq_len, q_width = q_proj.shape + _, _, k_width = k_proj.shape + q_heads = q_width // D_HEAD + k_heads = k_width // D_HEAD + + q_out = nl.ndarray( + (batch_size, q_heads, seq_len, D_HEAD), + dtype=q_proj.dtype, + buffer=nl.shared_hbm, + ) + k_out = nl.ndarray( + (batch_size, k_heads, seq_len, D_HEAD), + dtype=k_proj.dtype, + buffer=nl.shared_hbm, + ) + + _normalize_rope_store(q_proj, q_gamma, cos_cache, sin_cache, q_out, eps) + _normalize_rope_store(k_proj, k_gamma, cos_cache, sin_cache, k_out, eps) + + return q_out, k_out diff --git a/contrib/models/Qwen3.5-35B-A3B/test/__init__.py b/contrib/models/Qwen3.5-35B-A3B/test/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/contrib/models/Qwen3.5-35B-A3B/test/integration/__init__.py b/contrib/models/Qwen3.5-35B-A3B/test/integration/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/contrib/models/Qwen3.5-35B-A3B/test/integration/compile_vision_encoder.py b/contrib/models/Qwen3.5-35B-A3B/test/integration/compile_vision_encoder.py new file mode 100644 index 00000000..a66a1190 --- /dev/null +++ b/contrib/models/Qwen3.5-35B-A3B/test/integration/compile_vision_encoder.py @@ -0,0 +1,145 @@ +#!/usr/bin/env python +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +Compile the Qwen3.5 vision encoder (CPUVisionModel) to Neuron via +torch_neuronx.trace, one .pt per sequence-length bucket. The compiled +artifacts land at: + + /vision_encoder_.pt + +They can then be loaded by NeuronQwen35VisionModelWrapper.load_compiled(). + +Buckets: sequence length AFTER patch_embed + pos_embed. For Qwen3.5: + spatial_merge_size = 2, temporal_patch_size = 2, patch_size = 16 + H x W image with `image_grid_thw = [1, H//16, W//16]` produces + H//16 * W//16 patch tokens BEFORE merger. After merger there are + (H//16 // 2) * (W//16 // 2) merged tokens. + +Common image sizes (square): + 512x512 → 1024 patch tokens → 256 merged tokens + 1024x1024→ 4096 patch tokens → 1024 merged tokens + 2048x2048→16384 patch tokens → 4096 merged tokens + +We compile at the *patch-token* seq_len (i.e., the input to the ViT blocks). + +Usage: + source /opt/aws_neuronx_venv_pytorch_2_9_nxd_inference/bin/activate + python contrib/models/Qwen3.5-35B-A3B/test/integration/compile_vision_encoder.py \\ + --model-path /mnt/nvme/models/Qwen3.5-35B-A3B \\ + --out-dir /tmp/qwen35_35b_a3b_vl_bench/vision \\ + --buckets 1024 4096 16384 +""" + +import argparse +import gc +import json +import os +import sys +import time + +import torch + +_HERE = os.path.dirname(os.path.abspath(__file__)) +_CONTRIB_ROOT = os.path.abspath(os.path.join(_HERE, "..", "..")) +if _CONTRIB_ROOT not in sys.path: + sys.path.insert(0, _CONTRIB_ROOT) + + +def build_cpu_vision(model_path): + """Load CPUVisionModel with real weights from HF safetensors.""" + from src.modeling_qwen35_vision import CPUVisionModel + from types import SimpleNamespace + + with open(os.path.join(model_path, "config.json")) as f: + full = json.load(f) + vc = full["vision_config"] + vconf = SimpleNamespace( + depth=vc["depth"], + hidden_size=vc["hidden_size"], + num_heads=vc["num_heads"], + intermediate_size=vc["intermediate_size"], + patch_size=vc["patch_size"], + temporal_patch_size=vc.get("temporal_patch_size", 2), + spatial_merge_size=vc.get("spatial_merge_size", 2), + out_hidden_size=vc["out_hidden_size"], + num_position_embeddings=vc["num_position_embeddings"], + in_channels=vc.get("in_channels", 3), + ) + + # Reuse the wrapper's weight loader path + from src.modeling_qwen35_vision import NeuronQwen35VisionModelWrapper + w = NeuronQwen35VisionModelWrapper(config=vconf, model_cls=None, vision_seq_len_buckets=[16384]) + w.load_cpu_model(model_path) # populates w._cpu_model with a CPUVisionModel bf16 + return w._cpu_model, vconf + + +def compile_bucket(cpu_model, vconf, bucket_len, out_path): + """torch_neuronx.trace the CPUVisionModel at a specific patch-token seq_len.""" + import torch_neuronx + + dtype = torch.bfloat16 + hidden = vconf.hidden_size + num_heads = vconf.num_heads + head_dim = hidden // num_heads + + # Example inputs matching the CPUVisionModel.forward signature: + # (hidden_states, attention_mask, cos, sin) + hidden_states = torch.zeros((bucket_len, hidden), dtype=dtype) + attention_mask = torch.zeros((1, 1, bucket_len, bucket_len), dtype=dtype) + # cos/sin shape: (seq_len, head_dim//2 doubled to head_dim) — see + # CPUVisionModel._forward_attention. Actually we pass a shape matching + # what wrapper produces: (seq_len, head_dim) each, where head_dim entries + # are rope-emb-cos/sin values (padded via `torch.cat((rot,rot),-1)` in wrapper). + cos = torch.zeros((bucket_len, head_dim), dtype=dtype) + sin = torch.zeros((bucket_len, head_dim), dtype=dtype) + + print(f"[compile] bucket={bucket_len}, tracing...") + t0 = time.perf_counter() + traced = torch_neuronx.trace( + cpu_model, + (hidden_states, attention_mask, cos, sin), + compiler_workdir=f"/tmp/nxd_vision_ws_{bucket_len}", + compiler_args=[ + "--model-type=transformer", + "--auto-cast=none", + "-O1", + "--enable-mixed-precision-accumulation", + ], + ) + dt = time.perf_counter() - t0 + print(f"[compile] bucket={bucket_len}: done in {dt:.1f}s → {out_path}") + torch.jit.save(traced, out_path) + return dt + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--model-path", default="/mnt/nvme/models/Qwen3.5-35B-A3B") + ap.add_argument("--out-dir", default="/tmp/qwen35_35b_a3b_vl_bench/vision") + ap.add_argument("--buckets", nargs="+", type=int, default=[1024, 4096, 16384], + help="Patch-token seq lengths (before merger). " + "1024→512x512, 4096→1024x1024, 16384→2048x2048.") + ap.add_argument("--overwrite", action="store_true") + args = ap.parse_args() + + os.makedirs(args.out_dir, exist_ok=True) + cpu_model, vconf = build_cpu_vision(args.model_path) + cpu_model.eval() + print(f"[build] loaded vision encoder: depth={vconf.depth} " + f"hidden={vconf.hidden_size} num_heads={vconf.num_heads}") + + for bucket in args.buckets: + out_path = os.path.join(args.out_dir, f"vision_encoder_{bucket}.pt") + if os.path.exists(out_path) and not args.overwrite: + print(f"[skip] {out_path} exists (use --overwrite to force)") + continue + compile_bucket(cpu_model, vconf, bucket, out_path) + gc.collect() + + print("[done] all buckets compiled") + + +if __name__ == "__main__": + main() diff --git a/contrib/models/Qwen3.5-35B-A3B/test/integration/run_accuracy_check.py b/contrib/models/Qwen3.5-35B-A3B/test/integration/run_accuracy_check.py new file mode 100644 index 00000000..d3af9d04 --- /dev/null +++ b/contrib/models/Qwen3.5-35B-A3B/test/integration/run_accuracy_check.py @@ -0,0 +1,183 @@ +#!/usr/bin/env python +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +Cross-check Qwen3.5-35B-A3B accuracy on Neuron against a HuggingFace CPU reference. + +Runs a batch of prompts through both: + (a) The compiled Neuron model (loaded from --compiled-path) + (b) The HuggingFace text decoder on CPU (bf16, greedy) + +For each prompt, both are run greedy (top_k=1 / do_sample=False) for +--max-new-tokens tokens. Reports per-prompt token match rate and overall +match rate. + +Usage: + source /opt/aws_neuronx_venv_pytorch_2_9_nxd_inference/bin/activate + python contrib/models/Qwen3.5-35B-A3B/test/integration/run_accuracy_check.py \\ + --compiled-path /tmp/qwen35_35b_a3b_traced \\ + --max-new-tokens 16 +""" + +import argparse +import gc +import json +import os +import sys +import time + +import torch + +_HERE = os.path.dirname(os.path.abspath(__file__)) +_CONTRIB_ROOT = os.path.abspath(os.path.join(_HERE, "..", "..")) +if _CONTRIB_ROOT not in sys.path: + sys.path.insert(0, _CONTRIB_ROOT) + + +DEFAULT_PROMPTS = [ + "The capital of France is", + "The largest planet in our solar system is", + "Water boils at", + "A haiku about autumn leaves:", + "In one sentence, explain photosynthesis.", +] + + +def load_neuron(compiled_path: str): + from src.modeling_qwen35 import NeuronQwen35ForCausalLM + m = NeuronQwen35ForCausalLM(compiled_path) + m.load(compiled_path) + return m + + +def load_hf(model_path: str): + from transformers import AutoConfig, AutoModelForCausalLM + # transformers 4.57.6 may not have Qwen3_5ForConditionalGeneration + # registered by config's model_type. Try trust_remote_code. + cfg = AutoConfig.from_pretrained(model_path, trust_remote_code=True) + try: + # Load the text half only — vision would blow up CPU memory. + model = AutoModelForCausalLM.from_pretrained( + model_path, + torch_dtype=torch.bfloat16, + trust_remote_code=True, + ) + except Exception: + # Fall back to the text sub-config + from transformers import AutoModel + text_cfg = getattr(cfg, "text_config", cfg) + model = AutoModelForCausalLM.from_config(text_cfg, torch_dtype=torch.bfloat16) + raise RuntimeError( + "HF Qwen3.5-35B-A3B did not load via AutoModelForCausalLM; " + "cross-check on CPU is not supported for this model_type" + ) + model.eval() + return model + + +def generate_neuron(nxd_model, tok, prompt: str, max_new: int): + from transformers import GenerationConfig + import transformers + from neuronx_distributed_inference.utils.hf_adapter import ( + HuggingFaceGenerationAdapter, + ) + gen_cfg = GenerationConfig( + do_sample=True, top_k=1, + pad_token_id=tok.pad_token_id, eos_token_id=tok.eos_token_id, + ) + gen_cfg.transformers_version = transformers.__version__ + gen_model = HuggingFaceGenerationAdapter(nxd_model) + gen_model.generation_config.transformers_version = transformers.__version__ + + enc = tok(prompt, return_tensors="pt") + out = gen_model.generate( + enc.input_ids, generation_config=gen_cfg, max_new_tokens=max_new, + ) + new_ids = out[0].tolist()[enc.input_ids.shape[-1]:] + return new_ids, tok.decode(out[0], skip_special_tokens=True) + + +def generate_hf(hf_model, tok, prompt: str, max_new: int): + enc = tok(prompt, return_tensors="pt") + with torch.no_grad(): + out = hf_model.generate( + enc.input_ids, + do_sample=False, + max_new_tokens=max_new, + pad_token_id=tok.pad_token_id, + eos_token_id=tok.eos_token_id, + ) + new_ids = out[0].tolist()[enc.input_ids.shape[-1]:] + return new_ids, tok.decode(out[0], skip_special_tokens=True) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--model-path", default="/mnt/nvme/models/Qwen3.5-35B-A3B") + ap.add_argument("--compiled-path", default="/tmp/qwen35_35b_a3b_traced") + ap.add_argument("--max-new-tokens", type=int, default=16) + ap.add_argument("--prompts", nargs="*", default=None) + ap.add_argument("--skip-hf", action="store_true", help="Only run Neuron") + ap.add_argument("--out-json", default=None) + args = ap.parse_args() + + prompts = args.prompts if args.prompts else DEFAULT_PROMPTS + + from transformers import AutoTokenizer + tok = AutoTokenizer.from_pretrained(args.model_path, padding_side="right") + if tok.pad_token is None: + tok.pad_token = tok.eos_token + + print(f"[accuracy] loading Neuron model from {args.compiled_path}") + nxd = load_neuron(args.compiled_path) + + hf_model = None + if not args.skip_hf: + print(f"[accuracy] loading HF reference (CPU, bf16) from {args.model_path}") + try: + hf_model = load_hf(args.model_path) + except Exception as e: + print(f"[accuracy] HF load failed ({e}); running Neuron-only") + hf_model = None + + results = [] + total_match = 0 + total_tokens = 0 + + for i, prompt in enumerate(prompts): + print(f"\n=== prompt #{i}: {prompt!r} ===") + nxd_ids, nxd_txt = generate_neuron(nxd, tok, prompt, args.max_new_tokens) + entry = {"prompt": prompt, "neuron_tokens": nxd_ids, "neuron_text": nxd_txt} + print(f"[neuron] ids={nxd_ids}") + print(f"[neuron] text={nxd_txt!r}") + + if hf_model is not None: + hf_ids, hf_txt = generate_hf(hf_model, tok, prompt, args.max_new_tokens) + entry["hf_tokens"] = hf_ids + entry["hf_text"] = hf_txt + L = min(len(nxd_ids), len(hf_ids)) + match = sum(1 for a, b in zip(nxd_ids[:L], hf_ids[:L]) if a == b) + entry["match"] = match + entry["match_denom"] = L + total_match += match + total_tokens += L + print(f"[hf] ids={hf_ids}") + print(f"[hf] text={hf_txt!r}") + print(f"[match] {match}/{L}") + results.append(entry) + + if hf_model is not None: + rate = total_match / max(1, total_tokens) + print(f"\n=== TOTAL: {total_match}/{total_tokens} tokens match ({rate*100:.1f}%) ===") + + if args.out_json: + with open(args.out_json, "w") as f: + json.dump({"results": results, + "total_match": total_match, + "total_tokens": total_tokens}, f, indent=2) + print(f"[accuracy] wrote {args.out_json}") + + +if __name__ == "__main__": + main() diff --git a/contrib/models/Qwen3.5-35B-A3B/test/integration/run_benchmark.py b/contrib/models/Qwen3.5-35B-A3B/test/integration/run_benchmark.py new file mode 100644 index 00000000..43abb7de --- /dev/null +++ b/contrib/models/Qwen3.5-35B-A3B/test/integration/run_benchmark.py @@ -0,0 +1,151 @@ +#!/usr/bin/env python +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +TTFT / TPOT benchmark for Qwen3.5-35B-A3B on Neuron. + +For each prompt length, we run: + * a 1-token generation (measures TTFT = prefill + first-decode) + * an N-token generation (extracts TPOT = (elapsed - TTFT) / (N-1)) + +Results are averaged across --repeats runs; the first run is treated as warmup +and discarded. + +Usage: + source /opt/aws_neuronx_venv_pytorch_2_9_nxd_inference/bin/activate + python contrib/models/Qwen3.5-35B-A3B/test/integration/run_benchmark.py \\ + --compiled-path /tmp/qwen35_35b_a3b_traced \\ + --seq-len 512 --max-new-tokens 64 --repeats 5 +""" + +import argparse +import json +import os +import statistics +import sys +import time + +import torch + +_HERE = os.path.dirname(os.path.abspath(__file__)) +_CONTRIB_ROOT = os.path.abspath(os.path.join(_HERE, "..", "..")) +if _CONTRIB_ROOT not in sys.path: + sys.path.insert(0, _CONTRIB_ROOT) + + +def build_prompt_at_length(tok, target_len: int) -> str: + """Build a prompt whose tokenization length equals target_len (approx).""" + base = ("Once upon a time, there lived a curious explorer who traveled " + "across mountains and seas. ") + # Grow until at least target_len tokens + prompt = base + while True: + ids = tok(prompt, return_tensors="pt").input_ids[0] + if len(ids) >= target_len: + break + prompt += base + ids = tok(prompt, return_tensors="pt").input_ids[0][:target_len] + return tok.decode(ids, skip_special_tokens=True) + + +def bench_once(gen_model, tok, prompt: str, gen_cfg, max_new: int): + enc = tok(prompt, return_tensors="pt") + # TTFT + t0 = time.perf_counter() + _ = gen_model.generate(enc.input_ids, generation_config=gen_cfg, max_new_tokens=1) + ttft = (time.perf_counter() - t0) * 1000.0 + + # Full run + t1 = time.perf_counter() + out = gen_model.generate( + enc.input_ids, generation_config=gen_cfg, max_new_tokens=max_new, + ) + total = (time.perf_counter() - t1) * 1000.0 + n_new = out.shape[-1] - enc.input_ids.shape[-1] + n_decode = max(1, n_new - 1) + tpot = (total - ttft) / n_decode + return ttft, tpot, n_new + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--model-path", default="/mnt/nvme/models/Qwen3.5-35B-A3B") + ap.add_argument("--compiled-path", default="/tmp/qwen35_35b_a3b_traced") + ap.add_argument("--prompt-lens", type=int, nargs="+", default=[16, 64, 256]) + ap.add_argument("--max-new-tokens", type=int, default=64) + ap.add_argument("--repeats", type=int, default=5) + ap.add_argument("--out-json", default=None) + args = ap.parse_args() + + from transformers import AutoTokenizer, GenerationConfig + import transformers + from neuronx_distributed_inference.utils.hf_adapter import ( + HuggingFaceGenerationAdapter, + ) + from src.modeling_qwen35 import NeuronQwen35ForCausalLM + + print(f"[bench] loading from {args.compiled_path}") + m = NeuronQwen35ForCausalLM(args.compiled_path) + m.load(args.compiled_path) + + tok = AutoTokenizer.from_pretrained(args.model_path, padding_side="right") + if tok.pad_token is None: + tok.pad_token = tok.eos_token + + gen_cfg = GenerationConfig( + do_sample=True, top_k=1, + pad_token_id=tok.pad_token_id, eos_token_id=tok.eos_token_id, + ) + gen_cfg.transformers_version = transformers.__version__ + gen_model = HuggingFaceGenerationAdapter(m) + gen_model.generation_config.transformers_version = transformers.__version__ + + results = [] + for plen in args.prompt_lens: + prompt = build_prompt_at_length(tok, plen) + actual_plen = len(tok(prompt, return_tensors="pt").input_ids[0]) + + ttfts = [] + tpots = [] + for r in range(args.repeats + 1): # +1 warmup + ttft, tpot, n_new = bench_once( + gen_model, tok, prompt, gen_cfg, args.max_new_tokens + ) + if r == 0: + print(f"[bench] len={actual_plen} warmup: TTFT={ttft:.1f} ms TPOT={tpot:.2f} ms (discarded)") + continue + ttfts.append(ttft) + tpots.append(tpot) + print(f"[bench] len={actual_plen} r={r}: TTFT={ttft:.1f} ms TPOT={tpot:.2f} ms n_new={n_new}") + + r = { + "prompt_len": actual_plen, + "target_prompt_len": plen, + "ttft_ms_mean": statistics.mean(ttfts), + "ttft_ms_median": statistics.median(ttfts), + "ttft_ms_stdev": statistics.pstdev(ttfts) if len(ttfts) > 1 else 0.0, + "tpot_ms_mean": statistics.mean(tpots), + "tpot_ms_median": statistics.median(tpots), + "tpot_ms_stdev": statistics.pstdev(tpots) if len(tpots) > 1 else 0.0, + "throughput_tok_per_s_mean": 1000.0 / statistics.mean(tpots), + "max_new_tokens": args.max_new_tokens, + "repeats": args.repeats, + } + results.append(r) + + print("\n=== SUMMARY ===") + for r in results: + print(f" prompt_len={r['prompt_len']:4d}" + f" TTFT={r['ttft_ms_median']:7.1f} ms" + f" TPOT={r['tpot_ms_median']:6.2f} ms" + f" ({r['throughput_tok_per_s_mean']:6.1f} tok/s)") + + if args.out_json: + with open(args.out_json, "w") as f: + json.dump(results, f, indent=2) + print(f"[bench] wrote {args.out_json}") + + +if __name__ == "__main__": + main() diff --git a/contrib/models/Qwen3.5-35B-A3B/test/integration/run_hf_reference.py b/contrib/models/Qwen3.5-35B-A3B/test/integration/run_hf_reference.py new file mode 100644 index 00000000..37547805 --- /dev/null +++ b/contrib/models/Qwen3.5-35B-A3B/test/integration/run_hf_reference.py @@ -0,0 +1,87 @@ +#!/usr/bin/env python +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +Run the HuggingFace Qwen3.5-35B-A3B reference on CPU and dump greedy outputs to JSON. + +Because transformers 4.57.6 (shipped with the current NxDI DLAMI) predates Qwen3.5 +support, this must be run in an isolated venv with transformers>=5.13: + + python3 -m venv /tmp/hf_ref_venv + /tmp/hf_ref_venv/bin/pip install "transformers==5.13.0" "torch>=2.6" \\ + "safetensors" "sentencepiece" "accelerate" + /tmp/hf_ref_venv/bin/python \\ + contrib/models/Qwen3.5-35B-A3B/test/integration/run_hf_reference.py \\ + --model-path /mnt/nvme/models/Qwen3.5-35B-A3B \\ + --max-new-tokens 16 \\ + --out-json /tmp/qwen35_35b_a3b_hf_reference.json + +The output JSON can then be fed to compare_accuracy.py alongside a Neuron JSON. +""" + +import argparse +import json + +import torch + + +DEFAULT_PROMPTS = [ + "The capital of France is", + "The largest planet in our solar system is", + "Water boils at", + "A haiku about autumn leaves:", + "In one sentence, explain photosynthesis.", +] + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--model-path", default="/mnt/nvme/models/Qwen3.5-35B-A3B") + ap.add_argument("--max-new-tokens", type=int, default=16) + ap.add_argument("--prompts", nargs="*", default=None) + ap.add_argument("--out-json", required=True) + args = ap.parse_args() + + from transformers import AutoTokenizer, AutoModelForCausalLM + + tok = AutoTokenizer.from_pretrained(args.model_path) + if tok.pad_token is None: + tok.pad_token = tok.eos_token + + print(f"[hf] loading {args.model_path} on CPU (bf16)") + # For a VL model, AutoModelForCausalLM may or may not work. Try it first. + model = AutoModelForCausalLM.from_pretrained( + args.model_path, + torch_dtype=torch.bfloat16, + ) + model.eval() + + prompts = args.prompts if args.prompts else DEFAULT_PROMPTS + results = [] + for i, prompt in enumerate(prompts): + enc = tok(prompt, return_tensors="pt") + with torch.no_grad(): + out = model.generate( + enc.input_ids, + do_sample=False, + max_new_tokens=args.max_new_tokens, + pad_token_id=tok.pad_token_id, + eos_token_id=tok.eos_token_id, + ) + new_ids = out[0].tolist()[enc.input_ids.shape[-1]:] + text = tok.decode(out[0], skip_special_tokens=True) + results.append({ + "prompt": prompt, + "hf_new_tokens": new_ids, + "hf_text": text, + }) + print(f"[hf] #{i}: {prompt!r} -> {new_ids}") + + with open(args.out_json, "w") as f: + json.dump({"results": results, "max_new_tokens": args.max_new_tokens}, f, indent=2) + print(f"[hf] wrote {args.out_json}") + + +if __name__ == "__main__": + main() diff --git a/contrib/models/Qwen3.5-35B-A3B/test/integration/run_text_smoke.py b/contrib/models/Qwen3.5-35B-A3B/test/integration/run_text_smoke.py new file mode 100644 index 00000000..b80c3c9c --- /dev/null +++ b/contrib/models/Qwen3.5-35B-A3B/test/integration/run_text_smoke.py @@ -0,0 +1,181 @@ +#!/usr/bin/env python +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +Standalone text-only smoke test for Qwen3.5-35B-A3B on Neuron. + +Compiles the text backbone, runs a short generation, and prints TTFT + TPOT. +Use this before running the pytest suite to catch compile errors interactively. + +Usage: + source /opt/aws_neuronx_venv_pytorch_2_9_nxd_inference/bin/activate + python contrib/models/Qwen3.5-35B-A3B/test/integration/run_text_smoke.py \\ + --model-path /mnt/nvme/models/Qwen3.5-35B-A3B \\ + --tp 8 --seq-len 512 --max-new-tokens 32 +""" + +import argparse +import gc +import json +import os +import sys +import time + +import torch + +_HERE = os.path.dirname(os.path.abspath(__file__)) +_CONTRIB_ROOT = os.path.abspath(os.path.join(_HERE, "..", "..")) +if _CONTRIB_ROOT not in sys.path: + sys.path.insert(0, _CONTRIB_ROOT) + + +def build_config(model_path: str, tp: int, seq_len: int): + from neuronx_distributed_inference.models.config import ( + MoENeuronConfig, + OnDeviceSamplingConfig, + ) + from src.modeling_qwen35 import Qwen35InferenceConfig + + # 35B-A3B is a MoE variant (qwen3_5_moe_text). Use MoENeuronConfig so + # initialize_moe_module() finds the router / blockwise-matmul / moe_tp_degree + # settings. Force use_torch_block_wise=True because the shipped + # neuronx-cc 2.26.6360 DLAMI does not ship the shard-hidden LNC=2 + # blockwise-matmul NKI kernel. + neuron_config = MoENeuronConfig( + tp_degree=tp, + batch_size=1, + ctx_batch_size=1, + tkg_batch_size=1, + seq_len=seq_len, + torch_dtype=torch.bfloat16, + on_device_sampling_config=OnDeviceSamplingConfig(top_k=1), + enable_bucketing=False, + flash_decoding_enabled=False, + logical_nc_config=2, + save_sharded_checkpoint=True, + moe_tp_degree=tp, + moe_ep_degree=1, + normalize_top_k_affinities=True, + # Backend selection (empirical, TP=8, prompt=16, batch=1): + # use_shard_on_intermediate_dynamic_while → 480 ms TTFT (best short) + # patched _call_shard_hidden_kernel → 567 ms TTFT (flat vs prompt) + # use_torch_block_wise → 554 ms TTFT (fallback) + # shard-on-intermediate wins at short prompts (chat), the patched + # shard-hidden path is closer for long prompts (>256). Neither is a + # true NKI shard-hidden kernel — that path in the shipped DLAMI + # (neuronx-cc 2.26.6360) is a NotImplementedError stub. + # modeling_qwen35.py monkey-patches + # `_call_shard_hidden_kernel` to wire in nkilib's forward kernel so + # the default path is at least runnable; we still explicitly pick + # shard-on-intermediate here for best measured TTFT. + blockwise_matmul_config={"use_shard_on_intermediate_dynamic_while": True}, + ) + # Qwen3.5-MoE uses softmax over router logits. + if hasattr(neuron_config, "router_config"): + neuron_config.router_config.dtype = torch.float32 + neuron_config.router_config.act_fn = "softmax" + + with open(os.path.join(model_path, "config.json")) as f: + full = json.load(f) + text_cfg = full.get("text_config", full) + + cfg = dict(text_cfg) + cfg["pad_token_id"] = text_cfg.get("eos_token_id", 248044) + if "rope_parameters" in text_cfg: + rp = text_cfg["rope_parameters"] + cfg["rope_theta"] = rp.get("rope_theta", 10000000) + cfg["partial_rotary_factor"] = rp.get("partial_rotary_factor", 0.25) + cfg["mrope_section"] = rp.get("mrope_section", [11, 11, 10]) + cfg.setdefault("tie_word_embeddings", text_cfg.get("tie_word_embeddings", True)) + + return Qwen35InferenceConfig( + neuron_config=neuron_config, + use_hybrid_cache_manager=False, + **cfg, + ) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--model-path", default="/mnt/nvme/models/Qwen3.5-35B-A3B") + ap.add_argument("--compiled-path", default="/tmp/qwen35_35b_a3b_traced") + ap.add_argument("--tp", type=int, default=8) + ap.add_argument("--seq-len", type=int, default=512) + ap.add_argument("--max-new-tokens", type=int, default=32) + ap.add_argument("--prompt", default="The capital of France is") + ap.add_argument("--skip-compile", action="store_true", + help="Assume artifacts already exist at compiled-path") + args = ap.parse_args() + + from src.modeling_qwen35 import NeuronQwen35ForCausalLM + from transformers import AutoTokenizer, GenerationConfig + import transformers + from neuronx_distributed_inference.utils.hf_adapter import ( + HuggingFaceGenerationAdapter, + ) + + inf_config = build_config(args.model_path, args.tp, args.seq_len) + + neff = os.path.join(args.compiled_path, "model.pt") + if not args.skip_compile and not os.path.exists(neff): + print(f"[compile] → {args.compiled_path}") + m = NeuronQwen35ForCausalLM(args.model_path, inf_config) + t0 = time.perf_counter() + m.compile(args.compiled_path) + print(f"[compile] done in {(time.perf_counter()-t0):.1f} s") + del m + gc.collect() + + print(f"[load] ← {args.compiled_path}") + m = NeuronQwen35ForCausalLM(args.compiled_path) + m.load(args.compiled_path) + + tok = AutoTokenizer.from_pretrained(args.model_path, padding_side="right") + if tok.pad_token is None: + tok.pad_token = tok.eos_token + + gen_cfg = GenerationConfig( + do_sample=True, top_k=1, + pad_token_id=tok.pad_token_id, eos_token_id=tok.eos_token_id, + ) + gen_cfg.transformers_version = transformers.__version__ + gen_model = HuggingFaceGenerationAdapter(m) + gen_model.generation_config.transformers_version = transformers.__version__ + + enc = tok(args.prompt, return_tensors="pt") + + # Warmup (first call may include tracing overhead) + _ = gen_model.generate(enc.input_ids, generation_config=gen_cfg, max_new_tokens=1) + + # TTFT + t0 = time.perf_counter() + _ = gen_model.generate(enc.input_ids, generation_config=gen_cfg, max_new_tokens=1) + ttft = (time.perf_counter() - t0) * 1000 + + # Full run + t1 = time.perf_counter() + out = gen_model.generate( + enc.input_ids, + generation_config=gen_cfg, + max_new_tokens=args.max_new_tokens, + ) + total_ms = (time.perf_counter() - t1) * 1000 + + n_new = out.shape[-1] - enc.input_ids.shape[-1] + n_decode = max(1, n_new - 1) + tpot = (total_ms - ttft) / n_decode + + text = tok.decode(out[0], skip_special_tokens=True) + + print("=" * 72) + print(f"prompt : {args.prompt!r}") + print(f"output : {text!r}") + print(f"n_new : {n_new}") + print(f"TTFT : {ttft:.1f} ms") + print(f"TPOT : {tpot:.1f} ms ({1000.0/max(tpot,1e-6):.2f} tok/s)") + print("=" * 72) + + +if __name__ == "__main__": + main() diff --git a/contrib/models/Qwen3.5-35B-A3B/test/integration/run_vl_benchmark.py b/contrib/models/Qwen3.5-35B-A3B/test/integration/run_vl_benchmark.py new file mode 100644 index 00000000..ecf39e7d --- /dev/null +++ b/contrib/models/Qwen3.5-35B-A3B/test/integration/run_vl_benchmark.py @@ -0,0 +1,271 @@ +#!/usr/bin/env python +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +Vision-Language benchmark: TTFT / TPOT / output vs image size. + +Compiles ONCE with NxDI CTE bucketing enabled so the same model handles +512×512, 1024×1024, 2048×2048 images (input_ids ~280 / ~1048 / ~4120, +CTE buckets [512, 1024, 2048, 4096, 8192]). + +Usage: + source /opt/aws_neuronx_venv_pytorch_2_9_nxd_inference/bin/activate + python contrib/models/Qwen3.5-35B-A3B/test/integration/run_vl_benchmark.py \\ + --model-path /mnt/nvme/models/Qwen3.5-35B-A3B \\ + --compiled-path /tmp/qwen35_35b_a3b_vl_bench \\ + --images 512 1024 2048 \\ + --max-new-tokens 48 --repeats 3 +""" + +import argparse +import gc +import json +import os +import statistics +import sys +import time + +os.environ.setdefault("QWEN36_DELTANET_CTE_IMPL", "legacy_direct") +os.environ.setdefault("QWEN36_DELTANET_MULTIHEAD_CTE", "0") + +import torch + +_HERE = os.path.dirname(os.path.abspath(__file__)) +_CONTRIB_ROOT = os.path.abspath(os.path.join(_HERE, "..", "..")) +if _CONTRIB_ROOT not in sys.path: + sys.path.insert(0, _CONTRIB_ROOT) + + +DEFAULT_PROMPT = "What is in this image? Describe it briefly." + + +def build_vl_config(model_path: str, tp: int, max_seq_len: int, buckets): + # MoE variant needs MoENeuronConfig (see run_text_smoke.py comment block). + # The `_call_shard_hidden_kernel` stub in NxD is monkey-patched at import + # time by modeling_qwen35.py, so shard-on-intermediate remains the best + # explicit choice for prefill. + from neuronx_distributed_inference.models.config import ( + MoENeuronConfig, + OnDeviceSamplingConfig, + ) + from src.modeling_qwen35 import Qwen35InferenceConfig + from src.modeling_qwen35_vl import Qwen35VLInferenceConfig + + neuron_config = MoENeuronConfig( + tp_degree=tp, + batch_size=1, + ctx_batch_size=1, + tkg_batch_size=1, + seq_len=max_seq_len, + torch_dtype=torch.bfloat16, + on_device_sampling_config=OnDeviceSamplingConfig(top_k=1), + enable_bucketing=True, + context_encoding_buckets=buckets, + token_generation_buckets=[max_seq_len], + flash_decoding_enabled=False, + logical_nc_config=2, + save_sharded_checkpoint=True, + moe_tp_degree=tp, + moe_ep_degree=1, + normalize_top_k_affinities=True, + blockwise_matmul_config={"use_shard_on_intermediate_dynamic_while": True}, + ) + if hasattr(neuron_config, "router_config"): + neuron_config.router_config.dtype = torch.float32 + neuron_config.router_config.act_fn = "softmax" + + with open(os.path.join(model_path, "config.json")) as f: + full = json.load(f) + text_cfg = full.get("text_config", full) + cfg = dict(text_cfg) + cfg["pad_token_id"] = text_cfg.get("eos_token_id", 248044) + if "rope_parameters" in text_cfg: + rp = text_cfg["rope_parameters"] + cfg["rope_theta"] = rp.get("rope_theta", 10000000) + cfg["partial_rotary_factor"] = rp.get("partial_rotary_factor", 0.25) + cfg["mrope_section"] = rp.get("mrope_section", [11, 11, 10]) + cfg.setdefault("tie_word_embeddings", text_cfg.get("tie_word_embeddings", True)) + + text_config = Qwen35InferenceConfig( + neuron_config=neuron_config, + use_hybrid_cache_manager=False, + use_text_only_cte_inputs=False, + **cfg, + ) + + vision_config_dict = full["vision_config"] + vision_config_dict.setdefault("spatial_merge_size", 2) + vision_config_dict.setdefault("temporal_patch_size", 2) + + vl_config = Qwen35VLInferenceConfig( + text_config=text_config, + vision_config=vision_config_dict, + image_token_id=full.get("image_token_id", 248056), + video_token_id=full.get("video_token_id", 248057), + vision_start_token_id=full.get("vision_start_token_id", 248053), + vision_end_token_id=full.get("vision_end_token_id", 248054), + spatial_merge_size=vision_config_dict["spatial_merge_size"], + ) + return text_config, vl_config + + +def compile_and_load(model_path, compiled_path, text_config, vl_config, + skip_compile, vision_compiled_dir=None): + from src.modeling_qwen35 import NeuronQwen35ForCausalLM + from src.modeling_qwen35_vl import NeuronQwen35VLForCausalLM + + text_path = os.path.join(compiled_path, "text_model") + neff = os.path.join(text_path, "model.pt") + if not skip_compile and not os.path.exists(neff): + print(f"[compile-text] → {text_path}") + os.makedirs(text_path, exist_ok=True) + m = NeuronQwen35ForCausalLM(model_path, text_config) + t0 = time.perf_counter() + m.compile(text_path) + print(f"[compile-text] done in {(time.perf_counter()-t0):.1f} s") + del m + gc.collect() + + vl_model = NeuronQwen35VLForCausalLM( + model_path=model_path, + text_config=text_config, + vision_config=vl_config, + ) + vl_model.text_model.load(text_path) + + if vision_compiled_dir and os.path.isdir(vision_compiled_dir): + # Neuron-compiled vision encoder (multi-bucket .pt files) + # Also load CPU model as a fallback for seq_lens exceeding compiled buckets. + print(f"[load-vision] loading Neuron vision encoder from {vision_compiled_dir}") + vl_model.vision_model_wrapper.load_compiled(vision_compiled_dir) + vl_model.vision_model_wrapper.load_vision_weights_from_hf(model_path) + print("[load-vision] loading CPU vision encoder as oversize fallback") + vl_model.vision_model_wrapper.load_cpu_model(model_path) + else: + print("[load-vision] loading CPU vision encoder weights") + vl_model.vision_model_wrapper.load_cpu_model(model_path) + vl_model.vision_model_wrapper.load_vision_weights_from_hf(model_path) + return vl_model + + +def one_image_run(vl_model, processor, tok, image, prompt, max_new_tokens): + """Return (ttft_ms, tpot_ms, n_new, text).""" + msgs = [{"role":"user","content":[{"type":"image","image":image},{"type":"text","text":prompt}]}] + inp = processor.apply_chat_template(msgs, tokenize=True, add_generation_prompt=True, return_tensors="pt", return_dict=True) + + # TTFT: 1-token generate + t0 = time.perf_counter() + _ = vl_model.generate( + input_ids=inp["input_ids"], + attention_mask=inp.get("attention_mask", torch.ones_like(inp["input_ids"])), + pixel_values=inp["pixel_values"], + image_grid_thw=inp["image_grid_thw"], + max_new_tokens=1, + temperature=0.0, + ) + ttft_ms = (time.perf_counter() - t0) * 1000.0 + + # Full generate + t1 = time.perf_counter() + out = vl_model.generate( + input_ids=inp["input_ids"], + attention_mask=inp.get("attention_mask", torch.ones_like(inp["input_ids"])), + pixel_values=inp["pixel_values"], + image_grid_thw=inp["image_grid_thw"], + max_new_tokens=max_new_tokens, + temperature=0.0, + ) + total_ms = (time.perf_counter() - t1) * 1000.0 + + new_ids = out[0].tolist()[inp["input_ids"].shape[1]:] + n_new = len(new_ids) + tpot_ms = (total_ms - ttft_ms) / max(1, n_new - 1) + text = tok.decode(new_ids, skip_special_tokens=True) + return ttft_ms, tpot_ms, n_new, text + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--model-path", default="/mnt/nvme/models/Qwen3.5-35B-A3B") + ap.add_argument("--compiled-path", default="/tmp/qwen35_35b_a3b_vl_bench") + ap.add_argument("--tp", type=int, default=8) + ap.add_argument("--images", nargs="+", type=int, default=[512, 1024, 2048], + help="Image side sizes in pixels; expects /tmp/test_image_.jpg") + ap.add_argument("--buckets", nargs="+", type=int, + default=[512, 1024, 2048, 4096, 8192], + help="CTE bucket sizes to compile") + ap.add_argument("--max-new-tokens", type=int, default=48) + ap.add_argument("--repeats", type=int, default=3) + ap.add_argument("--prompt", default=DEFAULT_PROMPT) + ap.add_argument("--skip-compile", action="store_true") + ap.add_argument("--vision-compiled-dir", default=None, + help="Directory containing vision_encoder_.pt files " + "produced by compile_vision_encoder.py; if omitted, " + "runs vision encoder on CPU") + ap.add_argument("--out-json", default=None) + args = ap.parse_args() + + max_seq = max(args.buckets) + print(f"[bench] buckets={args.buckets}, max_seq={max_seq}, tp={args.tp}") + + text_config, vl_config = build_vl_config(args.model_path, args.tp, max_seq, args.buckets) + vl_model = compile_and_load( + args.model_path, args.compiled_path, text_config, vl_config, + args.skip_compile, vision_compiled_dir=args.vision_compiled_dir, + ) + + from transformers import AutoProcessor, AutoTokenizer + from PIL import Image + tok = AutoTokenizer.from_pretrained(args.model_path) + processor = AutoProcessor.from_pretrained(args.model_path) + + results = [] + for size in args.images: + img_path = f"/tmp/test_image_{size}.jpg" + image = Image.open(img_path).convert("RGB") + print(f"\n=== image {size}x{size} ({img_path}) ===") + + ttfts, tpots, texts = [], [], [] + for r in range(args.repeats + 1): # +1 warmup + ttft, tpot, n_new, text = one_image_run(vl_model, processor, tok, image, args.prompt, args.max_new_tokens) + tag = "warmup" if r == 0 else f"r{r}" + print(f" [{tag}] TTFT={ttft:7.1f} ms TPOT={tpot:6.1f} ms n_new={n_new}") + if r == 0: + continue # discard warmup + ttfts.append(ttft) + tpots.append(tpot) + texts.append(text) + + result = { + "size": size, + "prompt": args.prompt, + "ttft_ms_median": statistics.median(ttfts), + "ttft_ms_mean": statistics.mean(ttfts), + "tpot_ms_median": statistics.median(tpots), + "tpot_ms_mean": statistics.mean(tpots), + "throughput_tok_per_s": 1000.0 / statistics.mean(tpots), + "text_sample": texts[0], + } + results.append(result) + print(f" output: {texts[0]!r}") + + print("\n=== SUMMARY ===") + print(f"{'size':>6} {'TTFT (ms)':>10} {'TPOT (ms)':>10} {'tok/s':>7}") + for r in results: + print(f"{r['size']:>6} {r['ttft_ms_median']:>10.1f} {r['tpot_ms_median']:>10.1f} {r['throughput_tok_per_s']:>7.1f}") + + if args.out_json: + with open(args.out_json, "w") as f: + json.dump({ + "buckets": args.buckets, + "max_new_tokens": args.max_new_tokens, + "repeats": args.repeats, + "prompt": args.prompt, + "results": results, + }, f, indent=2) + print(f"\n[bench] wrote {args.out_json}") + + +if __name__ == "__main__": + main() diff --git a/contrib/models/Qwen3.5-35B-A3B/test/integration/run_vl_smoke.py b/contrib/models/Qwen3.5-35B-A3B/test/integration/run_vl_smoke.py new file mode 100644 index 00000000..8f68b492 --- /dev/null +++ b/contrib/models/Qwen3.5-35B-A3B/test/integration/run_vl_smoke.py @@ -0,0 +1,237 @@ +#!/usr/bin/env python +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +Vision-Language smoke test for Qwen3.5-35B-A3B on Neuron. + +Flow: + * Compile the text decoder with `use_text_only_cte_inputs=False` so it accepts + vision_embeddings, vision_mask, and mRoPE position_ids as CTE inputs. + * Load a CPU-side vision encoder (Qwen3-VL vision weights from HF safetensors). + * Preprocess a test image with `AutoProcessor`. + * Feed pixel_values → CPU vision → vision_embeddings. + * Run the Neuron-compiled text decoder with the vision embeddings scattered in. + * Print the generated caption. + +This uses CPU for vision (the Qwen3.5 ViT weights don't require huge compute +and the encoder is small). Tracing the vision encoder to Neuron is a +separate follow-up. + +Usage: + source /opt/aws_neuronx_venv_pytorch_2_9_nxd_inference/bin/activate + python contrib/models/Qwen3.5-35B-A3B/test/integration/run_vl_smoke.py \\ + --model-path /mnt/nvme/models/Qwen3.5-35B-A3B \\ + --compiled-path /tmp/qwen35_35b_a3b_vl_traced \\ + --image /path/to/image.jpg \\ + --prompt "Describe this image." +""" + +import argparse +import gc +import json +import os +import sys +import time + +# IMPORTANT: The default fused-multihead DeltaNet NKI kernel produces numerically +# unstable output on real vision embeddings (VL forward degenerates to repeated +# tokens). The legacy direct kernel is stable on VL inputs. Set these env vars +# BEFORE importing src.modeling_qwen35 so they apply during trace+compile. +os.environ.setdefault("QWEN36_DELTANET_CTE_IMPL", "legacy_direct") +os.environ.setdefault("QWEN36_DELTANET_MULTIHEAD_CTE", "0") + +import torch + +_HERE = os.path.dirname(os.path.abspath(__file__)) +_CONTRIB_ROOT = os.path.abspath(os.path.join(_HERE, "..", "..")) +if _CONTRIB_ROOT not in sys.path: + sys.path.insert(0, _CONTRIB_ROOT) + + +def build_vl_text_config(model_path: str, tp: int, seq_len: int): + """Build a Qwen35InferenceConfig with vision-aware CTE inputs enabled.""" + from neuronx_distributed_inference.models.config import ( + MoENeuronConfig, + OnDeviceSamplingConfig, + ) + from src.modeling_qwen35 import Qwen35InferenceConfig + + neuron_config = MoENeuronConfig( + tp_degree=tp, + batch_size=1, + ctx_batch_size=1, + tkg_batch_size=1, + seq_len=seq_len, + torch_dtype=torch.bfloat16, + on_device_sampling_config=OnDeviceSamplingConfig(top_k=1), + enable_bucketing=False, + flash_decoding_enabled=False, + logical_nc_config=2, + save_sharded_checkpoint=True, + moe_tp_degree=tp, + moe_ep_degree=1, + normalize_top_k_affinities=True, + blockwise_matmul_config={"use_shard_on_intermediate_dynamic_while": True}, + ) + if hasattr(neuron_config, "router_config"): + neuron_config.router_config.dtype = torch.float32 + neuron_config.router_config.act_fn = "softmax" + + with open(os.path.join(model_path, "config.json")) as f: + full = json.load(f) + text_cfg = full.get("text_config", full) + + cfg = dict(text_cfg) + cfg["pad_token_id"] = text_cfg.get("eos_token_id", 248044) + if "rope_parameters" in text_cfg: + rp = text_cfg["rope_parameters"] + cfg["rope_theta"] = rp.get("rope_theta", 10000000) + cfg["partial_rotary_factor"] = rp.get("partial_rotary_factor", 0.25) + cfg["mrope_section"] = rp.get("mrope_section", [11, 11, 10]) + cfg.setdefault("tie_word_embeddings", text_cfg.get("tie_word_embeddings", True)) + + return Qwen35InferenceConfig( + neuron_config=neuron_config, + use_hybrid_cache_manager=False, + use_text_only_cte_inputs=False, # ← accept vision + mrope inputs + **cfg, + ) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--model-path", default="/mnt/nvme/models/Qwen3.5-35B-A3B") + ap.add_argument("--compiled-path", default="/tmp/qwen35_35b_a3b_vl_traced") + ap.add_argument("--tp", type=int, default=8) + ap.add_argument("--seq-len", type=int, default=2048) + ap.add_argument("--max-new-tokens", type=int, default=64) + ap.add_argument("--image", default=None, + help="Path or URL to an image; if None, uses a random dummy") + ap.add_argument("--prompt", default="Describe this image.") + ap.add_argument("--skip-compile", action="store_true") + args = ap.parse_args() + + from src.modeling_qwen35 import NeuronQwen35ForCausalLM + from src.modeling_qwen35_vl import ( + NeuronQwen35VLForCausalLM, Qwen35VLInferenceConfig, + ) + from transformers import AutoProcessor, AutoTokenizer + + text_config = build_vl_text_config(args.model_path, args.tp, args.seq_len) + + # Extract vision_config dict from HF + with open(os.path.join(args.model_path, "config.json")) as f: + full = json.load(f) + vision_config_dict = full["vision_config"] + # Qwen3-VL VisionModel expects out_hidden_size / spatial_merge_size on the config + vision_config_dict.setdefault("spatial_merge_size", 2) + vision_config_dict.setdefault("temporal_patch_size", 2) + + vl_config = Qwen35VLInferenceConfig( + text_config=text_config, + vision_config=vision_config_dict, + image_token_id=full.get("image_token_id", 248056), + video_token_id=full.get("video_token_id", 248057), + vision_start_token_id=full.get("vision_start_token_id", 248053), + vision_end_token_id=full.get("vision_end_token_id", 248054), + spatial_merge_size=vision_config_dict["spatial_merge_size"], + ) + + # ---- Compile / load text model ---- + text_path = os.path.join(args.compiled_path, "text_model") + neff = os.path.join(text_path, "model.pt") + if not args.skip_compile and not os.path.exists(neff): + print(f"[compile-text] → {text_path}") + os.makedirs(text_path, exist_ok=True) + m = NeuronQwen35ForCausalLM(args.model_path, text_config) + t0 = time.perf_counter() + m.compile(text_path) + print(f"[compile-text] done in {(time.perf_counter()-t0):.1f} s") + del m + gc.collect() + + # ---- Build VL orchestrator ---- + vl_model = NeuronQwen35VLForCausalLM( + model_path=args.model_path, + text_config=text_config, + vision_config=vl_config, + ) + vl_model.text_model.load(text_path) + + # Load CPU vision weights (patch_embed, pos_embed, transformer blocks) + print("[load-vision] loading CPU vision encoder weights") + vl_model.vision_model_wrapper.load_cpu_model(args.model_path) + vl_model.vision_model_wrapper.load_vision_weights_from_hf(args.model_path) + + # ---- Prepare inputs ---- + processor = AutoProcessor.from_pretrained(args.model_path) + tok = AutoTokenizer.from_pretrained(args.model_path) + + # If no image, generate a dummy 224x224 RGB + if args.image is None: + print("[input] no --image provided; using dummy random image (224x224)") + from PIL import Image + import numpy as np + dummy = (np.random.rand(224, 224, 3) * 255).astype("uint8") + image = Image.fromarray(dummy) + else: + if args.image.startswith(("http://", "https://")): + import io + import urllib.request + with urllib.request.urlopen(args.image) as f: + data = f.read() + from PIL import Image + image = Image.open(io.BytesIO(data)).convert("RGB") + else: + from PIL import Image + image = Image.open(args.image).convert("RGB") + + messages = [ + {"role": "user", "content": [ + {"type": "image", "image": image}, + {"type": "text", "text": args.prompt}, + ]}, + ] + inputs = processor.apply_chat_template( + messages, + tokenize=True, + add_generation_prompt=True, + return_tensors="pt", + return_dict=True, + ) + input_ids = inputs["input_ids"] + attention_mask = inputs.get("attention_mask", torch.ones_like(input_ids)) + pixel_values = inputs["pixel_values"] + image_grid_thw = inputs["image_grid_thw"] + print(f"[input] input_ids shape={input_ids.shape}") + print(f"[input] pixel_values shape={pixel_values.shape}") + print(f"[input] image_grid_thw={image_grid_thw.tolist()}") + + # ---- Generate ---- + print("[generate] running VL generation") + t0 = time.perf_counter() + generated = vl_model.generate( + input_ids=input_ids, + attention_mask=attention_mask, + pixel_values=pixel_values, + image_grid_thw=image_grid_thw, + max_new_tokens=args.max_new_tokens, + temperature=0.0, + ) + dt = time.perf_counter() - t0 + + new_ids = generated[0].tolist()[input_ids.shape[1]:] + text = tok.decode(new_ids, skip_special_tokens=True) + + print("=" * 72) + print(f"prompt : {args.prompt!r}") + print(f"image : {args.image}") + print(f"n_new : {len(new_ids)}") + print(f"elapsed : {dt:.1f} s") + print(f"generated : {text!r}") + print("=" * 72) + + +if __name__ == "__main__": + main() diff --git a/contrib/models/Qwen3.5-35B-A3B/test/integration/test_model.py b/contrib/models/Qwen3.5-35B-A3B/test/integration/test_model.py new file mode 100644 index 00000000..e2c06dfc --- /dev/null +++ b/contrib/models/Qwen3.5-35B-A3B/test/integration/test_model.py @@ -0,0 +1,253 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +Integration tests for Qwen3.5-35B-A3B on Neuron (Trn2). + +Runs full compile + inference against the HuggingFace Qwen/Qwen3.5-35B-A3B weights. +The 2B model shares the hybrid DeltaNet + full-attention architecture of +Qwen3.5-27B (model_type = qwen3_5) at a smaller scale (24 layers, 18 linear ++ 6 full). This test reuses the modeling code (originally contributed for +Qwen3.6-27B in PR #173) at 2B scale. + +Env vars: + QWEN35_MODEL_PATH HF weights path (required) + QWEN35_COMPILED_PATH Where to write NEFF (default /tmp/qwen35_35b_a3b_traced) + QWEN35_TP_DEGREE TP degree (default 8 for 2B on Trn2) + QWEN35_SEQ_LEN Max seq len (default 512) + QWEN35_MAX_NEW_TOKENS Decode budget for latency measurement (default 32) + TTFT_THRESHOLD_MS Guard rail for TTFT (default 8000) + THROUGHPUT_THRESHOLD Min decode throughput tok/s (default 4.0) +""" + +import gc +import json +import os +import sys +import time + +import pytest +import torch + + +_CONTRIB_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) +if _CONTRIB_ROOT not in sys.path: + sys.path.insert(0, _CONTRIB_ROOT) + + +MODEL_PATH = os.environ.get("QWEN35_MODEL_PATH", "/mnt/nvme/models/Qwen3.5-35B-A3B") +COMPILED_PATH = os.environ.get("QWEN35_COMPILED_PATH", "/tmp/qwen35_35b_a3b_traced") +TP_DEGREE = int(os.environ.get("QWEN35_TP_DEGREE", "8")) +SEQ_LEN = int(os.environ.get("QWEN35_SEQ_LEN", "512")) +MAX_NEW_TOKENS = int(os.environ.get("QWEN35_MAX_NEW_TOKENS", "32")) +TTFT_THRESHOLD_MS = float(os.environ.get("TTFT_THRESHOLD_MS", "8000")) +THROUGHPUT_THRESHOLD = float(os.environ.get("THROUGHPUT_THRESHOLD", "4.0")) + +requires_model_path = pytest.mark.skipif( + not os.path.isdir(MODEL_PATH), + reason=f"Qwen3.5-35B-A3B weights not found at {MODEL_PATH}", +) + + +@pytest.fixture(scope="module") +def compiled_model(): + from neuronx_distributed_inference.models.config import ( + NeuronConfig, + OnDeviceSamplingConfig, + ) + from src.modeling_qwen35 import Qwen35InferenceConfig, NeuronQwen35ForCausalLM + + neuron_config = NeuronConfig( + tp_degree=TP_DEGREE, + batch_size=1, + ctx_batch_size=1, + tkg_batch_size=1, + seq_len=SEQ_LEN, + torch_dtype=torch.bfloat16, + on_device_sampling_config=OnDeviceSamplingConfig(top_k=1), + enable_bucketing=False, + flash_decoding_enabled=False, + logical_nc_config=2, + save_sharded_checkpoint=True, + ) + + with open(os.path.join(MODEL_PATH, "config.json")) as f: + full_config = json.load(f) + text_config = full_config.get("text_config", full_config) + cfg_dict = dict(text_config) + cfg_dict["pad_token_id"] = text_config.get("eos_token_id", 248044) + if "rope_parameters" in text_config: + cfg_dict["rope_theta"] = text_config["rope_parameters"].get("rope_theta", 10000000) + cfg_dict["partial_rotary_factor"] = text_config["rope_parameters"].get( + "partial_rotary_factor", 0.25 + ) + cfg_dict["mrope_section"] = text_config["rope_parameters"].get( + "mrope_section", [11, 11, 10] + ) + cfg_dict.setdefault("tie_word_embeddings", text_config.get("tie_word_embeddings", True)) + + inf_config = Qwen35InferenceConfig( + neuron_config=neuron_config, + use_hybrid_cache_manager=False, + **cfg_dict, + ) + + neff = os.path.join(COMPILED_PATH, "model.pt") + if not os.path.exists(neff): + print(f"[qwen35_35b_a3b] compiling → {COMPILED_PATH}") + m = NeuronQwen35ForCausalLM(MODEL_PATH, inf_config) + m.compile(COMPILED_PATH) + del m + gc.collect() + + print(f"[qwen35_35b_a3b] loading ← {COMPILED_PATH}") + m = NeuronQwen35ForCausalLM(COMPILED_PATH) + m.load(COMPILED_PATH) + return m + + +@pytest.fixture(scope="module") +def tokenizer(): + from transformers import AutoTokenizer + tok = AutoTokenizer.from_pretrained(MODEL_PATH, padding_side="right") + if tok.pad_token is None: + tok.pad_token = tok.eos_token + return tok + + +def _generate_and_time(model, tokenizer, prompt, max_new_tokens): + """Return (token_ids, decoded_text, ttft_ms, tpot_ms).""" + from transformers import GenerationConfig + import transformers + from neuronx_distributed_inference.utils.hf_adapter import ( + HuggingFaceGenerationAdapter, + ) + + gen_cfg = GenerationConfig( + do_sample=True, + top_k=1, + pad_token_id=tokenizer.pad_token_id, + eos_token_id=tokenizer.eos_token_id, + ) + gen_cfg.transformers_version = transformers.__version__ + + enc = tokenizer(prompt, padding=True, return_tensors="pt") + gen_model = HuggingFaceGenerationAdapter(model) + gen_model.generation_config.transformers_version = transformers.__version__ + + t0 = time.perf_counter() + out_first = gen_model.generate( + enc.input_ids, + generation_config=gen_cfg, + attention_mask=enc.attention_mask, + max_new_tokens=1, + ) + ttft_ms = (time.perf_counter() - t0) * 1000 + + t1 = time.perf_counter() + out_all = gen_model.generate( + enc.input_ids, + generation_config=gen_cfg, + attention_mask=enc.attention_mask, + max_new_tokens=max_new_tokens, + ) + elapsed_all = time.perf_counter() - t1 + + generated = out_all[0].tolist()[enc.input_ids.shape[-1]:] + # After TTFT, the remaining (n-1) tokens run in decode. TPOT is the + # per-decode-step latency. + n_new = max(1, len(generated)) + n_decode = max(1, n_new - 1) + ttft_end2end_ms = (t1 - t0) # approximate; use one-token pass above + tpot_ms = (elapsed_all * 1000 - ttft_ms) / n_decode + + text = tokenizer.decode(out_all[0], skip_special_tokens=True) + return out_all[0].tolist(), text, ttft_ms, tpot_ms + + +@requires_model_path +def test_generation_and_latency(compiled_model, tokenizer): + """Compile, load, and generate. Report TTFT + TPOT + text.""" + prompt = "The capital of France is" + tokens, text, ttft, tpot = _generate_and_time( + compiled_model, tokenizer, prompt, MAX_NEW_TOKENS + ) + print(f"\n[qwen35_35b_a3b] prompt : {prompt!r}") + print(f"[qwen35_35b_a3b] output : {text!r}") + print(f"[qwen35_35b_a3b] TTFT : {ttft:.1f} ms") + print(f"[qwen35_35b_a3b] TPOT : {tpot:.1f} ms ({1000.0/max(tpot,1e-6):.2f} tok/s)") + + assert text and len(text) > len(prompt), "Model produced no continuation" + assert ttft < TTFT_THRESHOLD_MS, ( + f"TTFT {ttft:.1f} ms > threshold {TTFT_THRESHOLD_MS} ms" + ) + tokens_per_sec = 1000.0 / max(tpot, 1e-6) + assert tokens_per_sec > THROUGHPUT_THRESHOLD, ( + f"decode throughput {tokens_per_sec:.2f} tok/s < {THROUGHPUT_THRESHOLD}" + ) + + +@requires_model_path +def test_accuracy_vs_hf(compiled_model, tokenizer): + """Greedy top-1 match against HuggingFace on CPU for short prompts. + + HF may be prohibitively slow at 2B on CPU, so we bail out if it takes + too long — this is a coarse coherence check, not a strict harness. + """ + prompts = [ + "The capital of France is", + "In one word, the color of the sky is", + ] + + hf_ok = False + try: + from transformers import AutoModelForCausalLM, AutoTokenizer + + hf_model = AutoModelForCausalLM.from_pretrained( + MODEL_PATH, torch_dtype=torch.bfloat16, trust_remote_code=True + ) + hf_model.eval() + hf_ok = True + except Exception as e: + print(f"[qwen35_35b_a3b] skipping HF cross-check: {e}") + pytest.skip("HuggingFace model load failed on CPU") + + from transformers import GenerationConfig + import transformers + from neuronx_distributed_inference.utils.hf_adapter import ( + HuggingFaceGenerationAdapter, + ) + + n_match = 0 + n_total = 0 + for p in prompts: + enc = tokenizer(p, return_tensors="pt") + with torch.no_grad(): + hf_out = hf_model.generate( + enc.input_ids, do_sample=False, max_new_tokens=16 + )[0].tolist() + + gen_cfg = GenerationConfig( + do_sample=True, top_k=1, + pad_token_id=tokenizer.pad_token_id, + eos_token_id=tokenizer.eos_token_id, + ) + gen_cfg.transformers_version = transformers.__version__ + gen_model = HuggingFaceGenerationAdapter(compiled_model) + gen_model.generation_config.transformers_version = transformers.__version__ + nxd_out = gen_model.generate( + enc.input_ids, + generation_config=gen_cfg, + max_new_tokens=16, + )[0].tolist() + + L = min(len(hf_out), len(nxd_out)) + for a, b in zip(hf_out[:L], nxd_out[:L]): + n_total += 1 + if a == b: + n_match += 1 + + print(f"[qwen35_35b_a3b] HF↔Neuron token match: {n_match}/{n_total}") + assert n_match / max(n_total, 1) >= 0.75, ( + f"HF↔Neuron accuracy too low: {n_match}/{n_total}" + ) diff --git a/contrib/models/Qwen3.5-35B-A3B/test/unit/__init__.py b/contrib/models/Qwen3.5-35B-A3B/test/unit/__init__.py new file mode 100644 index 00000000..e69de29b