diff --git a/README.md b/README.md index eb468f33b..af37e6c18 100644 --- a/README.md +++ b/README.md @@ -15,6 +15,7 @@ A lightweight vLLM implementation built from scratch. * 🚀 **Fast offline inference** - Comparable inference speeds to vLLM * 📖 **Readable codebase** - Clean implementation in ~ 1,200 lines of Python code * ⚡ **Optimization Suite** - Prefix caching, Tensor Parallelism, Torch compilation, CUDA graph, etc. +* 🖼️ **Multimodal support** - Qwen3.5 text + vision with CUDA graph decode ## Installation @@ -24,13 +25,20 @@ pip install git+https://github.com/GeeeekExplorer/nano-vllm.git ## Model Download -To download the model weights manually, use the following command: +### Qwen3 (text-only) ```bash huggingface-cli download --resume-download Qwen/Qwen3-0.6B \ --local-dir ~/huggingface/Qwen3-0.6B/ \ --local-dir-use-symlinks False ``` +### Qwen3.5 (multimodal) +```bash +huggingface-cli download --resume-download Qwen/Qwen3.5-0.8B \ + --local-dir ~/huggingface/Qwen3.5-0.8B/ \ + --local-dir-use-symlinks False +``` + ## Quick Start See `example.py` for usage. The API mirrors vLLM's interface with minor differences in the `LLM.generate` method: @@ -43,9 +51,47 @@ outputs = llm.generate(prompts, sampling_params) outputs[0]["text"] ``` +## Qwen3.5 Quick Start + +`example_qwen3_5.py` demonstrates both text-only and multimodal (image+text) inference. CUDA graph decode is enabled by default for maximum throughput. + +### Text-only batch +```bash +python example_qwen3_5.py --mode text --model ~/huggingface/Qwen3.5-0.8B/ +``` + +### Multimodal (image + text) +```bash +python example_qwen3_5.py --mode multimodal --model ~/huggingface/Qwen3.5-0.8B/ +``` + +### Python API +```python +from nanovllm import LLM, SamplingParams +from transformers import AutoProcessor + +llm = LLM( + "~/huggingface/Qwen3.5-0.8B/", + is_multimodal=True, + enforce_eager=False, # CUDA graph decode by default +) +processor = AutoProcessor.from_pretrained("~/huggingface/Qwen3.5-0.8B/") + +# Text-only +outputs = llm.generate([prompt_text], SamplingParams(temperature=0.6, max_tokens=256)) + +# Multimodal +requests = [{"text": chat_prompt, "images": [image]}] +outputs = llm.generate_multimodal( + requests, + [SamplingParams(temperature=0.6, max_tokens=256)], + processor, +) +``` + ## Benchmark -See `bench.py` for benchmark. +See `bench.py` for Qwen3 benchmark and `bench_qwen3_5.py` for Qwen3.5 benchmark. **Test Configuration:** - Hardware: RTX 4070 Laptop (8GB) @@ -54,12 +100,33 @@ See `bench.py` for benchmark. - Input Length: Randomly sampled between 100–1024 tokens - Output Length: Randomly sampled between 100–1024 tokens -**Performance Results:** +**Qwen3 Performance Results:** | Inference Engine | Output Tokens | Time (s) | Throughput (tokens/s) | |----------------|-------------|----------|-----------------------| | vLLM | 133,966 | 98.37 | 1361.84 | | Nano-vLLM | 133,966 | 93.41 | 1434.13 | +**Qwen3.5 Benchmark Commands:** +```bash +# Text-only throughput (same setup as Qwen3 benchmark) +python bench_qwen3_5.py --mode text --num-seqs 256 --temperature 0.6 --ignore-eos --seed 0 --model ~/huggingface/Qwen3.5-0.8B/ + +# Multimodal throughput +python bench_qwen3_5.py --mode multimodal --num-images 10 --model ~/huggingface/Qwen3.5-0.8B/ +``` + +**Test Configuration:** +- Hardware: HGX H20 (96GB) +- Model: Qwen3.5-0.8B +- Text: 256 sequences, input 100–1024 tokens, output 100–1024 tokens, temperature=0.6, ignore_eos=True +- Multimodal: 10 COCO images, max_new_tokens=2048, temperature=0.6 + +**Qwen3.5 Performance Results (Nano-vLLM only):** + +| Mode | Requests | Prompt Tokens | Generated Tokens | Time (s) | Throughput (tokens/s) | +|------|----------|---------------|------------------|----------|-----------------------| +| Text (CUDA graph) | 256 | 142,827 | 133,966 | 25.09 | 5,338.40 | +| Multimodal (CUDA graph) | 10 | 2,998 | 5,034 | 9.15 | 550.08 | ## Star History diff --git a/bench_qwen3_5.py b/bench_qwen3_5.py new file mode 100644 index 000000000..72b50d873 --- /dev/null +++ b/bench_qwen3_5.py @@ -0,0 +1,419 @@ +"""Benchmark script for nano-vllm Qwen3.5: multimodal (image+text) or pure text. + +Text mode mirrors nano-vllm/bench.py: random prompt token ids, batched llm.generate(). +Defaults to CUDA graph decode for maximum throughput. +""" + +import argparse +import io +import os +import random +import time +import urllib.request +from dataclasses import dataclass +from typing import Iterable + +import torch +from PIL import Image + +from nanovllm import LLM, SamplingParams + + +DEFAULT_IMAGE_URLS: tuple[str, ...] = ( + "http://images.cocodataset.org/val2017/000000000285.jpg", + "http://images.cocodataset.org/val2017/000000000632.jpg", + "http://images.cocodataset.org/val2017/000000000724.jpg", + "http://images.cocodataset.org/val2017/000000000776.jpg", + "http://images.cocodataset.org/val2017/000000001000.jpg", + "http://images.cocodataset.org/val2017/000000001268.jpg", + "http://images.cocodataset.org/val2017/000000006012.jpg", + "http://images.cocodataset.org/val2017/000000190236.jpg", + "http://images.cocodataset.org/val2017/000000331352.jpg", + "http://images.cocodataset.org/val2017/000000517069.jpg", +) + +DEFAULT_PROMPT = ( + "Please describe the scene in the image and highlight the primary objects." +) + + +def get_env_float(name: str, default: float) -> float: + value = os.getenv(name) + if value is None: + return default + try: + return float(value) + except ValueError: + return default + + +def get_env_int(name: str, default: int) -> int: + value = os.getenv(name) + if value is None: + return default + try: + return int(value) + except ValueError: + return default + + +def download_image(url: str) -> Image.Image: + with urllib.request.urlopen(url) as response: + payload = response.read() + return Image.open(io.BytesIO(payload)).convert("RGB") + + +@dataclass +class BenchmarkResult: + num_requests: int + total_prompt_tokens: int + total_generated_tokens: int + latency: float + + @property + def tok_per_sec(self) -> float: + if self.latency <= 0: + return 0.0 + return self.total_generated_tokens / self.latency + + +def build_requests( + processor, + image_urls: Iterable[str], + prompt: str, +) -> list[dict]: + requests = [] + for url in image_urls: + image = download_image(url) + chat_prompt = processor.apply_chat_template( + [ + { + "role": "user", + "content": [ + {"type": "image", "image": image}, + {"type": "text", "text": prompt}, + ], + } + ], + tokenize=False, + add_generation_prompt=True, + ) + requests.append( + { + "text": chat_prompt, + "images": [image], + "meta": {"url": url}, + } + ) + return requests + + +def run_benchmark_text( + model_path: str, + num_seqs: int, + min_prompt_len: int, + max_prompt_len: int, + min_output_len: int, + max_output_len: int, + temperature: float, + seed: int, + ignore_eos: bool, + enforce_eager: bool, +) -> BenchmarkResult: + """Pure-text throughput test (same spirit as nano-vllm/bench.py).""" + from transformers import AutoTokenizer + + random.seed(seed) + tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) + vocab_size = getattr(tokenizer, "vocab_size", None) or len(tokenizer) + hi = min(10_000, max(vocab_size - 1, 1)) + + prompt_token_ids = [ + [random.randint(0, hi) for _ in range(random.randint(min_prompt_len, max_prompt_len))] + for _ in range(num_seqs) + ] + sampling_params_list = [ + SamplingParams( + temperature=temperature, + max_tokens=random.randint(min_output_len, max_output_len), + ignore_eos=ignore_eos, + ) + for _ in range(num_seqs) + ] + + max_prompt_toks = max(len(p) for p in prompt_token_ids) + max_out_toks = max(sp.max_tokens for sp in sampling_params_list) + max_model_len = max(2048, max_prompt_toks + max_out_toks + 32) + max_num_batched_tokens = max(16384, num_seqs * max_prompt_len + 4096) + + if not enforce_eager: + os.environ["NANOVLLM_ALLOW_QWEN35_DECODE_CUDAGRAPH"] = "1" + + llm = LLM( + model_path, + enforce_eager=enforce_eager, + tensor_parallel_size=1, + is_multimodal=False, + max_model_len=max_model_len, + max_num_batched_tokens=max_num_batched_tokens, + max_num_seqs=max(16, num_seqs), + ) + + # Warm-up + llm.generate( + ["Benchmark: "], + SamplingParams(temperature=temperature, max_tokens=8), + use_tqdm=False, + ) + + if torch.cuda.is_available(): + torch.cuda.synchronize() + start = time.perf_counter() + outputs = llm.generate( + prompt_token_ids, + sampling_params_list, + use_tqdm=False, + ) + if torch.cuda.is_available(): + torch.cuda.synchronize() + latency = time.perf_counter() - start + + total_generated = sum(len(item["token_ids"]) for item in outputs) + total_prompt_tokens = sum(len(p) for p in prompt_token_ids) + + return BenchmarkResult( + num_requests=num_seqs, + total_prompt_tokens=total_prompt_tokens, + total_generated_tokens=total_generated, + latency=latency, + ) + + +def run_benchmark_multimodal( + model_path: str, + max_new_tokens: int, + temperature: float, + image_urls: Iterable[str], + enforce_eager: bool = False, +) -> BenchmarkResult: + from transformers import AutoProcessor + + urls = list(image_urls) + num_urls = len(urls) + + if not enforce_eager: + os.environ["NANOVLLM_ALLOW_QWEN35_DECODE_CUDAGRAPH"] = "1" + + llm = LLM( + model_path, + enforce_eager=enforce_eager, + tensor_parallel_size=1, + is_multimodal=True, + max_model_len=2048, + max_num_batched_tokens=16384, + max_num_seqs=max(16, num_urls), + ) + + processor = AutoProcessor.from_pretrained(model_path) + + requests = build_requests(processor, urls, DEFAULT_PROMPT) + num_requests = len(requests) + sampling_params = SamplingParams( + temperature=temperature, + max_tokens=max_new_tokens, + ) + sampling_params_list = [sampling_params] * num_requests + + # Warm-up with a single request + llm.generate_multimodal( + [requests[0]], + sampling_params_list[:1], + processor, + use_tqdm=False, + ) + + if torch.cuda.is_available(): + torch.cuda.synchronize() + start = time.perf_counter() + outputs = llm.generate_multimodal( + requests, + sampling_params_list, + processor, + use_tqdm=False, + ) + if torch.cuda.is_available(): + torch.cuda.synchronize() + latency = time.perf_counter() - start + + total_generated = sum(len(item["token_ids"]) for item in outputs) + + prompt_token_lengths = 0 + for req in requests: + encoded = processor( + text=[req["text"]], + images=req["images"], + return_tensors="pt", + ) + prompt_token_lengths += encoded["input_ids"].shape[-1] + + return BenchmarkResult( + num_requests=num_requests, + total_prompt_tokens=prompt_token_lengths, + total_generated_tokens=total_generated, + latency=latency, + ) + + +def parse_args() -> argparse.Namespace: + default_out_seq = get_env_int("out_seq_length", 2048) + default_temperature = get_env_float("temperature", 0.2) + + parser = argparse.ArgumentParser( + description="Benchmark nano-vllm Qwen3.5 (multimodal or pure text)." + ) + parser.add_argument( + "--model", + type=str, + default=None, + help="Path to the Qwen3.5 model directory.", + ) + parser.add_argument( + "--mode", + type=str, + choices=("multimodal", "text"), + default="text", + help=( + "multimodal: image+COCO URLs via generate_multimodal; " + "text: random prompt ids + llm.generate (like nano-vllm/bench.py)." + ), + ) + parser.add_argument( + "--max-new-tokens", + type=int, + default=default_out_seq, + help="(multimodal) Maximum new tokens per request.", + ) + parser.add_argument( + "--temperature", + type=float, + default=default_temperature, + help="Sampling temperature.", + ) + parser.add_argument( + "--num-images", + type=int, + default=len(DEFAULT_IMAGE_URLS), + help="(multimodal) Number of images (picked from a fixed COCO subset).", + ) + parser.add_argument( + "--num-seqs", + type=int, + default=64, + help="(text) Number of concurrent random prompts.", + ) + parser.add_argument( + "--min-prompt-len", + type=int, + default=100, + help="(text) Inclusive lower bound for random prompt length.", + ) + parser.add_argument( + "--max-prompt-len", + type=int, + default=1024, + help="(text) Inclusive upper bound for random prompt length.", + ) + parser.add_argument( + "--min-output-len", + type=int, + default=100, + help="(text) Inclusive lower bound for max_tokens per sequence.", + ) + parser.add_argument( + "--max-output-len", + type=int, + default=1024, + help="(text) Inclusive upper bound for max_tokens per sequence.", + ) + parser.add_argument( + "--seed", + type=int, + default=0, + help="(text) RNG seed for reproducible random prompts/sampling params.", + ) + parser.add_argument( + "--ignore-eos", + dest="ignore_eos", + action="store_true", + help="(text) Force decoding until max_tokens (higher sustained tok/s).", + ) + parser.add_argument( + "--no-ignore-eos", + dest="ignore_eos", + action="store_false", + help="(text) Stop at EOS.", + ) + parser.add_argument( + "--enforce-eager", + action="store_true", + default=False, + help="Disable CUDA graph capture (enforce eager mode).", + ) + parser.set_defaults(ignore_eos=True) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + args.model = os.path.expanduser(args.model) + + if args.mode == "text": + if args.num_seqs < 1: + raise ValueError("num-seqs must be positive.") + if args.min_prompt_len > args.max_prompt_len: + raise ValueError("min-prompt-len must be <= max-prompt-len.") + if args.min_output_len > args.max_output_len: + raise ValueError("min-output-len must be <= max-output-len.") + result = run_benchmark_text( + model_path=args.model, + num_seqs=args.num_seqs, + min_prompt_len=args.min_prompt_len, + max_prompt_len=args.max_prompt_len, + min_output_len=args.min_output_len, + max_output_len=args.max_output_len, + temperature=args.temperature, + seed=args.seed, + ignore_eos=args.ignore_eos, + enforce_eager=args.enforce_eager, + ) + print("=== nano-vllm Qwen3.5 text benchmark ===") + print( + "Mode : random prompt token ids, llm.generate() " + "(see nano-vllm/bench.py)" + ) + else: + image_urls = DEFAULT_IMAGE_URLS[: args.num_images] + if not image_urls: + raise ValueError("num-images must be positive.") + result = run_benchmark_multimodal( + model_path=args.model, + max_new_tokens=args.max_new_tokens, + temperature=args.temperature, + image_urls=image_urls, + enforce_eager=args.enforce_eager, + ) + print("=== nano-vllm Qwen3.5 multimodal benchmark ===") + print( + "Mode : one generate_multimodal() over all requests " + "(batched in engine)" + ) + + print(f"Requests : {result.num_requests}") + print(f"Prompt tokens : {result.total_prompt_tokens}") + print(f"Generated tokens : {result.total_generated_tokens}") + print(f"Latency : {result.latency:.2f}s") + print(f"Throughput : {result.tok_per_sec:.2f} tok/s") + + +if __name__ == "__main__": + main() diff --git a/example_qwen3_5.py b/example_qwen3_5.py new file mode 100644 index 000000000..88134f017 --- /dev/null +++ b/example_qwen3_5.py @@ -0,0 +1,207 @@ +import argparse +import io +import os +import urllib.request + +from PIL import Image +from nanovllm import LLM, SamplingParams +from transformers import AutoProcessor + + +def get_env_float(name: str, default: float) -> float: + value = os.getenv(name) + if value is None: + return default + try: + return float(value) + except ValueError: + return default + + +def get_env_int(name: str, default: int) -> int: + value = os.getenv(name) + if value is None: + return default + try: + return int(value) + except ValueError: + return default + + +def download_image(url: str) -> Image.Image: + """Download an image from URL and return as PIL Image.""" + with urllib.request.urlopen(url) as response: + payload = response.read() + return Image.open(io.BytesIO(payload)).convert("RGB") + + +def parse_args() -> argparse.Namespace: + default_max_tokens = get_env_int("out_seq_length", 2048) + default_temperature = get_env_float("temperature", 0.6) + parser = argparse.ArgumentParser( + description="nano-vllm Qwen3.5 multimodal example (text and/or image+text)." + ) + parser.add_argument( + "--model", + type=str, + default=None, + help="Model directory (same as AutoProcessor.from_pretrained).", + ) + parser.add_argument( + "--max-model-len", + type=int, + default=2048, + help="LLM max_model_len (KV / context budget).", + ) + parser.add_argument( + "--max-num-batched-tokens", + type=int, + default=8192, + help="Scheduler max tokens per forward step (all sequences).", + ) + parser.add_argument( + "--max-num-seqs", + type=int, + default=8, + help="Max concurrent sequences in one batch.", + ) + parser.add_argument( + "--enforce-eager", + action="store_true", + default=False, + help="Disable CUDA graph capture (enforce eager mode).", + ) + parser.add_argument( + "--tensor-parallel-size", + type=int, + default=1, + help="Tensor parallel size.", + ) + parser.add_argument( + "--max-tokens", + type=int, + default=default_max_tokens, + help="Sampling max_tokens (env: out_seq_length).", + ) + parser.add_argument( + "--temperature", + type=float, + default=default_temperature, + help="Sampling temperature (env: temperature).", + ) + parser.add_argument( + "--mode", + type=str, + choices=("text", "multimodal", "all"), + default="text", + help="Which mode to run: text-only batch, multimodal, or both.", + ) + return parser.parse_args() + + +def example_text(llm: LLM, processor: AutoProcessor, sampling_params: SamplingParams) -> None: + print("=" * 50) + print("Example 1: Text-only prompts") + print("=" * 50) + + text_prompts = [ + "Give me a short introduction to large language models.", + "introduce yourself", + "list all prime numbers within 100", + ] + chat_prompts = [ + processor.apply_chat_template( + [{"role": "user", "content": prompt}], + tokenize=False, + add_generation_prompt=True, + ) + for prompt in text_prompts + ] + outputs = llm.generate(chat_prompts, sampling_params) + + for prompt, output in zip(chat_prompts, outputs): + print(f"\nPrompt: {prompt!r}") + print(f"Completion: {output['text']!r}") + + +def example_multimodal( + llm: LLM, + processor: AutoProcessor, + sampling_params: SamplingParams, +) -> None: + print("\n" + "=" * 50) + print("Example 2: Multimodal prompts with images") + print("=" * 50) + + image_urls = [ + "http://images.cocodataset.org/val2017/000000000285.jpg", + "http://images.cocodataset.org/val2017/000000000632.jpg", + ] + + requests = [] + for url in image_urls: + image = download_image(url) + prompt_text = ( + "Please describe the scene in the image and highlight the primary objects." + ) + + chat_prompt = processor.apply_chat_template( + [ + { + "role": "user", + "content": [ + {"type": "image", "image": image}, + {"type": "text", "text": prompt_text}, + ], + } + ], + tokenize=False, + add_generation_prompt=True, + ) + requests.append({ + "text": chat_prompt, + "images": [image], + }) + + sampling_params_list = [sampling_params] * len(requests) + outputs = llm.generate_multimodal( + requests, + sampling_params_list, + processor, + use_tqdm=True, + ) + + for req, output in zip(requests, outputs): + print(f"\nPrompt: {req['text'][:100]}...") + print(f"Completion: {output['text']!r}") + + +def main() -> None: + args = parse_args() + path = os.path.expanduser(args.model) + + llm = LLM( + path, + enforce_eager=args.enforce_eager, + tensor_parallel_size=args.tensor_parallel_size, + is_multimodal=True, + max_model_len=args.max_model_len, + max_num_batched_tokens=args.max_num_batched_tokens, + max_num_seqs=args.max_num_seqs, + ) + + processor = AutoProcessor.from_pretrained(path) + + sampling_params = SamplingParams( + temperature=args.temperature, + max_tokens=args.max_tokens, + ) + + if args.mode in ("text", "all"): + example_text(llm, processor, sampling_params) + if args.mode in ("multimodal", "all"): + example_multimodal(llm, processor, sampling_params) + + +if __name__ == "__main__": + main() diff --git a/nanovllm/config.py b/nanovllm/config.py index 7066cbeb5..21b1a6ae3 100644 --- a/nanovllm/config.py +++ b/nanovllm/config.py @@ -3,7 +3,7 @@ from transformers import AutoConfig -@dataclass(slots=True) +@dataclass class Config: model: str max_num_batched_tokens: int = 16384 @@ -12,6 +12,7 @@ class Config: gpu_memory_utilization: float = 0.9 tensor_parallel_size: int = 1 enforce_eager: bool = False + is_multimodal: bool = False hf_config: AutoConfig | None = None eos: int = -1 kvcache_block_size: int = 256 @@ -22,4 +23,30 @@ def __post_init__(self): assert self.kvcache_block_size % 256 == 0 assert 1 <= self.tensor_parallel_size <= 8 self.hf_config = AutoConfig.from_pretrained(self.model) - self.max_model_len = min(self.max_model_len, self.hf_config.max_position_embeddings) + + # Auto-detect multimodal from model_type + _model_type = getattr(self.hf_config, "model_type", None) or "" + if _model_type in ( + "qwen3_vl", "qwen3_vl_moe", "qwen3_5", "qwen3_5_moe" + ): + self.is_multimodal = True + + # Multimodal: text config in text_config + text_config = getattr( + self.hf_config, "text_config", self.hf_config + ) + + max_position_embeddings = getattr( + text_config, "max_position_embeddings", None + ) + if max_position_embeddings is not None: + self.max_model_len = min( + self.max_model_len, max_position_embeddings + ) + + # eos may be defined within the text config + eos_token_id = getattr(text_config, "eos_token_id", None) + if eos_token_id is not None: + self.eos = eos_token_id + + assert self.max_num_batched_tokens >= self.max_model_len diff --git a/nanovllm/engine/block_manager.py b/nanovllm/engine/block_manager.py index a48989ca0..2aad16659 100644 --- a/nanovllm/engine/block_manager.py +++ b/nanovllm/engine/block_manager.py @@ -25,12 +25,13 @@ def reset(self): class BlockManager: - def __init__(self, num_blocks: int, block_size: int): + def __init__(self, num_blocks: int, block_size: int, non_cache_token_ids: set[int] | None = None): self.block_size = block_size self.blocks: list[Block] = [Block(i) for i in range(num_blocks)] self.hash_to_block_id: dict[int, int] = dict() self.free_block_ids: deque[int] = deque(range(num_blocks)) self.used_block_ids: set[int] = set() + self.non_cache_token_ids = non_cache_token_ids or set() @classmethod def compute_hash(cls, token_ids: list[int], prefix: int = -1): @@ -40,8 +41,11 @@ def compute_hash(cls, token_ids: list[int], prefix: int = -1): h.update(np.array(token_ids).tobytes()) return h.intdigest() - def _allocate_block(self) -> int: - block_id = self.free_block_ids.popleft() + def _allocate_block(self, block_id: int | None = None) -> int: + if block_id is None: + block_id = self.free_block_ids.popleft() + else: + self.free_block_ids.remove(block_id) block = self.blocks[block_id] assert block.ref_count == 0 if block.hash != -1 and self.hash_to_block_id.get(block.hash) == block_id: @@ -55,11 +59,21 @@ def _deallocate_block(self, block_id: int): self.used_block_ids.remove(block_id) self.free_block_ids.append(block_id) + def _block_has_non_cache_tokens(self, seq: Sequence, block_idx: int) -> bool: + """Check if a block contains any non-cacheable tokens (e.g. vision tokens).""" + if not self.non_cache_token_ids: + return False + token_ids = seq.block(block_idx) + return bool(self.non_cache_token_ids.intersection(token_ids)) + def can_allocate(self, seq: Sequence) -> int: h = -1 num_cached_blocks = 0 num_new_blocks = seq.num_blocks for i in range(seq.num_blocks - 1): + # Skip caching for blocks with non-cacheable tokens + if self._block_has_non_cache_tokens(seq, i): + break token_ids = seq.block(i) h = self.compute_hash(token_ids, h) block_id = self.hash_to_block_id.get(h, -1) @@ -113,6 +127,10 @@ def hash_blocks(self, seq: Sequence): if start == end: return h = self.blocks[seq.block_table[start - 1]].hash if start > 0 else -1 for i in range(start, end): + # Skip caching for blocks with non-cacheable tokens + if self._block_has_non_cache_tokens(seq, i): + h = -1 + continue block = self.blocks[seq.block_table[i]] token_ids = seq.block(i) h = self.compute_hash(token_ids, h) diff --git a/nanovllm/engine/llm_engine.py b/nanovllm/engine/llm_engine.py index 3685094c9..a0a2b7b61 100644 --- a/nanovllm/engine/llm_engine.py +++ b/nanovllm/engine/llm_engine.py @@ -3,6 +3,7 @@ from time import perf_counter from tqdm.auto import tqdm from transformers import AutoTokenizer +import torch import torch.multiprocessing as mp from nanovllm.config import Config @@ -18,7 +19,6 @@ def __init__(self, model, **kwargs): config_fields = {field.name for field in fields(Config)} config_kwargs = {k: v for k, v in kwargs.items() if k in config_fields} config = Config(model, **config_kwargs) - Sequence.block_size = config.kvcache_block_size self.ps = [] self.events = [] ctx = mp.get_context("spawn") @@ -29,29 +29,168 @@ def __init__(self, model, **kwargs): self.ps.append(process) self.events.append(event) self.model_runner = ModelRunner(config, 0, self.events) - self.tokenizer = AutoTokenizer.from_pretrained(config.model, use_fast=True) + # Load tokenizer without use_fast to match HF behavior (same tokenization as tokenizer([text], return_tensors="pt")). + self.tokenizer = AutoTokenizer.from_pretrained(config.model, trust_remote_code=True) config.eos = self.tokenizer.eos_token_id self.scheduler = Scheduler(config) atexit.register(self.exit) + def _expand_vision_placeholders( + self, + input_ids: list[int], + image_grid_thw: torch.Tensor, + ) -> tuple[list[int], list[int], list[tuple[int, int]]]: + hf_config = self.model_runner.config.hf_config + vision_config = hf_config.vision_config + merge_size = vision_config.spatial_merge_size + + image_token_id = getattr(hf_config, "image_token_id", None) + vision_start_token_id = getattr(hf_config, "vision_start_token_id", None) + vision_end_token_id = getattr(hf_config, "vision_end_token_id", None) + + if None in (image_token_id, vision_start_token_id, vision_end_token_id): + raise ValueError("缺少视觉占位符相关的 token id 配置") + + if image_grid_thw.dim() != 2 or image_grid_thw.size(-1) != 3: + raise ValueError("image_grid_thw 形状不正确，期望为 [num_images, 3]") + + grids = image_grid_thw.tolist() + expected_counts = [ + int(t * h * w // (merge_size**2)) for t, h, w in grids + ] + + original_count = input_ids.count(image_token_id) + new_input_ids: list[int] = [] + i = 0 + image_idx = 0 + total_images = len(expected_counts) + length = len(input_ids) + + placeholder_ranges: list[tuple[int, int]] = [] + + while i < length: + token = input_ids[i] + if token == vision_start_token_id and image_idx < total_images: + new_input_ids.append(token) + i += 1 + # Skip original contents until matching vision_end_token_id + while i < length and input_ids[i] != vision_end_token_id: + i += 1 + if i == length: + raise ValueError("vision_start_token 后未找到匹配的 vision_end_token") + + required = expected_counts[image_idx] + start_offset = len(new_input_ids) + new_input_ids.extend([image_token_id] * required) + new_input_ids.append(vision_end_token_id) + placeholder_ranges.append((start_offset, required)) + i += 1 # Skip the original vision_end token + image_idx += 1 + else: + new_input_ids.append(token) + i += 1 + + if image_idx != total_images: + raise ValueError( + f"存在 {total_images - image_idx} 个图像没有匹配的占位符" + ) + + # new_count = new_input_ids.count(image_token_id) + # print( + # "[LLMEngine] expand placeholders:", + # { + # "original_count": original_count, + # "expected_counts": expected_counts, + # "new_count": new_count, + # }, + # ) + + return new_input_ids, expected_counts, placeholder_ranges + def exit(self): self.model_runner.call("exit") del self.model_runner for p in self.ps: p.join() - def add_request(self, prompt: str | list[int], sampling_params: SamplingParams): + def add_request( + self, + prompt: str | list[int], + sampling_params: SamplingParams, + images=None, + pixel_values=None, + image_grid_thw=None, + vision_counts=None, + vision_placeholders=None, + ): if isinstance(prompt, str): - prompt = self.tokenizer.encode(prompt) - seq = Sequence(prompt, sampling_params) + # Match HF: chat template already contains <|im_start|>, <|im_end|>, etc.; + # do not add extra BOS/special tokens so tokenization matches tokenizer([text], return_tensors="pt"). + prompt = self.tokenizer.encode(prompt, add_special_tokens=False) + + # Robustness: clamp the requested generation length so that + # (prompt_len + max_tokens) never exceeds `config.max_model_len`. + # This avoids cudagraph capture shape mismatches (e.g. block_tables + # width 8 vs 9) when users pass `max_tokens` larger than the + # remaining context budget. + max_model_len = self.model_runner.config.max_model_len + input_len = len(prompt) + if input_len > max_model_len: + raise ValueError( + f"Input length ({input_len}) exceeds model's maximum " + f"context length ({max_model_len})." + ) + remaining_budget = max_model_len - input_len + if remaining_budget < 1: + raise ValueError( + f"Remaining generation budget is {remaining_budget} " + f"(input_len={input_len}, max_model_len={max_model_len}). " + f"Please reduce the prompt length." + ) + if sampling_params.max_tokens > remaining_budget: + sampling_params = SamplingParams( + temperature=sampling_params.temperature, + max_tokens=remaining_budget, + ignore_eos=sampling_params.ignore_eos, + ) + seq = Sequence( + prompt, + sampling_params, + images=images, + pixel_values=pixel_values, + image_grid_thw=image_grid_thw, + vision_counts=vision_counts, + vision_placeholders=vision_placeholders, + ) + # if self.model_runner.config.hf_config is not None: + # image_token_id = getattr(self.model_runner.config.hf_config, "image_token_id", None) + # if image_token_id is not None: + # placeholder_count = seq.token_ids.count(image_token_id) + # print( + # "[LLMEngine] sequence placeholder count", + # {"seq_id": seq.seq_id, "placeholder_count": placeholder_count, "vision_counts": vision_counts}, + # ) self.scheduler.add(seq) def step(self): seqs, is_prefill = self.scheduler.schedule() - num_tokens = sum(seq.num_scheduled_tokens for seq in seqs) if is_prefill else -len(seqs) token_ids = self.model_runner.call("run", seqs, is_prefill) - self.scheduler.postprocess(seqs, token_ids, is_prefill) - outputs = [(seq.seq_id, seq.completion_token_ids) for seq in seqs if seq.is_finished] + self.scheduler.postprocess(seqs, token_ids) + # Clean up GDN recurrent/conv states for finished sequences to avoid + # memory leaks and stale state reuse when seq_ids are recycled. + finished_seqs = [seq for seq in seqs if seq.is_finished] + if finished_seqs: + self.model_runner.call("cleanup_seq_states", [seq.seq_id for seq in finished_seqs]) + outputs = [ + (seq.seq_id, seq.completion_token_ids) + for seq in seqs + if seq.is_finished + ] + num_tokens = ( + sum(len(seq) for seq in seqs) + if is_prefill + else -len(seqs) + ) return outputs, num_tokens def is_finished(self): @@ -63,7 +202,12 @@ def generate( sampling_params: SamplingParams | list[SamplingParams], use_tqdm: bool = True, ) -> list[str]: - pbar = tqdm(total=len(prompts), desc="Generating", dynamic_ncols=True, disable=not use_tqdm) + if use_tqdm: + pbar = tqdm( + total=len(prompts), + desc="Generating", + dynamic_ncols=True, + ) if not isinstance(sampling_params, list): sampling_params = [sampling_params] * len(prompts) for prompt, sp in zip(prompts, sampling_params): @@ -73,18 +217,159 @@ def generate( while not self.is_finished(): t = perf_counter() output, num_tokens = self.step() - if num_tokens > 0: - prefill_throughput = num_tokens / (perf_counter() - t) - else: - decode_throughput = -num_tokens / (perf_counter() - t) - pbar.set_postfix({ - "Prefill": f"{int(prefill_throughput)}tok/s", - "Decode": f"{int(decode_throughput)}tok/s", - }) + if use_tqdm: + if num_tokens > 0: + prefill_throughput = num_tokens / (perf_counter() - t) + else: + decode_throughput = -num_tokens / (perf_counter() - t) + pbar.set_postfix( + { + "Prefill": f"{int(prefill_throughput)}tok/s", + "Decode": f"{int(decode_throughput)}tok/s", + } + ) for seq_id, token_ids in output: outputs[seq_id] = token_ids - pbar.update(1) - pbar.close() - outputs = [outputs[seq_id] for seq_id in sorted(outputs.keys())] - outputs = [{"text": self.tokenizer.decode(token_ids), "token_ids": token_ids} for token_ids in outputs] + if use_tqdm: + pbar.update(1) + outputs = [ + outputs[seq_id] + for seq_id in sorted(outputs.keys()) + ] + outputs = [ + { + "text": self.tokenizer.decode(token_ids), + "token_ids": token_ids, + } + for token_ids in outputs + ] + if use_tqdm: + pbar.close() return outputs + + def generate_multimodal( + self, + requests: list[dict], + sampling_params: SamplingParams | list[SamplingParams], + processor, + use_tqdm: bool = True, + ) -> list[str]: + if use_tqdm: + pbar = tqdm( + total=len(requests), + desc="Generating", + dynamic_ncols=True, + ) + + if not isinstance(sampling_params, list): + sampling_params = [sampling_params] * len(requests) + + for request, sp in zip(requests, sampling_params): + messages = request.get("messages") + text = request.get("text") + images = request.get("images") + + if text is None: + if messages is None: + raise ValueError( + "multimodal request requires 'text' or 'messages'" + ) + + text = processor.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=True, + ) + + if images is None: + extracted_images = [] + for message in messages: + for content in message.get("content", []): + is_image = content.get("type") == "image" + has_payload = "image" in content + if is_image and has_payload: + extracted_images.append(content["image"]) + images = extracted_images if extracted_images else None + + if images is not None and not isinstance(images, (list, tuple)): + images = [images] + + processor_kwargs = { + "text": [text], + "return_tensors": "pt", + "padding": True, + } + if images: + processor_kwargs["images"] = images + + processor_outputs = processor(**processor_kwargs) + + input_ids = processor_outputs["input_ids"][0].tolist() + pixel_values = processor_outputs.get("pixel_values") + image_grid_thw = processor_outputs.get("image_grid_thw") + + vision_counts = [] + vision_placeholders = [] + if image_grid_thw is not None: + expanded_input_ids, vision_counts, vision_placeholders = self._expand_vision_placeholders( + input_ids, + image_grid_thw.squeeze(0) if image_grid_thw.dim() == 3 else image_grid_thw, + ) + input_ids = expanded_input_ids + + if pixel_values is not None: + pixel_values = pixel_values.contiguous().cpu() + + if image_grid_thw is not None: + image_grid_thw = image_grid_thw.contiguous().cpu() + + self.add_request( + input_ids, + sp, + pixel_values=pixel_values, + image_grid_thw=image_grid_thw, + vision_counts=vision_counts, + vision_placeholders=vision_placeholders, + ) + + outputs = {} + prefill_throughput = decode_throughput = 0. + while not self.is_finished(): + t = perf_counter() + output, num_tokens = self.step() + if use_tqdm: + if num_tokens > 0: + prefill_throughput = num_tokens / (perf_counter() - t) + else: + decode_throughput = -num_tokens / (perf_counter() - t) + pbar.set_postfix( + { + "Prefill": f"{int(prefill_throughput)}tok/s", + "Decode": f"{int(decode_throughput)}tok/s", + } + ) + for seq_id, token_ids in output: + outputs[seq_id] = token_ids + if use_tqdm: + pbar.update(1) + + outputs = [ + outputs[seq_id] + for seq_id in sorted(outputs.keys()) + ] + results = [ + { + "text": self.tokenizer.decode( + token_ids, + skip_special_tokens=True, + clean_up_tokenization_spaces=False, + ), + "token_ids": token_ids, + } + for token_ids in outputs + ] + + if use_tqdm: + pbar.close() + + return results diff --git a/nanovllm/engine/model_runner.py b/nanovllm/engine/model_runner.py index 71d9883c9..eaf18c97f 100644 --- a/nanovllm/engine/model_runner.py +++ b/nanovllm/engine/model_runner.py @@ -1,3 +1,4 @@ +import os import pickle import torch import torch.distributed as dist @@ -11,6 +12,52 @@ from nanovllm.utils.context import set_context, get_context, reset_context from nanovllm.utils.loader import load_model +try: + from nanovllm.models.qwen3_vl import load_qwen3_vl_model + QWEN3_VL_AVAILABLE = True +except ImportError: + QWEN3_VL_AVAILABLE = False + load_qwen3_vl_model = None +try: + from nanovllm.models.qwen3_5 import load_qwen3_5_model + QWEN3_5_AVAILABLE = True +except ImportError: + QWEN3_5_AVAILABLE = False + load_qwen3_5_model = None +MULTIMODAL_AVAILABLE = QWEN3_VL_AVAILABLE or QWEN3_5_AVAILABLE + +try: + from nanovllm.models.qwen3_next import Qwen3NextForCausalLM + QWEN3_NEXT_AVAILABLE = True +except ImportError: + QWEN3_NEXT_AVAILABLE = False + Qwen3NextForCausalLM = None + +if not MULTIMODAL_AVAILABLE: + print("[ModelRunner] 多模态模块不可用") + + +class GDNSlotManager: + """Manages slot allocation for GDN state pool buffers.""" + + def __init__(self, max_slots: int): + self.max_slots = max_slots + self._seq_to_slot: dict[int, int] = {} + self._free_slots: list[int] = list(range(max_slots - 1, -1, -1)) + + def allocate(self, seq_id: int) -> int: + slot = self._free_slots.pop() + self._seq_to_slot[seq_id] = slot + return slot + + def get_slot(self, seq_id: int) -> int | None: + return self._seq_to_slot.get(seq_id) + + def release(self, seq_id: int): + slot = self._seq_to_slot.pop(seq_id, None) + if slot is not None: + self._free_slots.append(slot) + class ModelRunner: @@ -23,15 +70,95 @@ def __init__(self, config: Config, rank: int, event: Event | list[Event]): self.rank = rank self.event = event - dist.init_process_group("nccl", "tcp://localhost:2333", world_size=self.world_size, rank=rank) - torch.cuda.set_device(rank) + self.use_cuda = torch.cuda.is_available() + self.device = torch.device("cuda", rank) if self.use_cuda else torch.device("cpu") + + # Even with world_size=1, layers call dist.get_world_size()/get_rank(). + # Initialize a process group with a CPU-capable backend when CUDA isn't available. + backend = "nccl" if (self.use_cuda and dist.is_nccl_available()) else "gloo" + dist.init_process_group( + backend, + "tcp://localhost:2333", + world_size=self.world_size, + rank=rank, + ) + + if self.use_cuda: + torch.cuda.set_device(rank) default_dtype = torch.get_default_dtype() - torch.set_default_dtype(hf_config.dtype) - torch.set_default_device("cuda") - self.model = Qwen3ForCausalLM(hf_config) - load_model(self.model, config.model) + torch_dtype = getattr(hf_config, "torch_dtype", None) + if torch_dtype is None and hasattr(hf_config, "text_config"): + torch_dtype = getattr(hf_config.text_config, "torch_dtype", None) + if isinstance(torch_dtype, str): + resolved = getattr(torch, torch_dtype, None) + if resolved is None: + alias_map = {"bf16": torch.bfloat16, "fp16": torch.float16} + resolved = alias_map.get(torch_dtype.lower()) + torch_dtype = resolved + if torch_dtype is None: + torch_dtype = torch.bfloat16 + torch.set_default_dtype(torch_dtype) + torch.set_default_device(self.device.type) + + # 根据 model_type 区分多模态 vs 纯文本，以及 qwen3_vl vs qwen3_5 / qwen3_next vs qwen3 + model_type = getattr(hf_config, "model_type", None) or "" + self.is_multimodal = ( + getattr(config, "is_multimodal", False) and MULTIMODAL_AVAILABLE + ) + if self.is_multimodal: + # 多模态：按 model_type 选择 qwen3_5 或 qwen3_vl + if model_type in ("qwen3_5", "qwen3_5_moe") and QWEN3_5_AVAILABLE and load_qwen3_5_model: + print("[ModelRunner] 加载 Qwen3.5 多模态模型 (model_type=%s)" % model_type) + self.model = load_qwen3_5_model(config.model, config) + elif QWEN3_VL_AVAILABLE and load_qwen3_vl_model: + print("[ModelRunner] 加载 Qwen3-VL 多模态模型 (model_type=%s)" % model_type) + self.model = load_qwen3_vl_model(config.model, config) + else: + raise RuntimeError( + "多模态已开启但无法加载: model_type=%s, QWEN3_5_AVAILABLE=%s, QWEN3_VL_AVAILABLE=%s" + % (model_type, QWEN3_5_AVAILABLE, QWEN3_VL_AVAILABLE) + ) + else: + # 纯文本：按 text_config.model_type 选择 Qwen3Next 或 Qwen3 + text_config = getattr(hf_config, "text_config", hf_config) + text_model_type = getattr(text_config, "model_type", None) or model_type + if text_model_type == "qwen3_next" and QWEN3_NEXT_AVAILABLE and Qwen3NextForCausalLM is not None: + print("[ModelRunner] 加载纯文本模型 Qwen3Next (model_type={})".format(text_model_type)) + print("[ModelRunner] 正在初始化模型结构...") + self.model = Qwen3NextForCausalLM(text_config) + print("[ModelRunner] 模型结构初始化完成，正在加载权重...") + load_model(self.model, config.model) + print("[ModelRunner] 权重加载完成") + else: + print("[ModelRunner] 加载纯文本模型 Qwen3 (model_type={})".format(text_model_type or "qwen3")) + print("[ModelRunner] 正在初始化模型结构...") + self.model = Qwen3ForCausalLM(text_config) + print("[ModelRunner] 模型结构初始化完成，正在加载权重...") + load_model(self.model, config.model) + print("[ModelRunner] 权重加载完成") + + embed_module = getattr(self.model, "language_model", self.model) + if hasattr(embed_module, "model"): + embed_module = embed_module.model + self.model_dtype = embed_module.embed_tokens.weight.dtype + # GatedDeltaNet now supports CUDA graph decode via persistent state buffers + # (_graph_conv_state / _graph_recurrent_state) with in-place updates that + # are correctly captured and replayed by the CUDA graph. No need to force + # eager decode for Qwen3.5 anymore. + self.sampler = Sampler() + # Initialize GDN layer references before warmup (run_model checks self._gdn_layers) + self._gdn_layers = [] + if hasattr(self.model, 'language_model') and hasattr(self.model.language_model, 'model'): + for layer in self.model.language_model.model.layers: + if hasattr(layer, 'linear_attn') and layer.linear_attn is not None: + self._gdn_layers.append(layer.linear_attn) + self._gdn_slot_manager = GDNSlotManager(max_slots=512) + self._gdn_slot_tensor = torch.zeros(512, dtype=torch.int64, device=self.device) self.warmup_model() + # Reset GatedDeltaNet states after warmup to avoid polluting real sequences + for gdn in self._gdn_layers: + gdn.reset_state() self.allocate_kv_cache() if not self.enforce_eager: self.capture_cudagraph() @@ -40,38 +167,112 @@ def __init__(self, config: Config, rank: int, event: Event | list[Event]): if self.world_size > 1: if rank == 0: - self.shm = SharedMemory(name="nanovllm", create=True, size=2**20) - dist.barrier() + try: + try: + self.shm = SharedMemory(name="nanovllm", create=True, size=2**20) + except FileExistsError: + # 共享内存已存在，先尝试清理再创建 + try: + existing_shm = SharedMemory(name="nanovllm") + existing_shm.unlink() + self.shm = SharedMemory(name="nanovllm", create=True, size=2**20) + except Exception: + # 如果清理失败，直接打开现有的 + self.shm = SharedMemory(name="nanovllm") + # 确保 barrier 前所有初始化都成功 + if dist.is_initialized(): + dist.barrier() + except Exception as e: + print(f"[R{rank}] Failed to initialize shared memory: {e}") + import traceback + traceback.print_exc() + # 尝试清理并重新抛出异常 + if dist.is_initialized(): + try: + dist.destroy_process_group() + except Exception: + pass + raise else: - dist.barrier() - self.shm = SharedMemory(name="nanovllm") - self.loop() + try: + if dist.is_initialized(): + dist.barrier() + self.shm = SharedMemory(name="nanovllm") + self.loop() + except Exception as e: + print(f"[R{rank}] Failed to initialize worker: {e}") + import traceback + traceback.print_exc() + if dist.is_initialized(): + try: + dist.destroy_process_group() + except Exception: + pass + raise def exit(self): if self.world_size > 1: - self.shm.close() - dist.barrier() - if self.rank == 0: - self.shm.unlink() - if not self.enforce_eager: + if hasattr(self, "shm"): + try: + self.shm.close() + except Exception: + pass # 忽略关闭时的异常 + if dist.is_initialized(): + try: + dist.barrier() + except Exception: + pass # 忽略 barrier 时的异常 + if self.rank == 0 and hasattr(self, "shm"): + try: + self.shm.unlink() + except (FileNotFoundError, Exception): + pass # 已经被清理或不存在 + if not self.enforce_eager and hasattr(self, "graphs"): del self.graphs, self.graph_pool - torch.cuda.synchronize() - dist.destroy_process_group() + if self.use_cuda: + torch.cuda.synchronize() + if dist.is_initialized(): + try: + dist.destroy_process_group() + except Exception: + pass # Process group 可能已经被销毁 def loop(self): while True: - method_name, args = self.read_shm() - self.call(method_name, *args) - if method_name == "exit": + try: + method_name, args = self.read_shm() + try: + self.call(method_name, *args) + except Exception as e: + # 捕获执行时的异常，记录但不崩溃 + print(f"[R{self.rank}] Error executing {method_name}: {e}") + import traceback + traceback.print_exc() + # 继续循环，等待下一个指令 + if method_name == "exit": + break + except Exception as e: + # 捕获 read_shm 或其他顶层异常 + print(f"[R{self.rank}] Fatal error in loop: {e}") + import traceback + traceback.print_exc() break def read_shm(self): assert self.world_size > 1 and self.rank > 0 self.event.wait() - n = int.from_bytes(self.shm.buf[0:4], "little") - method_name, *args = pickle.loads(self.shm.buf[4:n+4]) - self.event.clear() - return method_name, args + try: + n = int.from_bytes(self.shm.buf[0:4], "little") + # 验证数据长度合理性 + if n < 0 or n > len(self.shm.buf) - 4: + raise ValueError(f"Invalid data length: {n}") + method_name, *args = pickle.loads(self.shm.buf[4:n+4]) + self.event.clear() + return method_name, args + except Exception as e: + self.event.clear() # 确保清除 event，避免死锁 + print(f"[R{self.rank}] Error reading shared memory: {e}") + raise def write_shm(self, method_name, *args): assert self.world_size > 1 and self.rank == 0 @@ -85,45 +286,131 @@ def write_shm(self, method_name, *args): def call(self, method_name, *args): if self.world_size > 1 and self.rank == 0: self.write_shm(method_name, *args) - method = getattr(self, method_name, None) - return method(*args) + try: + method = getattr(self, method_name, None) + if method is None: + raise AttributeError(f"Method '{method_name}' not found") + return method(*args) + except Exception as e: + # 记录异常但不让进程崩溃 + print(f"[R{self.rank}] Error in call({method_name}): {e}") + import traceback + traceback.print_exc() + # 如果是 rank 0，需要通知 workers 退出，避免他们永远等待 + if self.world_size > 1 and self.rank == 0: + try: + # 尝试发送 exit 信号给 workers + self.write_shm("exit") + except Exception: + pass # 如果连 exit 都发不出去，说明系统已经严重损坏 + raise # 重新抛出异常，让上层处理 def warmup_model(self): - torch.cuda.empty_cache() - torch.cuda.reset_peak_memory_stats() + if self.use_cuda: + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() max_num_batched_tokens, max_model_len = self.config.max_num_batched_tokens, self.config.max_model_len - seq_len = min(max_num_batched_tokens, max_model_len) - num_seqs = min(max_num_batched_tokens // seq_len, self.config.max_num_seqs) - seqs = [Sequence([0] * seq_len) for _ in range(num_seqs)] - for seq in seqs: - seq.num_scheduled_tokens = seq_len + num_seqs = min(max_num_batched_tokens // max_model_len, self.config.max_num_seqs) + seqs = [Sequence([0] * max_model_len) for _ in range(num_seqs)] self.run(seqs, True) - torch.cuda.empty_cache() + if self.use_cuda: + torch.cuda.empty_cache() + + def _reset_gdn_states(self): + """Reset Qwen3.5/Qwen3Next linear-attn recurrent states if present.""" + if hasattr(self.model, "language_model") and hasattr( + self.model.language_model, "model" + ): + for layer in self.model.language_model.model.layers: + if hasattr(layer, "linear_attn") and layer.linear_attn is not None: + if hasattr(layer.linear_attn, "reset_state"): + layer.linear_attn.reset_state() + + def cleanup_seq_states(self, seq_ids: list[int]): + """Release pool slots and clean residual state_cache for finished sequences.""" + for sid in seq_ids: + slot = self._gdn_slot_manager.get_slot(sid) + if slot is not None: + for gdn in self._gdn_layers: + gdn._pool_conv_state[slot].zero_() + gdn._graph_recurrent_state[slot].zero_() + self._gdn_slot_manager.release(sid) + # Also clean residual state_cache entries (from prefill if not yet migrated) + for gdn in self._gdn_layers: + for sid in seq_ids: + gdn.state_cache.pop(sid, None) + if hasattr(gdn, '_decode_step_counter'): + gdn._decode_step_counter.pop(sid, None) def allocate_kv_cache(self): config = self.config hf_config = config.hf_config - free, total = torch.cuda.mem_get_info() - used = total - free - peak = torch.cuda.memory_stats()["allocated_bytes.all.peak"] - current = torch.cuda.memory_stats()["allocated_bytes.all.current"] - num_kv_heads = hf_config.num_key_value_heads // self.world_size - head_dim = getattr(hf_config, "head_dim", hf_config.hidden_size // hf_config.num_attention_heads) - block_bytes = 2 * hf_config.num_hidden_layers * self.block_size * num_kv_heads * head_dim * hf_config.dtype.itemsize - config.num_kvcache_blocks = int(total * config.gpu_memory_utilization - used - peak + current) // block_bytes - assert config.num_kvcache_blocks > 0 - self.kv_cache = torch.empty(2, hf_config.num_hidden_layers, config.num_kvcache_blocks, self.block_size, num_kv_heads, head_dim) + text_config = getattr(hf_config, "text_config", hf_config) + num_kv_heads = text_config.num_key_value_heads // self.world_size + head_dim = getattr(text_config, "head_dim", text_config.hidden_size // text_config.num_attention_heads) + dtype = getattr(getattr(hf_config, "text_config", hf_config), "torch_dtype", getattr(hf_config, "torch_dtype", torch.bfloat16)) + if isinstance(dtype, str): + dtype = getattr(torch, dtype, torch.bfloat16) + + # 对于混合架构（如 Qwen3-Next），只计算需要 KV cache 的层 + # 只有 full_attention 层需要 KV cache，linear_attention (Mamba/GatedDeltaNet) 不需要 + layer_types = getattr(text_config, "layer_types", None) + if layer_types is not None: + num_attention_layers = sum(1 for lt in layer_types if lt == "full_attention") + else: + num_attention_layers = text_config.num_hidden_layers + + block_bytes = 2 * num_attention_layers * self.block_size * num_kv_heads * head_dim * dtype.itemsize + if self.use_cuda: + free, total = torch.cuda.mem_get_info() + used = total - free + peak = torch.cuda.memory_stats()["allocated_bytes.all.peak"] + current = torch.cuda.memory_stats()["allocated_bytes.all.current"] + # Use simpler formula: only account for currently used memory + available = int(total * config.gpu_memory_utilization - used) + config.num_kvcache_blocks = available // block_bytes + else: + # CPU debug path: allocate enough blocks to cover max_model_len. + if config.num_kvcache_blocks == -1: + config.num_kvcache_blocks = (config.max_model_len + self.block_size - 1) // self.block_size + + assert config.num_kvcache_blocks > 0, ( + "Insufficient KV cache blocks " + f"(num_kvcache_blocks={config.num_kvcache_blocks}, block_bytes={block_bytes}, " + f"num_attention_layers={num_attention_layers}, total_layers={text_config.num_hidden_layers})" + ) + + self.kv_cache = torch.empty( + 2, + num_attention_layers, + config.num_kvcache_blocks, + self.block_size, + num_kv_heads, + head_dim, + dtype=dtype, + device=self.device, + ) layer_id = 0 for module in self.model.modules(): if hasattr(module, "k_cache") and hasattr(module, "v_cache"): + assert layer_id < num_attention_layers, ( + f"More attention layers found ({layer_id + 1}) than expected ({num_attention_layers}). " + f"This may indicate a mismatch between layer_types configuration and actual model structure." + ) module.k_cache = self.kv_cache[0, layer_id] module.v_cache = self.kv_cache[1, layer_id] layer_id += 1 + + # 验证分配的层数是否正确 + assert layer_id == num_attention_layers, ( + f"KV cache allocation mismatch: expected {num_attention_layers} attention layers, " + f"but found {layer_id} layers with k_cache/v_cache attributes." + ) def prepare_block_tables(self, seqs: list[Sequence]): max_len = max(len(seq.block_table) for seq in seqs) block_tables = [seq.block_table + [-1] * (max_len - len(seq.block_table)) for seq in seqs] - block_tables = torch.tensor(block_tables, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True) + block_tables = torch.tensor(block_tables, dtype=torch.int32, device=self.device) return block_tables def prepare_prefill(self, seqs: list[Sequence]): @@ -136,36 +423,31 @@ def prepare_prefill(self, seqs: list[Sequence]): slot_mapping = [] block_tables = None for seq in seqs: - start = seq.num_cached_tokens - seqlen_q = seq.num_scheduled_tokens - end = start + seqlen_q - seqlen_k = end - input_ids.extend(seq[start:end]) - positions.extend(range(start, end)) + seqlen = len(seq) + input_ids.extend(seq[seq.num_cached_tokens:]) + positions.extend(list(range(seq.num_cached_tokens, seqlen))) + seqlen_q = seqlen - seq.num_cached_tokens + seqlen_k = seqlen cu_seqlens_q.append(cu_seqlens_q[-1] + seqlen_q) cu_seqlens_k.append(cu_seqlens_k[-1] + seqlen_k) max_seqlen_q = max(seqlen_q, max_seqlen_q) max_seqlen_k = max(seqlen_k, max_seqlen_k) if not seq.block_table: # warmup continue - start_block = start // self.block_size - end_block = (end + self.block_size - 1) // self.block_size - for i in range(start_block, end_block): - slot_start = seq.block_table[i] * self.block_size - if i == start_block: - slot_start += start % self.block_size - if i != end_block - 1: - slot_end = seq.block_table[i] * self.block_size + self.block_size + for i in range(seq.num_cached_blocks, seq.num_blocks): + start = seq.block_table[i] * self.block_size + if i != seq.num_blocks - 1: + end = start + self.block_size else: - slot_end = seq.block_table[i] * self.block_size + end - i * self.block_size - slot_mapping.extend(range(slot_start, slot_end)) + end = start + seq.last_block_num_tokens + slot_mapping.extend(list(range(start, end))) if cu_seqlens_k[-1] > cu_seqlens_q[-1]: # prefix cache block_tables = self.prepare_block_tables(seqs) - input_ids = torch.tensor(input_ids, dtype=torch.int64, pin_memory=True).cuda(non_blocking=True) - positions = torch.tensor(positions, dtype=torch.int64, pin_memory=True).cuda(non_blocking=True) - cu_seqlens_q = torch.tensor(cu_seqlens_q, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True) - cu_seqlens_k = torch.tensor(cu_seqlens_k, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True) - slot_mapping = torch.tensor(slot_mapping, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True) + input_ids = torch.tensor(input_ids, dtype=torch.int64, device=self.device) + positions = torch.tensor(positions, dtype=torch.int64, device=self.device) + cu_seqlens_q = torch.tensor(cu_seqlens_q, dtype=torch.int32, device=self.device) + cu_seqlens_k = torch.tensor(cu_seqlens_k, dtype=torch.int32, device=self.device) + slot_mapping = torch.tensor(slot_mapping, dtype=torch.int32, device=self.device) set_context(True, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, slot_mapping, None, block_tables) return input_ids, positions @@ -179,74 +461,383 @@ def prepare_decode(self, seqs: list[Sequence]): positions.append(len(seq) - 1) context_lens.append(len(seq)) slot_mapping.append(seq.block_table[-1] * self.block_size + seq.last_block_num_tokens - 1) - input_ids = torch.tensor(input_ids, dtype=torch.int64, pin_memory=True).cuda(non_blocking=True) - positions = torch.tensor(positions, dtype=torch.int64, pin_memory=True).cuda(non_blocking=True) - slot_mapping = torch.tensor(slot_mapping, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True) - context_lens = torch.tensor(context_lens, dtype=torch.int32, pin_memory=True).cuda(non_blocking=True) + input_ids = torch.tensor(input_ids, dtype=torch.int64, device=self.device) + positions = torch.tensor(positions, dtype=torch.int64, device=self.device) + slot_mapping = torch.tensor(slot_mapping, dtype=torch.int32, device=self.device) + context_lens = torch.tensor(context_lens, dtype=torch.int32, device=self.device) block_tables = self.prepare_block_tables(seqs) set_context(False, slot_mapping=slot_mapping, context_lens=context_lens, block_tables=block_tables) return input_ids, positions def prepare_sample(self, seqs: list[Sequence]): - temperatures = [seq.temperature for seq in seqs] - temperatures = torch.tensor(temperatures, dtype=torch.float32, pin_memory=True).cuda(non_blocking=True) + temperatures = [] + for seq in seqs: + temperatures.append(seq.temperature) + temperatures = torch.tensor(temperatures, dtype=torch.float32, device=self.device) return temperatures + def _gather_gdn_states(self, sequence_ids: list[int], bs: int): + """Vectorized gather: assign pool slots and copy conv state into graph buffers.""" + if not self._gdn_layers: + return + # Allocate slots for new sequences (first decode after prefill) + slots = [] + for seq_id in sequence_ids: + slot = self._gdn_slot_manager.get_slot(seq_id) + if slot is None: + slot = self._gdn_slot_manager.allocate(seq_id) + # Migrate from state_cache (set during prefill) to pool + for gdn in self._gdn_layers: + cached = gdn.state_cache.pop(seq_id, (None, None)) + if cached[0] is not None: + gdn._pool_conv_state[slot].copy_(cached[0].squeeze(0)) + else: + gdn._pool_conv_state[slot].zero_() + if cached[1] is not None: + gdn._graph_recurrent_state[slot].copy_(cached[1].squeeze(0)) + else: + gdn._graph_recurrent_state[slot].zero_() + slots.append(slot) + + # Build index tensor + self._gdn_slot_tensor[:bs].copy_(torch.tensor(slots, dtype=torch.int64, device=self.device)) + idx = self._gdn_slot_tensor[:bs] + + # Vectorized: gather conv state + set ssm_state_indices for each layer + for gdn in self._gdn_layers: + gdn._graph_conv_state[:bs] = gdn._pool_conv_state[idx] + gdn._graph_ssm_state_indices[:bs] = idx + + def _scatter_gdn_states(self, sequence_ids: list[int], bs: int): + """Scatter conv state back to pool. Recurrent state is written in-place by Triton.""" + if not self._gdn_layers: + return + idx = self._gdn_slot_tensor[:bs] + for gdn in self._gdn_layers: + gdn._pool_conv_state[idx] = gdn._graph_conv_state[:bs] + + def _debug_prefill_logits(self, hidden_states: torch.Tensor, logits: torch.Tensor): + """Print last-token hidden norm and logits top-5 when NANOVLLM_DEBUG_LOGITS=1.""" + if self.rank != 0 or hidden_states is None or logits is None: + return + ctx = get_context() + if not getattr(ctx, "is_prefill", False) or not hasattr(ctx, "cu_seqlens_q"): + return + last_indices = ctx.cu_seqlens_q[1:] - 1 + last_h = hidden_states[last_indices] + norm_val = last_h.norm().item() + has_nan = torch.isnan(hidden_states).any().item() or torch.isnan(logits).any().item() + has_inf = torch.isinf(logits).any().item() + top5 = torch.topk(logits[0].float().cpu(), min(5, logits.shape[-1])) + print("[NANOVLLM_DEBUG_LOGITS] last_token_hidden_norm=%.6f has_nan=%s has_inf=%s" % (norm_val, has_nan, has_inf)) + print("[NANOVLLM_DEBUG_LOGITS] argmax_id=%s top5_ids=%s top5_vals=%s" % ( + logits[0].argmax().item(), + top5.indices.tolist(), + [round(v, 4) for v in top5.values.tolist()], + )) + @torch.inference_mode() - def run_model(self, input_ids: torch.Tensor, positions: torch.Tensor, is_prefill: bool): - if is_prefill or self.enforce_eager or input_ids.size(0) > 512: - return self.model.compute_logits(self.model(input_ids, positions)) - else: - bs = input_ids.size(0) - context = get_context() - graph = self.graphs[next(x for x in self.graph_bs if x >= bs)] - graph_vars = self.graph_vars - graph_vars["input_ids"][:bs] = input_ids - graph_vars["positions"][:bs] = positions - graph_vars["slot_mapping"].fill_(-1) - graph_vars["slot_mapping"][:bs] = context.slot_mapping - graph_vars["context_lens"].zero_() - graph_vars["context_lens"][:bs] = context.context_lens - graph_vars["block_tables"][:bs, :context.block_tables.size(1)] = context.block_tables - graph.replay() - return self.model.compute_logits(graph_vars["outputs"][:bs]) + def run_model( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + is_prefill: bool, + sequence_lengths: list[int] | None = None, + sequence_ids: list[int] | None = None, + vision_slices_per_seq: list[list[dict]] | None = None, + image_grid_thw_per_seq: list | None = None, + ): + if ( + is_prefill + or self.enforce_eager + or input_ids.size(0) > 512 + ): + model_kwargs = {} + # HF-style: pass sequence_lengths for vision alignment and GDN per-sequence isolation + if self.is_multimodal: + if sequence_lengths is not None: + model_kwargs["sequence_lengths"] = sequence_lengths + if sequence_ids is not None: + model_kwargs["sequence_ids"] = sequence_ids + if vision_slices_per_seq is not None: + model_kwargs["vision_slices_per_seq"] = vision_slices_per_seq + if image_grid_thw_per_seq is not None: + model_kwargs["image_grid_thw_per_seq"] = image_grid_thw_per_seq + elif ( + self._gdn_layers + and sequence_lengths is not None + and len(sequence_lengths) > 1 + ): + # Qwen3.5/GDN batch: GDN needs per-seq isolation + model_kwargs["sequence_lengths"] = sequence_lengths + if sequence_ids is not None: + model_kwargs["sequence_ids"] = sequence_ids + outputs = self.model(input_ids, positions, **model_kwargs) + # Prefill: do NOT index here. ParallelLMHead.forward() already does + # x[context.cu_seqlens_q[1:]-1] to take the last token per sequence. + # Indexing here would shrink outputs to (num_seqs, hidden), then + # lm_head would index again with full-sequence indices -> OOB. + logits = self.model.compute_logits(outputs) + if os.environ.get("NANOVLLM_DEBUG_LOGITS"): + self._debug_prefill_logits(outputs, logits) + return logits + + bs = input_ids.size(0) + context = get_context() + graph = self.graphs[next(x for x in self.graph_bs if x >= bs)] + graph_vars = self.graph_vars + graph_vars["input_ids"][:bs] = input_ids + graph_vars["positions"][:bs] = positions + graph_vars["slot_mapping"].fill_(-1) + graph_vars["slot_mapping"][:bs] = context.slot_mapping + graph_vars["context_lens"].zero_() + graph_vars["context_lens"][:bs] = context.context_lens + graph_vars["block_tables"][:bs, :context.block_tables.size(1)] = context.block_tables + # Gather per-sequence GDN states into graph buffers before replay. + if sequence_ids is not None: + self._gather_gdn_states(sequence_ids, bs) + graph.replay() + # Scatter updated GDN states back to per-sequence state_cache after replay. + if sequence_ids is not None: + self._scatter_gdn_states(sequence_ids, bs) + return self.model.compute_logits(graph_vars["outputs"][:bs]) def run(self, seqs: list[Sequence], is_prefill: bool) -> list[int]: input_ids, positions = self.prepare_prefill(seqs) if is_prefill else self.prepare_decode(seqs) + + # HF-style: per-sequence GDN state isolation. + # Use actual seq_id from each Sequence so GDN state_cache keys are + # stable across batch composition changes (e.g. when a sequence finishes + # and the remaining sequences are re-batched with different indices). + sequence_ids = [seq.seq_id for seq in seqs] + if is_prefill: + sequence_lengths = [len(seq) - seq.num_cached_tokens for seq in seqs] + else: + sequence_lengths = [1] * len(seqs) # decode: 1 token per sequence + vision_slices_per_seq = None + image_grid_thw_per_seq = None + if is_prefill and self.is_multimodal: + image_grid_thw_per_seq = [getattr(seq, "image_grid_thw", None) for seq in seqs] + image_token_id = getattr(self.config.hf_config, "image_token_id", None) + if image_token_id is not None: + seq_stats = [] + for seq in seqs: + seq_stats.append( + { + "seq_id": seq.seq_id, + "len": len(seq), + "cached": seq.num_cached_tokens, + "placeholder_count": seq.token_ids.count(image_token_id), + "vision_counts": seq.vision_counts, + "vision_consumed": seq.vision_consumed, + } + ) + # print("[ModelRunner] prefill seq stats", seq_stats) + + vision_slices_per_seq = [] + has_slices = False + for seq in seqs: + self._ensure_vision_cache(seq) + slices_for_seq: list[dict] = [] + window_start = seq.num_cached_tokens + window_end = len(seq) + + for placeholder_idx, (offset, length) in enumerate(seq.vision_placeholders): + if placeholder_idx >= len(seq.vision_counts): + continue + consumed = seq.vision_consumed[placeholder_idx] + total_len = length + if consumed >= total_len: + continue + + range_start = offset + range_end = offset + total_len + + overlap_start = max(range_start, window_start) + overlap_end = min(range_end, window_end) + if overlap_end <= overlap_start: + continue + + slice_offset = max(consumed, overlap_start - range_start) + remaining = total_len - slice_offset + overlap_available = overlap_end - overlap_start + take = min(remaining, overlap_available) + if take <= 0: + continue + + target_offset = overlap_start - window_start + + chunk_tokens = seq.cached_vision_tokens[placeholder_idx] + token_slice = chunk_tokens[ + slice_offset : slice_offset + take + ].to( + device="cuda", + dtype=self.model_dtype, + non_blocking=True, + ).contiguous() + + debug_mean = float(token_slice.float().abs().mean().item()) + # print( + # "[ModelRunner] slice", + # { + # "seq_id": seq.seq_id, + # "placeholder": placeholder_idx, + # "slice_offset": slice_offset, + # "length": take, + # "target_offset": target_offset, + # "mean_abs": round(debug_mean, 6), + # }, + # ) + + slices_for_seq.append( + { + "tokens": token_slice, + "length": take, + "target_offset": target_offset, + "placeholder_idx": placeholder_idx, + } + ) + has_slices = True + + vision_slices_per_seq.append(slices_for_seq) + + if not has_slices: + vision_slices_per_seq = None + + def _advance_vision_offsets(): + if not is_prefill or not self.is_multimodal: + return + if vision_slices_per_seq is None: + return + for seq, slices in zip(seqs, vision_slices_per_seq): + for slice_info in slices: + length = slice_info["length"] + placeholder_idx = slice_info["placeholder_idx"] + if placeholder_idx < len(seq.vision_consumed): + span = seq.vision_placeholders[placeholder_idx][1] + seq.vision_consumed[placeholder_idx] += length + seq.vision_consumed[placeholder_idx] = min( + seq.vision_consumed[placeholder_idx], span + ) + if seq.vision_placeholders: + all_consumed = all( + seq.vision_consumed[idx] >= span + for idx, (_, span) in enumerate(seq.vision_placeholders) + ) + else: + all_consumed = True + if all_consumed: + seq.cached_vision_tokens = None + temperatures = self.prepare_sample(seqs) if self.rank == 0 else None - logits = self.run_model(input_ids, positions, is_prefill) + # vLLM-style: batch handling is done via cu_seqlens in global context + logits = self.run_model( + input_ids, + positions, + is_prefill, + sequence_lengths=sequence_lengths, + sequence_ids=sequence_ids, + vision_slices_per_seq=vision_slices_per_seq, + image_grid_thw_per_seq=image_grid_thw_per_seq, + ) + _advance_vision_offsets() token_ids = self.sampler(logits, temperatures).tolist() if self.rank == 0 else None reset_context() return token_ids @torch.inference_mode() def capture_cudagraph(self): + if not self.use_cuda: + # CPU mode: CUDA graph replay isn't applicable. + return config = self.config hf_config = config.hf_config + text_config = getattr(hf_config, "text_config", hf_config) + max_bs = min(self.config.max_num_seqs, 512) - max_num_blocks = (config.max_model_len + self.block_size - 1) // self.block_size + max_num_blocks = ( + config.max_model_len + self.block_size - 1 + ) // self.block_size + input_ids = torch.zeros(max_bs, dtype=torch.int64) positions = torch.zeros(max_bs, dtype=torch.int64) slot_mapping = torch.zeros(max_bs, dtype=torch.int32) context_lens = torch.zeros(max_bs, dtype=torch.int32) - block_tables = torch.zeros(max_bs, max_num_blocks, dtype=torch.int32) - outputs = torch.zeros(max_bs, hf_config.hidden_size) - self.graph_bs = [1, 2, 4, 8] + list(range(16, max_bs + 1, 16)) + block_tables = torch.zeros( + max_bs, max_num_blocks, dtype=torch.int32 + ) + + hidden_size = getattr(text_config, "hidden_size", None) + if hidden_size is None: + # Fallback: infer from token embedding dim. + embed_module = getattr( + self.model, "language_model", self.model + ) + if hasattr(embed_module, "model"): + embed_module = embed_module.model + hidden_size = int( + embed_module.embed_tokens.weight.shape[-1] + ) + + outputs = torch.zeros(max_bs, hidden_size) + self.graph_bs = [1, 2, 4, 8] + list( + range(16, max_bs + 1, 16) + ) + # Only capture graphs that fit the allocated dummy buffers. + # When max_bs is small (e.g. max_num_seqs=1), we must not capture + # larger bs values because input_ids[:bs] will still have length + # max_bs, but sequence_lengths would have length bs, causing empty + # sequence slices inside Qwen3.5 linear_attention. + self.graph_bs = [x for x in self.graph_bs if x <= max_bs] self.graphs = {} self.graph_pool = None for bs in reversed(self.graph_bs): + # Capture uses dummy tokens; ensure it doesn't leak recurrent states + # into subsequent captures or real inference. + self._reset_gdn_states() + # Ensure identity mapping for ssm_state_indices during capture + for gdn in self._gdn_layers: + gdn._graph_ssm_state_indices.copy_( + torch.arange(gdn._graph_ssm_state_indices.shape[0], dtype=torch.int64, device=gdn._graph_ssm_state_indices.device) + ) graph = torch.cuda.CUDAGraph() - set_context(False, slot_mapping=slot_mapping[:bs], context_lens=context_lens[:bs], block_tables=block_tables[:bs]) - outputs[:bs] = self.model(input_ids[:bs], positions[:bs]) # warmup + set_context( + False, + slot_mapping=slot_mapping[:bs], + context_lens=context_lens[:bs], + block_tables=block_tables[:bs], + ) + # Decode-only: each "sequence" in the batch is exactly one token. + # Qwen3.5/GDN uses `sequence_lengths/sequence_ids` to provide per-seq + # recurrent state isolation; if we omit it here, CUDA graph replay + # will use the wrong recurrence semantics and can produce repeated + # tokens / garbled outputs. + extra_kwargs = {} + if self._gdn_layers: + extra_kwargs["sequence_lengths"] = [1] * bs + extra_kwargs["use_graph"] = True + outputs[:bs] = self.model( + input_ids[:bs], + positions[:bs], + **extra_kwargs, + ) # warmup with torch.cuda.graph(graph, self.graph_pool): - outputs[:bs] = self.model(input_ids[:bs], positions[:bs]) # capture + outputs[:bs] = self.model( + input_ids[:bs], + positions[:bs], + **extra_kwargs, + ) # capture if self.graph_pool is None: self.graph_pool = graph.pool() self.graphs[bs] = graph - torch.cuda.synchronize() + if self.use_cuda: + torch.cuda.synchronize() reset_context() + # Final cleanup before serving real requests. + self._reset_gdn_states() + self.graph_vars = dict( input_ids=input_ids, positions=positions, @@ -255,3 +846,26 @@ def capture_cudagraph(self): block_tables=block_tables, outputs=outputs, ) + + def _ensure_vision_cache(self, seq: Sequence): + if seq.cached_vision_tokens is not None: + return + if seq.pixel_values is None or seq.image_grid_thw is None: + seq.cached_vision_tokens = [] + return + pixel = seq.pixel_values.to( + device=self.device, + dtype=self.model_dtype, + non_blocking=True, + ).contiguous() + grid = seq.image_grid_thw.to( + device=self.device, + dtype=torch.int32, + non_blocking=True, + ).contiguous() + image_tokens = self.model.visual(pixel, grid) + seq.cached_vision_tokens = [ + image_tokens[i].detach().cpu() for i in range(image_tokens.size(0)) + ] + seq.pixel_values = None + seq.image_grid_thw = None diff --git a/nanovllm/engine/scheduler.py b/nanovllm/engine/scheduler.py index d15979d3b..4984db8a1 100644 --- a/nanovllm/engine/scheduler.py +++ b/nanovllm/engine/scheduler.py @@ -11,8 +11,17 @@ def __init__(self, config: Config): self.max_num_seqs = config.max_num_seqs self.max_num_batched_tokens = config.max_num_batched_tokens self.eos = config.eos - self.block_size = config.kvcache_block_size - self.block_manager = BlockManager(config.num_kvcache_blocks, config.kvcache_block_size) + hf_config = config.hf_config or {} + non_cache_token_ids = set() + for token_attr in ("image_token_id", "vision_start_token_id", "vision_end_token_id"): + token_id = getattr(hf_config, token_attr, None) + if token_id is not None: + non_cache_token_ids.add(int(token_id)) + self.block_manager = BlockManager( + config.num_kvcache_blocks, + config.kvcache_block_size, + non_cache_token_ids=non_cache_token_ids, + ) self.waiting: deque[Sequence] = deque() self.running: deque[Sequence] = deque() @@ -23,39 +32,27 @@ def add(self, seq: Sequence): self.waiting.append(seq) def schedule(self) -> tuple[list[Sequence], bool]: + # prefill scheduled_seqs = [] + num_seqs = 0 num_batched_tokens = 0 - - # prefill - while self.waiting and len(scheduled_seqs) < self.max_num_seqs: + while self.waiting and num_seqs < self.max_num_seqs: seq = self.waiting[0] - remaining = self.max_num_batched_tokens - num_batched_tokens - if remaining == 0: + num_cached_blocks = self.block_manager.can_allocate(seq) + if num_batched_tokens + len(seq) > self.max_num_batched_tokens or num_cached_blocks == -1: break - if not seq.block_table: - num_cached_blocks = self.block_manager.can_allocate(seq) - if num_cached_blocks == -1: - break - num_tokens = seq.num_tokens - num_cached_blocks * self.block_size - else: - num_tokens = seq.num_tokens - seq.num_cached_tokens - if remaining < num_tokens and scheduled_seqs: # only allow chunked prefill for the first seq - break - if not seq.block_table: - self.block_manager.allocate(seq, num_cached_blocks) - seq.num_scheduled_tokens = min(num_tokens, remaining) - num_batched_tokens += seq.num_scheduled_tokens - if seq.num_cached_tokens + seq.num_scheduled_tokens == seq.num_tokens: - seq.status = SequenceStatus.RUNNING - self.waiting.popleft() - self.running.append(seq) + num_seqs += 1 + self.block_manager.allocate(seq, num_cached_blocks) + num_batched_tokens += len(seq) - seq.num_cached_tokens + seq.status = SequenceStatus.RUNNING + self.waiting.popleft() + self.running.append(seq) scheduled_seqs.append(seq) - if scheduled_seqs: return scheduled_seqs, True # decode - while self.running and len(scheduled_seqs) < self.max_num_seqs: + while self.running and num_seqs < self.max_num_seqs: seq = self.running.popleft() while not self.block_manager.can_append(seq): if self.running: @@ -64,8 +61,7 @@ def schedule(self) -> tuple[list[Sequence], bool]: self.preempt(seq) break else: - seq.num_scheduled_tokens = 1 - seq.is_prefill = False + num_seqs += 1 self.block_manager.may_append(seq) scheduled_seqs.append(seq) assert scheduled_seqs @@ -74,17 +70,11 @@ def schedule(self) -> tuple[list[Sequence], bool]: def preempt(self, seq: Sequence): seq.status = SequenceStatus.WAITING - seq.is_prefill = True self.block_manager.deallocate(seq) self.waiting.appendleft(seq) - def postprocess(self, seqs: list[Sequence], token_ids: list[int], is_prefill: bool): + def postprocess(self, seqs: list[Sequence], token_ids: list[int]) -> list[bool]: for seq, token_id in zip(seqs, token_ids): - self.block_manager.hash_blocks(seq) - seq.num_cached_tokens += seq.num_scheduled_tokens - seq.num_scheduled_tokens = 0 - if is_prefill and seq.num_cached_tokens < seq.num_tokens: - continue seq.append_token(token_id) if (not seq.ignore_eos and token_id == self.eos) or seq.num_completion_tokens == seq.max_tokens: seq.status = SequenceStatus.FINISHED diff --git a/nanovllm/engine/sequence.py b/nanovllm/engine/sequence.py index 4decfce5f..b9aa0979b 100644 --- a/nanovllm/engine/sequence.py +++ b/nanovllm/engine/sequence.py @@ -2,6 +2,8 @@ from enum import Enum, auto from itertools import count +import torch + from nanovllm.sampling_params import SamplingParams @@ -15,7 +17,9 @@ class Sequence: block_size = 256 counter = count() - def __init__(self, token_ids: list[int], sampling_params = SamplingParams()): + def __init__(self, token_ids: list[int], sampling_params = SamplingParams(), + images=None, pixel_values=None, image_grid_thw=None, + vision_counts=None, vision_placeholders=None): self.seq_id = next(Sequence.counter) self.status = SequenceStatus.WAITING self.token_ids = copy(token_ids) @@ -30,6 +34,16 @@ def __init__(self, token_ids: list[int], sampling_params = SamplingParams()): self.max_tokens = sampling_params.max_tokens self.ignore_eos = sampling_params.ignore_eos + # Multimodal fields + self.images = images + self.pixel_values = pixel_values + self.image_grid_thw = image_grid_thw + self.vision_placeholders = vision_placeholders or [] + self.vision_counts = vision_counts or [] + self.vision_consumed = [0] * len(self.vision_placeholders) + self.cached_vision_tokens = None + self.vision_offset = 0 + def __len__(self): return self.num_tokens @@ -56,6 +70,10 @@ def completion_token_ids(self): def num_blocks(self): return (self.num_tokens + self.block_size - 1) // self.block_size + @property + def num_cached_blocks(self): + return self.num_cached_tokens // self.block_size + @property def last_block_num_tokens(self): return self.num_tokens - (self.num_blocks - 1) * self.block_size diff --git a/nanovllm/layers/attention.py b/nanovllm/layers/attention.py index e416139ea..145f21fe0 100644 --- a/nanovllm/layers/attention.py +++ b/nanovllm/layers/attention.py @@ -1,43 +1,96 @@ import torch from torch import nn -import triton -import triton.language as tl -from flash_attn import flash_attn_varlen_func, flash_attn_with_kvcache +# Triton is optional. When it's not available, we fall back to a torch KV-cache writer. +try: + import triton # type: ignore + import triton.language as tl # type: ignore + _HAS_TRITON = True +except ModuleNotFoundError: + triton = None + tl = None + _HAS_TRITON = False + +# flash-attn is optional. When it's unavailable, we fall back to a torch attention implementation. +try: + from flash_attn import flash_attn_varlen_func, flash_attn_with_kvcache # type: ignore + _HAS_FLASH_ATTN = True +except ModuleNotFoundError: + flash_attn_varlen_func = None + flash_attn_with_kvcache = None + _HAS_FLASH_ATTN = False from nanovllm.utils.context import get_context -@triton.jit -def store_kvcache_kernel( - key_ptr, - key_stride, - value_ptr, - value_stride, - k_cache_ptr, - v_cache_ptr, - slot_mapping_ptr, - D: tl.constexpr, +if _HAS_TRITON: + @triton.jit + def store_kvcache_kernel( + key_ptr, + key_stride, + value_ptr, + value_stride, + k_cache_ptr, + v_cache_ptr, + slot_mapping_ptr, + D: tl.constexpr, + ): + idx = tl.program_id(0) + slot = tl.load(slot_mapping_ptr + idx) + if slot == -1: + return + key_offsets = idx * key_stride + tl.arange(0, D) + value_offsets = idx * value_stride + tl.arange(0, D) + key = tl.load(key_ptr + key_offsets) + value = tl.load(value_ptr + value_offsets) + cache_offsets = slot * D + tl.arange(0, D) + tl.store(k_cache_ptr + cache_offsets, key) + tl.store(v_cache_ptr + cache_offsets, value) + + +def store_kvcache( + key: torch.Tensor, + value: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + slot_mapping: torch.Tensor, ): - idx = tl.program_id(0) - slot = tl.load(slot_mapping_ptr + idx) - if slot == -1: return - key_offsets = idx * key_stride + tl.arange(0, D) - value_offsets = idx * value_stride + tl.arange(0, D) - key = tl.load(key_ptr + key_offsets) - value = tl.load(value_ptr + value_offsets) - cache_offsets = slot * D + tl.arange(0, D) - tl.store(k_cache_ptr + cache_offsets, key) - tl.store(v_cache_ptr + cache_offsets, value) - - -def store_kvcache(key: torch.Tensor, value: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, slot_mapping: torch.Tensor): + """ + Write (K, V) into KV cache at positions given by `slot_mapping`. + + key/value: (N, num_kv_heads, head_dim) + slot_mapping: (N,) where each entry is an absolute slot index, or -1 for "ignore". + """ N, num_heads, head_dim = key.shape D = num_heads * head_dim - assert key.stride(-1) == 1 and value.stride(-1) == 1 - assert key.stride(1) == head_dim and value.stride(1) == head_dim - assert k_cache.stride(1) == D and v_cache.stride(1) == D assert slot_mapping.numel() == N - store_kvcache_kernel[(N,)](key, key.stride(0), value, value.stride(0), k_cache, v_cache, slot_mapping, D) + + if _HAS_TRITON: + assert key.stride(-1) == 1 and value.stride(-1) == 1 + assert key.stride(1) == head_dim and value.stride(1) == head_dim + assert k_cache.stride(1) == D and v_cache.stride(1) == D + store_kvcache_kernel[(N,)]( + key, + key.stride(0), + value, + value.stride(0), + k_cache, + v_cache, + slot_mapping, + D, + ) + return + + # Torch fallback: flatten caches to (num_slots, D) and assign. + key_flat = key.reshape(N, D) + value_flat = value.reshape(N, D) + k_cache_flat = k_cache.contiguous().view(-1, D) + v_cache_flat = v_cache.contiguous().view(-1, D) + + mask = slot_mapping != -1 + if mask.any(): + slots = slot_mapping[mask].to(torch.long) + k_cache_flat[slots] = key_flat[mask] + v_cache_flat[slots] = value_flat[mask] class Attention(nn.Module): @@ -62,14 +115,122 @@ def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor): if k_cache.numel() and v_cache.numel(): store_kvcache(k, v, k_cache, v_cache, context.slot_mapping) if context.is_prefill: - if context.block_tables is not None: # prefix cache - k, v = k_cache, v_cache - o = flash_attn_varlen_func(q, k, v, - max_seqlen_q=context.max_seqlen_q, cu_seqlens_q=context.cu_seqlens_q, - max_seqlen_k=context.max_seqlen_k, cu_seqlens_k=context.cu_seqlens_k, - softmax_scale=self.scale, causal=True, block_table=context.block_tables) - else: # decode - o = flash_attn_with_kvcache(q.unsqueeze(1), k_cache, v_cache, - cache_seqlens=context.context_lens, block_table=context.block_tables, - softmax_scale=self.scale, causal=True) + if _HAS_FLASH_ATTN: + if context.block_tables is not None: # prefix cache + k, v = k_cache, v_cache + o = flash_attn_varlen_func( + q, + k, + v, + max_seqlen_q=context.max_seqlen_q, + cu_seqlens_q=context.cu_seqlens_q, + max_seqlen_k=context.max_seqlen_k, + cu_seqlens_k=context.cu_seqlens_k, + softmax_scale=self.scale, + causal=True, + block_table=context.block_tables, + ) + else: + # Correctness-oriented torch fallback. + # Note: this path currently doesn't support prefix-cache (block_tables != None). + if context.block_tables is not None: + raise NotImplementedError( + "Torch flash_attn fallback does not support prefix-cache (block_tables != None). " + "Install flash-attn or disable prefix caching for this debug run." + ) + assert context.cu_seqlens_q is not None and context.cu_seqlens_k is not None + cu_q = context.cu_seqlens_q.tolist() + cu_k = context.cu_seqlens_k.tolist() + assert len(cu_q) == len(cu_k), "cu_seqlens_q/cu_seqlens_k length mismatch" + + outs = [] + # q: (total_q, num_heads, head_dim) + # k/v: (total_k, num_kv_heads, head_dim) + for start_q, end_q, start_k, end_k in zip(cu_q[:-1], cu_q[1:], cu_k[:-1], cu_k[1:]): + q_i = q[start_q:end_q] # (L, H, D) + k_i = k[start_k:end_k] # (L, Hk, D) + v_i = v[start_k:end_k] # (L, Hk, D) + + if self.num_heads != self.num_kv_heads: + if self.num_heads % self.num_kv_heads != 0: + raise ValueError("num_heads must be a multiple of num_kv_heads for torch fallback") + group = self.num_heads // self.num_kv_heads + k_i = k_i.repeat_interleave(group, dim=1) + v_i = v_i.repeat_interleave(group, dim=1) + + # scaled_dot_product_attention expects (B, H, L, D) + q_t = q_i.transpose(0, 1).unsqueeze(0) + k_t = k_i.transpose(0, 1).unsqueeze(0) + v_t = v_i.transpose(0, 1).unsqueeze(0) + + q_dtype = q_t.dtype + if q_t.device.type != "cuda" and q_dtype in (torch.float16, torch.bfloat16): + q_t = q_t.float() + k_t = k_t.float() + v_t = v_t.float() + + attn = torch.nn.functional.scaled_dot_product_attention( + q_t, k_t, v_t, attn_mask=None, dropout_p=0.0, is_causal=True + ) + out_i = attn.squeeze(0).transpose(0, 1).to(q_dtype) + outs.append(out_i) + o = torch.cat(outs, dim=0) + else: + # decode + if _HAS_FLASH_ATTN: + o = flash_attn_with_kvcache( + q.unsqueeze(1), + k_cache, + v_cache, + cache_seqlens=context.context_lens, + block_table=context.block_tables, + softmax_scale=self.scale, + causal=True, + ) + else: + # Correctness-oriented torch fallback for decode. + assert context.context_lens is not None and context.block_tables is not None + bs = q.shape[0] + block_size = k_cache.shape[1] + k_cache_flat = k_cache.contiguous().view(-1, self.num_kv_heads, self.head_dim) + v_cache_flat = v_cache.contiguous().view(-1, self.num_kv_heads, self.head_dim) + + outs = [] + for i in range(bs): + seqlen = int(context.context_lens[i].item()) + block_ids = context.block_tables[i].tolist() + block_ids = [bid for bid in block_ids if bid != -1] + if not block_ids: + raise RuntimeError("decode torch fallback: empty block_ids") + + # token positions [0..seqlen-1] map to slots: block_id * block_size + offset + pos = torch.arange(seqlen, device=q.device, dtype=torch.long) + block_idx = (pos // block_size).to(torch.long) + within = (pos % block_size).to(torch.long) + + block_id_tensor = torch.tensor(block_ids, device=q.device, dtype=torch.long) + slot_indices = block_id_tensor[block_idx] * block_size + within + + K = k_cache_flat[slot_indices] # (seqlen, Hk, D) + V = v_cache_flat[slot_indices] # (seqlen, Hk, D) + if self.num_heads != self.num_kv_heads: + if self.num_heads % self.num_kv_heads != 0: + raise ValueError("num_heads must be a multiple of num_kv_heads for torch fallback") + group = self.num_heads // self.num_kv_heads + K = K.repeat_interleave(group, dim=1) + V = V.repeat_interleave(group, dim=1) + + q_i = q[i] # (H, D) + q_dtype = q_i.dtype + if q_i.device.type != "cuda" and q_dtype in (torch.float16, torch.bfloat16): + q_i = q_i.float() + K = K.float() + V = V.float() + + # scores: (H, seqlen) + scores = torch.einsum("hd,thd->ht", q_i, K) * self.scale + attn = torch.softmax(scores.float(), dim=-1).to(V.dtype) + out_i = torch.einsum("ht,thd->hd", attn, V).to(q_dtype) + outs.append(out_i) + o = torch.stack(outs, dim=0) return o diff --git a/nanovllm/layers/ops/__init__.py b/nanovllm/layers/ops/__init__.py new file mode 100644 index 000000000..07e74c9dd --- /dev/null +++ b/nanovllm/layers/ops/__init__.py @@ -0,0 +1,29 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang + +""" +Optimized GatedDeltaNet operations using Triton kernels. +""" + +import torch + +# Set Triton allocator to use PyTorch CUDA caching allocator +try: + import triton + triton.set_allocator(lambda size, device, stream: torch.cuda.caching_allocator_alloc(size)) +except Exception: + pass + +try: + from nanovllm.layers.ops.fused_recurrent import fused_recurrent_gated_delta_rule + __all__ = ["fused_recurrent_gated_delta_rule"] +except ImportError: + fused_recurrent_gated_delta_rule = None + __all__ = [] + +try: + from nanovllm.layers.ops.chunk import chunk_gated_delta_rule + __all__.append("chunk_gated_delta_rule") +except ImportError: + chunk_gated_delta_rule = None diff --git a/nanovllm/layers/ops/chunk.py b/nanovllm/layers/ops/chunk.py new file mode 100644 index 000000000..deca7774d --- /dev/null +++ b/nanovllm/layers/ops/chunk.py @@ -0,0 +1,228 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang +# +# This file contains code copied from the flash-linear-attention project. +# The original source code was licensed under the MIT license and included +# the following copyright notice: +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang +# ruff: noqa: E501 +import torch +from einops import rearrange + +from nanovllm.layers.ops.chunk_delta_h import chunk_gated_delta_rule_fwd_h +from nanovllm.layers.ops.chunk_o import chunk_fwd_o +from nanovllm.layers.ops.chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd +from nanovllm.layers.ops.cumsum import chunk_local_cumsum +from nanovllm.layers.ops.l2norm import l2norm_fwd +from nanovllm.layers.ops.solve_tril import solve_tril +from nanovllm.layers.ops.utils import SUPPRESS_LEVEL, input_guard +from nanovllm.layers.ops.wy_fast import recompute_w_u_fwd + + +def chunk_gated_delta_rule_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.LongTensor | None = None, +): + g = chunk_local_cumsum(g, chunk_size=64, cu_seqlens=cu_seqlens, head_first=False) + # obtain WY representation. u is actually the new v. + A = chunk_scaled_dot_kkt_fwd( + k=k, beta=beta, g=g, cu_seqlens=cu_seqlens, output_dtype=torch.float32 + ) + A = solve_tril(A=A, cu_seqlens=cu_seqlens, output_dtype=k.dtype) + w, u = recompute_w_u_fwd( + k=k, + v=v, + beta=beta, + A=A, + g_cumsum=g, + cu_seqlens=cu_seqlens, + ) + h, v_new, final_state = chunk_gated_delta_rule_fwd_h( + k=k, + w=w, + u=u, + g=g, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + ) + o = chunk_fwd_o( + q=q, + k=k, + v=v_new, + h=h, + g=g, + scale=scale, + cu_seqlens=cu_seqlens, + ) + if SUPPRESS_LEVEL < 3: + return g, o, A, final_state, None, None, None + elif SUPPRESS_LEVEL >= 3: + return g, o, A, final_state, w, h, v_new + + +class ChunkGatedDeltaRuleFunction(torch.autograd.Function): + @staticmethod + @input_guard + @torch.amp.custom_fwd(device_type="cuda") + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool, + cu_seqlens: torch.LongTensor | None = None, + use_qk_l2norm_in_kernel: bool = False, + ): + if use_qk_l2norm_in_kernel: + q = l2norm_fwd(q) + k = l2norm_fwd(k) + + g, o, A, final_state, w, h, v_new = chunk_gated_delta_rule_fwd( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + cu_seqlens=cu_seqlens, + ) + ctx.scale = scale + ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel + return o.to(q.dtype), final_state + + +@torch.compiler.disable +def chunk_gated_delta_rule( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + cu_seqlens: torch.LongTensor | None = None, + use_qk_l2norm_in_kernel: bool = False, +): + r""" + Chunk-based Gated Delta Rule implementation. + Only supports time-first format: (B, T, H, K). + + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, H, V]`. + g (torch.Tensor): + (forget) gating tensor (in log space!) of shape `[B, T, H]`. + beta (torch.Tensor): + betas of shape `[B, T, H]`. + scale (Optional[int]): + Scale factor for the RetNet attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[N, H, V, K]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, H, V, K]`. Default: `False`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + use_qk_l2norm_in_kernel (bool): + Whether to use L2 normalization for q and k in the kernel. Default: `False`. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, H, V]`. + final_state (torch.Tensor): + Final state of shape `[N, H, V, K]` if `output_final_state=True` else `None`. + + Examples:: + >>> import torch + >>> import torch.nn.functional as F + >>> from einops import rearrange + >>> from fla.ops.gated_delta_rule import chunk_gated_delta_rule + # inputs with equal lengths + >>> B, T, H, K, V = 4, 2048, 4, 512, 512 + >>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda') + >>> k = F.normalize(torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda'), p=2, dim=-1) + >>> v = torch.randn(B, T, H, V, dtype=torch.bfloat16, device='cuda') + >>> beta = torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda').sigmoid() + >>> g = F.logsigmoid(torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda')) + >>> h0 = torch.randn(B, H, V, K, dtype=torch.bfloat16, device='cuda') + >>> o, ht = chunk_gated_delta_rule( + q, k, v, g, beta, + initial_state=h0, + output_final_state=True + ) + # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required + >>> q, k, v, beta, g = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, beta, g)) + # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected + >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) + >>> o_var, ht_var = chunk_gated_delta_rule( + q, k, v, g, beta, + initial_state=h0, + output_final_state=True, + cu_seqlens=cu_seqlens + ) + """ + assert q.dtype == k.dtype == v.dtype + assert q.dtype != torch.float32, ( + "ChunkGatedDeltaRuleFunction does not support float32. Please use bfloat16." + ) + assert len(beta.shape) == 3, ( + "beta must be of shape [B, T, H]." + ) + + # Validate input format: must be (B, T, H, K) format + assert len(q.shape) == 4, f"q must be 4D tensor (B, T, H, K), got shape {q.shape}" + assert len(k.shape) == 4, f"k must be 4D tensor (B, T, H, K), got shape {k.shape}" + assert len(v.shape) == 4, f"v must be 4D tensor (B, T, H, V), got shape {v.shape}" + assert q.shape[0] == k.shape[0] == v.shape[0], "Batch size mismatch" + assert q.shape[1] == k.shape[1] == v.shape[1], "Sequence length mismatch" + assert q.shape[2] == k.shape[2], "Number of heads mismatch between q and k" + if cu_seqlens is not None: + if q.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing." + ) + if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}." + ) + if scale is None: + scale = k.shape[-1] ** -0.5 + o, final_state = ChunkGatedDeltaRuleFunction.apply( + q, + k, + v, + g, + beta, + scale, + initial_state, + output_final_state, + cu_seqlens, + use_qk_l2norm_in_kernel, + ) + # Output is always in (B, T, H, V) format + return o, final_state diff --git a/nanovllm/layers/ops/chunk_delta_h.py b/nanovllm/layers/ops/chunk_delta_h.py new file mode 100644 index 000000000..92c09d7c4 --- /dev/null +++ b/nanovllm/layers/ops/chunk_delta_h.py @@ -0,0 +1,345 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang +# +# This file contains code copied from the flash-linear-attention project. +# The original source code was licensed under the MIT license and included +# the following copyright notice: +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang +# ruff: noqa: E501 + +import torch + +import triton +import triton.language as tl + +from nanovllm.layers.ops.index import prepare_chunk_indices, prepare_chunk_offsets +from nanovllm.layers.ops.op import exp +from nanovllm.layers.ops.utils import use_cuda_graph + +NUM_WARPS = [2, 4, 8, 16] + + +@triton.heuristics( + { + "USE_G": lambda args: args["g"] is not None, + "USE_GK": lambda args: args["gk"] is not None, + "USE_INITIAL_STATE": lambda args: args["h0"] is not None, + "STORE_FINAL_STATE": lambda args: args["ht"] is not None, + "SAVE_NEW_VALUE": lambda args: args["v_new"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +@triton.autotune( + configs=[ + triton.Config({"BV": BV}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4] + for num_stages in [2, 3, 4] + for BV in [32, 64] + ], + key=["H", "K", "V", "BT"], + use_cuda_graph=use_cuda_graph, +) +@triton.jit(do_not_specialize=["T"]) +def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( + k, + v, + w, + v_new, + g, + gk, + h, + h0, + ht, + cu_seqlens, + chunk_offsets, + T, + H: tl.constexpr, + Hg: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + USE_GK: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, + STORE_FINAL_STATE: tl.constexpr, + SAVE_NEW_VALUE: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_nh = tl.program_id(0), tl.program_id(1) + i_n, i_h = i_nh // H, i_nh % H + if IS_VARLEN: + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + NT = tl.cdiv(T, BT) + boh = tl.load(chunk_offsets + i_n).to(tl.int32) + else: + bos, eos = i_n * T, i_n * T + T + NT = tl.cdiv(T, BT) + boh = i_n * NT + + # [BV, BK] + b_h1 = tl.zeros([BV, 64], dtype=tl.float32) + if K > 64: + b_h2 = tl.zeros([BV, 64], dtype=tl.float32) + if K > 128: + b_h3 = tl.zeros([BV, 64], dtype=tl.float32) + if K > 192: + b_h4 = tl.zeros([BV, 64], dtype=tl.float32) + + # calculate offset + h += ((boh * H + i_h) * V * K).to(tl.int64) + v += ((bos * H + i_h) * V).to(tl.int64) + k += ((bos * Hg + i_h // (H // Hg)) * K).to(tl.int64) + w += ((bos * H + i_h) * K).to(tl.int64) + if SAVE_NEW_VALUE: + v_new += ((bos * H + i_h) * V).to(tl.int64) + stride_v = H * V + stride_h = H * V * K + stride_k = Hg * K + stride_w = H * K + if USE_INITIAL_STATE: + h0 = h0 + i_nh * V * K + if STORE_FINAL_STATE: + ht = ht + i_nh * V * K + + # load initial state + if USE_INITIAL_STATE: + p_h0_1 = tl.make_block_ptr(h0, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)) + b_h1 += tl.load(p_h0_1, boundary_check=(0, 1)).to(tl.float32) + if K > 64: + p_h0_2 = tl.make_block_ptr( + h0, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0) + ) + b_h2 += tl.load(p_h0_2, boundary_check=(0, 1)).to(tl.float32) + if K > 128: + p_h0_3 = tl.make_block_ptr( + h0, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0) + ) + b_h3 += tl.load(p_h0_3, boundary_check=(0, 1)).to(tl.float32) + if K > 192: + p_h0_4 = tl.make_block_ptr( + h0, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0) + ) + b_h4 += tl.load(p_h0_4, boundary_check=(0, 1)).to(tl.float32) + + # main recurrence + for i_t in range(NT): + p_h1 = tl.make_block_ptr( + h + i_t * stride_h, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0) + ) + tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), boundary_check=(0, 1)) + if K > 64: + p_h2 = tl.make_block_ptr( + h + i_t * stride_h, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0) + ) + tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), boundary_check=(0, 1)) + if K > 128: + p_h3 = tl.make_block_ptr( + h + i_t * stride_h, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0) + ) + tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), boundary_check=(0, 1)) + if K > 192: + p_h4 = tl.make_block_ptr( + h + i_t * stride_h, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0) + ) + tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), boundary_check=(0, 1)) + + p_w = tl.make_block_ptr( + w, (T, K), (stride_w, 1), (i_t * BT, 0), (BT, 64), (1, 0) + ) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_v = tl.dot(b_w, tl.trans(b_h1).to(b_w.dtype)) + if K > 64: + p_w = tl.make_block_ptr( + w, (T, K), (stride_w, 1), (i_t * BT, 64), (BT, 64), (1, 0) + ) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_v += tl.dot(b_w, tl.trans(b_h2).to(b_w.dtype)) + if K > 128: + p_w = tl.make_block_ptr( + w, (T, K), (stride_w, 1), (i_t * BT, 128), (BT, 64), (1, 0) + ) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_v += tl.dot(b_w, tl.trans(b_h3).to(b_w.dtype)) + if K > 192: + p_w = tl.make_block_ptr( + w, (T, K), (stride_w, 1), (i_t * BT, 192), (BT, 64), (1, 0) + ) + b_w = tl.load(p_w, boundary_check=(0, 1)) + b_v += tl.dot(b_w, tl.trans(b_h4).to(b_w.dtype)) + p_v = tl.make_block_ptr( + v, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) + ) + b_v = tl.load(p_v, boundary_check=(0, 1)) - b_v + + if SAVE_NEW_VALUE: + p_v = tl.make_block_ptr( + v_new, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) + ) + tl.store(p_v, b_v.to(p_v.dtype.element_ty), boundary_check=(0, 1)) + + last_idx = min((i_t + 1) * BT, T) - 1 + if USE_G: + m_t = (i_t * BT + tl.arange(0, BT)) < T + b_g_last = tl.load(g + bos * H + last_idx * H + i_h) + p_g = tl.make_block_ptr( + g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,) + ) + b_g = tl.load(p_g, boundary_check=(0,)) + b_v = b_v * tl.where(m_t, exp(b_g_last - b_g), 0)[:, None] + b_g_last = exp(b_g_last) + b_h1 *= b_g_last + if K > 64: + b_h2 *= b_g_last + if K > 128: + b_h3 *= b_g_last + if K > 192: + b_h4 *= b_g_last + + if USE_GK: + o_k1 = tl.arange(0, 64) + b_gk_last1 = tl.load( + gk + (bos + last_idx) * H * K + i_h * K + o_k1, + mask=(o_k1 < K), + other=0.0, + ) + b_h1 *= exp(b_gk_last1)[None, :] + if K > 64: + o_k2 = 64 + o_k1 + b_gk_last2 = tl.load( + gk + (bos + last_idx) * H * K + i_h * K + o_k2, + mask=(o_k2 < K), + other=0.0, + ) + b_h2 *= exp(b_gk_last2)[None, :] + if K > 128: + o_k3 = 128 + o_k1 + b_gk_last3 = tl.load( + gk + (bos + last_idx) * H * K + i_h * K + o_k3, + mask=(o_k3 < K), + other=0.0, + ) + b_h3 *= exp(b_gk_last3)[None, :] + if K > 192: + o_k4 = 192 + o_k1 + b_gk_last4 = tl.load( + gk + (bos + last_idx) * H * K + i_h * K + o_k4, + mask=(o_k4 < K), + other=0.0, + ) + b_h4 *= exp(b_gk_last4)[None, :] + b_v = b_v.to(k.dtype.element_ty) + + p_k = tl.make_block_ptr( + k, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1) + ) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_h1 += tl.trans(tl.dot(b_k, b_v)) + if K > 64: + p_k = tl.make_block_ptr( + k, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1) + ) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_h2 += tl.trans(tl.dot(b_k, b_v)) + if K > 128: + p_k = tl.make_block_ptr( + k, (K, T), (1, stride_k), (128, i_t * BT), (64, BT), (0, 1) + ) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_h3 += tl.trans(tl.dot(b_k, b_v)) + if K > 192: + p_k = tl.make_block_ptr( + k, (K, T), (1, stride_k), (192, i_t * BT), (64, BT), (0, 1) + ) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_h4 += tl.trans(tl.dot(b_k, b_v)) + # epilogue + if STORE_FINAL_STATE: + p_ht = tl.make_block_ptr(ht, (V, K), (K, 1), (i_v * BV, 0), (BV, 64), (1, 0)) + tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + if K > 64: + p_ht = tl.make_block_ptr( + ht, (V, K), (K, 1), (i_v * BV, 64), (BV, 64), (1, 0) + ) + tl.store(p_ht, b_h2.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + if K > 128: + p_ht = tl.make_block_ptr( + ht, (V, K), (K, 1), (i_v * BV, 128), (BV, 64), (1, 0) + ) + tl.store(p_ht, b_h3.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + if K > 192: + p_ht = tl.make_block_ptr( + ht, (V, K), (K, 1), (i_v * BV, 192), (BV, 64), (1, 0) + ) + tl.store(p_ht, b_h4.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_gated_delta_rule_fwd_h( + k: torch.Tensor, + w: torch.Tensor, + u: torch.Tensor, + g: torch.Tensor | None = None, + gk: torch.Tensor | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + chunk_size: int = 64, # SY: remove this argument and force chunk size 64? + save_new_value: bool = True, + cu_seqlens: torch.LongTensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + # This kernel is slightly different from fla to support Q/K with different head numbers. + # In fla, Q/K always have the same head number, so Hg is always equal to H. + B, T, Hg, K, V = *k.shape, u.shape[-1] + H = u.shape[-2] + BT = chunk_size + + chunk_indices = ( + prepare_chunk_indices(cu_seqlens, chunk_size) + if cu_seqlens is not None + else None + ) + # N: the actual number of sequences in the batch with either equal or variable lengths + if cu_seqlens is None: + N, NT, chunk_offsets = B, triton.cdiv(T, BT), None + else: + N, NT, chunk_offsets = ( + len(cu_seqlens) - 1, + len(chunk_indices), + prepare_chunk_offsets(cu_seqlens, BT), + ) + assert K <= 256, "current kernel does not support head dimension larger than 256." + + h = k.new_empty(B, NT, H, V, K) + final_state = ( + k.new_empty(N, H, V, K, dtype=torch.float32) if output_final_state else None + ) + + v_new = torch.empty_like(u) if save_new_value else None + + def grid(meta): + return (triton.cdiv(V, meta["BV"]), N * H) + + chunk_gated_delta_rule_fwd_kernel_h_blockdim64[grid]( + k=k, + v=u, + w=w, + v_new=v_new, + g=g, + gk=gk, + h=h, + h0=initial_state, + ht=final_state, + cu_seqlens=cu_seqlens, + chunk_offsets=chunk_offsets, + T=T, + H=H, + Hg=Hg, + K=K, + V=V, + BT=BT, + ) + return h, v_new, final_state diff --git a/nanovllm/layers/ops/chunk_o.py b/nanovllm/layers/ops/chunk_o.py new file mode 100644 index 000000000..6345ffcb3 --- /dev/null +++ b/nanovllm/layers/ops/chunk_o.py @@ -0,0 +1,184 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang +# +# This file contains code copied from the flash-linear-attention project. +# The original source code was licensed under the MIT license and included +# the following copyright notice: +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +# ruff: noqa: E501 + + +import torch + +import triton +import triton.language as tl + +from nanovllm.layers.ops.index import prepare_chunk_indices +from nanovllm.layers.ops.op import exp +from nanovllm.layers.ops.utils import FLA_GDN_FIX_BT, check_shared_mem, is_nvidia_hopper + +BKV_LIST = [64, 128] if check_shared_mem() else [32, 64] +NUM_WARPS = [2, 4] if is_nvidia_hopper else [2, 4, 8] + + +@triton.heuristics( + { + "USE_G": lambda args: args["g"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +@triton.autotune( + configs=[ + triton.Config({"BK": BK, "BV": BV}, num_warps=num_warps, num_stages=num_stages) + for BK in BKV_LIST + for BV in BKV_LIST + for num_warps in NUM_WARPS + for num_stages in [2, 3, 4] + ], + key=["H", "K", "V", "BT"], +) +@triton.jit(do_not_specialize=["T"]) +def chunk_fwd_kernel_o( + q, + k, + v, + h, + g, + o, + cu_seqlens, + chunk_indices, + scale, + T, + H: tl.constexpr, + Hg: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_G: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + + if IS_VARLEN: + i_tg = i_t + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + NT = tl.cdiv(T, BT) + else: + NT = tl.cdiv(T, BT) + i_tg = i_b * NT + i_t + bos, eos = i_b * T, i_b * T + T + + # offset calculation + q += (bos * Hg + i_h // (H // Hg)) * K + k += (bos * Hg + i_h // (H // Hg)) * K + v += (bos * H + i_h) * V + o += (bos * H + i_h) * V + h += (i_tg * H + i_h).to(tl.int64) * V * K + + b_o = tl.zeros([BT, BV], dtype=tl.float32) + b_A = tl.zeros([BT, BT], dtype=tl.float32) + + for i_k in range(tl.cdiv(K, BK)): + p_q = tl.make_block_ptr( + q, (T, K), (Hg * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0) + ) + p_k = tl.make_block_ptr( + k, (K, T), (1, Hg * K), (i_k * BK, i_t * BT), (BK, BT), (0, 1) + ) + p_h = tl.make_block_ptr( + h, (V, K), (K, 1), (i_v * BV, i_k * BK), (BV, BK), (1, 0) + ) + # [BT, BK] + b_q = tl.load(p_q, boundary_check=(0, 1)) + # [BK, BT] + b_k = tl.load(p_k, boundary_check=(0, 1)) + # [BV, BK] + b_h = tl.load(p_h, boundary_check=(0, 1)) + + # [BT, BK] @ [BK, BV] -> [BT, BV] + b_o += tl.dot(b_q, tl.trans(b_h)) + # [BT, BK] @ [BK, BT] -> [BT, BT] + b_A += tl.dot(b_q, b_k) + + if USE_G: + g += bos * H + i_h + p_g = tl.make_block_ptr(g, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_g = tl.load(p_g, boundary_check=(0,)) + b_o = b_o * exp(b_g)[:, None] + b_A = b_A * exp(b_g[:, None] - b_g[None, :]) + + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + m_A = (o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t) + b_A = tl.where(m_A, b_A, 0) + + p_v = tl.make_block_ptr( + v, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) + ) + p_o = tl.make_block_ptr( + o, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0) + ) + b_v = tl.load(p_v, boundary_check=(0, 1)) + + # to fix mma -> mma layout conversion + # already solved by triton v3.2 or higher + b_o = b_o * scale + tl.dot(b_A.to(b_v.dtype), b_v) * scale + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_fwd_o( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + h: torch.Tensor, + g: torch.Tensor | None = None, # cumsum of log decay + scale: float | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, +) -> torch.Tensor: + B, T, Hg, K, V = *q.shape, v.shape[-1] + H = v.shape[-2] + BT = 64 if FLA_GDN_FIX_BT else min(chunk_size, max(16, triton.next_power_of_2(T))) + chunk_indices = ( + prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + ) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + if scale is None: + scale = k.shape[-1] ** -0.5 + + o = torch.empty_like(v) + + def grid(meta): + return (triton.cdiv(V, meta["BV"]), NT, B * H) + + chunk_fwd_kernel_o[grid]( + q, + k, + v, + h, + g, + o, + cu_seqlens, + chunk_indices, + scale, + T=T, + H=H, + Hg=Hg, + K=K, + V=V, + BT=BT, + ) + return o diff --git a/nanovllm/layers/ops/chunk_scaled_dot_kkt.py b/nanovllm/layers/ops/chunk_scaled_dot_kkt.py new file mode 100644 index 000000000..d449d8084 --- /dev/null +++ b/nanovllm/layers/ops/chunk_scaled_dot_kkt.py @@ -0,0 +1,155 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang +# +# This file contains code copied from the flash-linear-attention project. +# The original source code was licensed under the MIT license and included +# the following copyright notice: +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang +# ruff: noqa: E501 + +import torch + +import triton +import triton.language as tl + +from nanovllm.layers.ops.index import prepare_chunk_indices +from nanovllm.layers.ops.op import exp + + +@triton.heuristics( + { + "USE_G": lambda args: args["g"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +@triton.autotune( + configs=[ + triton.Config({"BK": BK}, num_warps=num_warps, num_stages=num_stages) + for BK in [32, 64, 128] + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=["H", "K", "BT", "IS_VARLEN"], +) +@triton.jit(do_not_specialize=["T"]) +def chunk_scaled_dot_kkt_fwd_kernel( + k, + beta, + g, + A, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + Hg: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_G: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + o_t = i_t * BT + tl.arange(0, BT) + m_t = o_t < T + + p_beta = tl.make_block_ptr( + beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,) + ) + b_beta = tl.load(p_beta, boundary_check=(0,)) + + b_A = tl.zeros([BT, BT], dtype=tl.float32) + for i_k in range(tl.cdiv(K, BK)): + p_k = tl.make_block_ptr( + k + (bos * Hg + i_h // (H // Hg)) * K, + (T, K), + (Hg * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_kb = b_k * b_beta[:, None] + b_A += tl.dot(b_kb.to(b_k.dtype), tl.trans(b_k)) + + if USE_G: + p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + b_g = tl.load(p_g, boundary_check=(0,)) + b_g_diff = b_g[:, None] - b_g[None, :] + b_A = b_A * exp(b_g_diff) + + m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t) + b_A = tl.where(m_A, b_A, 0) + p_A = tl.make_block_ptr( + A + (bos * H + i_h) * BT, (T, BT), (BT * H, 1), (i_t * BT, 0), (BT, BT), (1, 0) + ) + tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_scaled_dot_kkt_fwd( + k: torch.Tensor, + g: torch.Tensor | None = None, + beta: torch.Tensor | None = None, + cu_seqlens: torch.LongTensor | None = None, + chunk_size: int = 64, + output_dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + r""" + Compute beta * K * K^T. + + Args: + k (torch.Tensor): + The key tensor of shape `[B, T, H, K]`. + beta (torch.Tensor): + The beta tensor of shape `[B, T, H]`. + g (torch.Tensor): + The cumulative sum of the gate tensor of shape `[B, T, H]`. Default: `None`. + cu_seqlens (torch.LongTensor): + The cumulative sequence lengths of the input tensor. + Default: None + chunk_size (int): + The chunk size. Default: 64. + output_dtype (torch.dtype): + The dtype of the output tensor. Default: `torch.float32` + + Returns: + beta * K * K^T of shape `[B, T, H, BT]` where `BT` is the chunk size. + """ + # This kernel is slightly different from fla to support Q/K with different head numbers. + # In fla, Q/K always have the same head number, so Hg is always equal to H. + B, T, Hg, K = k.shape + H = beta.shape[-1] + BT = chunk_size + chunk_indices = ( + prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + ) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + + A = torch.empty(B, T, H, BT, device=k.device, dtype=output_dtype) + chunk_scaled_dot_kkt_fwd_kernel[(NT, B * H)]( + k=k, + g=g, + beta=beta, + A=A, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + Hg=Hg, + K=K, + BT=BT, + ) + return A diff --git a/nanovllm/layers/ops/cumsum.py b/nanovllm/layers/ops/cumsum.py new file mode 100644 index 000000000..3cf52d8ef --- /dev/null +++ b/nanovllm/layers/ops/cumsum.py @@ -0,0 +1,273 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang +# +# This file contains code copied from the flash-linear-attention project. +# The original source code was licensed under the MIT license and included +# the following copyright notice: +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang +# ruff: noqa: E501 +import torch + +import triton +import triton.language as tl + +from nanovllm.layers.ops.index import prepare_chunk_indices +from nanovllm.layers.ops.utils import check_shared_mem, input_guard + +BS_LIST = [32, 64] if check_shared_mem() else [16, 32] + + +@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) +@triton.autotune( + configs=[triton.Config({}, num_warps=num_warps) for num_warps in [1, 2, 4, 8]], + key=["B", "H", "BT", "IS_VARLEN", "REVERSE"], +) +@triton.jit(do_not_specialize=["T"]) +def chunk_local_cumsum_scalar_kernel( + s, + o, + cu_seqlens, + chunk_indices, + T, + B: tl.constexpr, + H: tl.constexpr, + BT: tl.constexpr, + REVERSE: tl.constexpr, + IS_VARLEN: tl.constexpr, + HEAD_FIRST: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + if HEAD_FIRST: + p_s = tl.make_block_ptr( + s + bos * H + i_h * T, (T,), (1,), (i_t * BT,), (BT,), (0,) + ) + p_o = tl.make_block_ptr( + o + bos * H + i_h * T, (T,), (1,), (i_t * BT,), (BT,), (0,) + ) + else: + p_s = tl.make_block_ptr(s + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_o = tl.make_block_ptr(o + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)) + # [BT] + b_s = tl.load(p_s, boundary_check=(0,)).to(tl.float32) + b_o = tl.cumsum(b_s, axis=0) + if REVERSE: + b_z = tl.sum(b_s, axis=0) + b_o = -b_o + b_z[None] + b_s + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0,)) + + +@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) +@triton.autotune( + configs=[ + triton.Config({"BS": BS}, num_warps=num_warps) + for BS in BS_LIST + for num_warps in [2, 4, 8] + ], + key=["B", "H", "S", "BT", "IS_VARLEN", "REVERSE"], +) +@triton.jit(do_not_specialize=["T"]) +def chunk_local_cumsum_vector_kernel( + s, + o, + cu_seqlens, + chunk_indices, + T, + B: tl.constexpr, + H: tl.constexpr, + S: tl.constexpr, + BT: tl.constexpr, + BS: tl.constexpr, + REVERSE: tl.constexpr, + IS_VARLEN: tl.constexpr, + HEAD_FIRST: tl.constexpr, +): + i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + o_i = tl.arange(0, BT) + if REVERSE: + m_s = tl.where(o_i[:, None] <= o_i[None, :], 1.0, 0.0) + else: + m_s = tl.where(o_i[:, None] >= o_i[None, :], 1.0, 0.0) + + if HEAD_FIRST: + p_s = tl.make_block_ptr( + s + (bos * H + i_h * T) * S, + (T, S), + (S, 1), + (i_t * BT, i_s * BS), + (BT, BS), + (1, 0), + ) + p_o = tl.make_block_ptr( + o + (bos * H + i_h * T) * S, + (T, S), + (S, 1), + (i_t * BT, i_s * BS), + (BT, BS), + (1, 0), + ) + else: + p_s = tl.make_block_ptr( + s + (bos * H + i_h) * S, + (T, S), + (H * S, 1), + (i_t * BT, i_s * BS), + (BT, BS), + (1, 0), + ) + p_o = tl.make_block_ptr( + o + (bos * H + i_h) * S, + (T, S), + (H * S, 1), + (i_t * BT, i_s * BS), + (BT, BS), + (1, 0), + ) + # [BT, BS] + b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32) + b_o = tl.dot(m_s, b_s, allow_tf32=False) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + +def chunk_local_cumsum_scalar( + g: torch.Tensor, + chunk_size: int, + reverse: bool = False, + cu_seqlens: torch.Tensor | None = None, + head_first: bool = False, + output_dtype: torch.dtype | None = torch.float, +) -> torch.Tensor: + if head_first: + B, H, T = g.shape + else: + B, T, H = g.shape + assert chunk_size == 2 ** (chunk_size.bit_length() - 1), ( + "chunk_size must be a power of 2" + ) + BT = chunk_size + chunk_indices = ( + prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + ) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype) + grid = (NT, B * H) + chunk_local_cumsum_scalar_kernel[grid]( + g_org, + g, + cu_seqlens, + chunk_indices, + T=T, + B=B, + H=H, + BT=BT, + HEAD_FIRST=head_first, + REVERSE=reverse, + ) + return g + + +def chunk_local_cumsum_vector( + g: torch.Tensor, + chunk_size: int, + reverse: bool = False, + cu_seqlens: torch.Tensor | None = None, + head_first: bool = False, + output_dtype: torch.dtype | None = torch.float, +) -> torch.Tensor: + if head_first: + B, H, T, S = g.shape + else: + B, T, H, S = g.shape + BT = chunk_size + chunk_indices = ( + prepare_chunk_indices(cu_seqlens, chunk_size) + if cu_seqlens is not None + else None + ) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + assert chunk_size == 2 ** (chunk_size.bit_length() - 1), ( + "chunk_size must be a power of 2" + ) + + g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype) + + def grid(meta): + return (triton.cdiv(meta["S"], meta["BS"]), NT, B * H) + + # keep cumulative normalizer in fp32 + # this kernel is equivalent to + # g = g.view(B, H, NT, BT, -1).cumsum(-2).view(B, H, T, -1) + chunk_local_cumsum_vector_kernel[grid]( + g_org, + g, + cu_seqlens, + chunk_indices, + T=T, + B=B, + H=H, + S=S, + BT=BT, + HEAD_FIRST=head_first, + REVERSE=reverse, + ) + return g + + +@input_guard +def chunk_local_cumsum( + g: torch.Tensor, + chunk_size: int, + reverse: bool = False, + cu_seqlens: torch.Tensor | None = None, + head_first: bool = False, + output_dtype: torch.dtype | None = torch.float, + **kwargs, +) -> torch.Tensor: + # Only support head_first=False (time-first format: B, T, H) + assert not head_first, "Only head_first=False (time-first format) is supported" + if cu_seqlens is not None: + assert g.shape[0] == 1, ( + "Only batch size 1 is supported when cu_seqlens are provided" + ) + if len(g.shape) == 3: + return chunk_local_cumsum_scalar( + g, chunk_size, reverse, cu_seqlens, head_first, output_dtype + ) + elif len(g.shape) == 4: + return chunk_local_cumsum_vector( + g, chunk_size, reverse, cu_seqlens, head_first, output_dtype + ) + else: + raise ValueError( + f"Unsupported input shape {g.shape}. " + f"which should be (B, T, H, D) if `head_first=False` " + f"or (B, H, T, D) otherwise" + ) diff --git a/nanovllm/layers/ops/fused_recurrent.py b/nanovllm/layers/ops/fused_recurrent.py new file mode 100644 index 000000000..3e8503c26 --- /dev/null +++ b/nanovllm/layers/ops/fused_recurrent.py @@ -0,0 +1,394 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang +# +# This file contains code copied from the flash-linear-attention project. +# The original source code was licensed under the MIT license and included +# the following copyright notice: +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang +# ruff: noqa: E501 + +import torch + +import triton +import triton.language as tl + +from nanovllm.layers.ops.op import exp + + +@triton.heuristics( + { + "USE_INITIAL_STATE": lambda args: args["h0"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + "IS_CONTINUOUS_BATCHING": lambda args: args["ssm_state_indices"] is not None, + "IS_SPEC_DECODING": lambda args: args["num_accepted_tokens"] is not None, + } +) +@triton.jit(do_not_specialize=["N", "T"]) +def fused_recurrent_gated_delta_rule_fwd_kernel( + q, + k, + v, + g, + beta, + o, + h0, + ht, + cu_seqlens, + ssm_state_indices, + num_accepted_tokens, + scale, + N: tl.int64, # num of sequences + T: tl.int64, # num of tokens + B: tl.constexpr, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + stride_init_state_token: tl.constexpr, + stride_final_state_token: tl.constexpr, + stride_indices_seq: tl.constexpr, + stride_indices_tok: tl.constexpr, + USE_INITIAL_STATE: tl.constexpr, # whether to use initial state + INPLACE_FINAL_STATE: tl.constexpr, # whether to store final state inplace + IS_BETA_HEADWISE: tl.constexpr, # whether beta is headwise vector or scalar, + USE_QK_L2NORM_IN_KERNEL: tl.constexpr, + IS_VARLEN: tl.constexpr, + IS_CONTINUOUS_BATCHING: tl.constexpr, + IS_SPEC_DECODING: tl.constexpr, + IS_KDA: tl.constexpr, +): + i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + i_n, i_hv = i_nh // HV, i_nh % HV + i_h = i_hv // (HV // H) + if IS_VARLEN: + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int64), + tl.load(cu_seqlens + i_n + 1).to(tl.int64), + ) + all = T + T = eos - bos + else: + bos, eos = i_n * T, i_n * T + T + all = B * T + + if T == 0: + # no tokens to process for this sequence + return + + o_k = i_k * BK + tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + + p_q = q + (bos * H + i_h) * K + o_k + p_k = k + (bos * H + i_h) * K + o_k + p_v = v + (bos * HV + i_hv) * V + o_v + if IS_BETA_HEADWISE: + p_beta = beta + (bos * HV + i_hv) * V + o_v + else: + p_beta = beta + bos * HV + i_hv + + if not IS_KDA: + p_g = g + bos * HV + i_hv + else: + p_gk = g + (bos * HV + i_hv) * K + o_k + + p_o = o + ((i_k * all + bos) * HV + i_hv) * V + o_v + + mask_k = o_k < K + mask_v = o_v < V + mask_h = mask_v[:, None] & mask_k[None, :] + + b_h = tl.zeros([BV, BK], dtype=tl.float32) + if USE_INITIAL_STATE: + if IS_CONTINUOUS_BATCHING: + if IS_SPEC_DECODING: + i_t = tl.load(num_accepted_tokens + i_n).to(tl.int64) - 1 + else: + i_t = 0 + # Load state index and check for PAD_SLOT_ID (-1) + state_idx = tl.load(ssm_state_indices + i_n * stride_indices_seq + i_t).to( + tl.int64 + ) + # Skip if state index is invalid (PAD_SLOT_ID = -1) + if state_idx < 0: + return + p_h0 = h0 + state_idx * stride_init_state_token + else: + p_h0 = h0 + bos * HV * V * K + p_h0 = p_h0 + i_hv * V * K + o_v[:, None] * K + o_k[None, :] + b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32) + + for i_t in range(0, T): + b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32) + b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32) + b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32) + + if USE_QK_L2NORM_IN_KERNEL: + b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6) + b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6) + b_q = b_q * scale + # [BV, BK] + if not IS_KDA: + b_g = tl.load(p_g).to(tl.float32) + b_h *= exp(b_g) + else: + b_gk = tl.load(p_gk).to(tl.float32) + b_h *= exp(b_gk[None, :]) + # [BV] + b_v -= tl.sum(b_h * b_k[None, :], 1) + if IS_BETA_HEADWISE: + b_beta = tl.load(p_beta, mask=mask_v, other=0).to(tl.float32) + else: + b_beta = tl.load(p_beta).to(tl.float32) + b_v *= b_beta + # [BV, BK] + b_h += b_v[:, None] * b_k[None, :] + # [BV] + b_o = tl.sum(b_h * b_q[None, :], 1) + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v) + + # keep the states for multi-query tokens + if INPLACE_FINAL_STATE: + # Load state index and check for PAD_SLOT_ID (-1) + final_state_idx = tl.load( + ssm_state_indices + i_n * stride_indices_seq + i_t + ).to(tl.int64) + # Only store if state index is valid (not PAD_SLOT_ID) + if final_state_idx >= 0: + p_ht = ht + final_state_idx * stride_final_state_token + p_ht = p_ht + i_hv * V * K + o_v[:, None] * K + o_k[None, :] + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) + else: + p_ht = ht + (bos + i_t) * stride_final_state_token + p_ht = p_ht + i_hv * V * K + o_v[:, None] * K + o_k[None, :] + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h) + + p_q += H * K + p_k += H * K + p_o += HV * V + p_v += HV * V + if not IS_KDA: + p_g += HV + else: + p_gk += HV * K + p_beta += HV * (V if IS_BETA_HEADWISE else 1) + + +def fused_recurrent_gated_delta_rule_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + inplace_final_state: bool = True, + cu_seqlens: torch.LongTensor | None = None, + ssm_state_indices: torch.Tensor | None = None, + num_accepted_tokens: torch.Tensor | None = None, + use_qk_l2norm_in_kernel: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, H, K, V = *k.shape, v.shape[-1] + HV = v.shape[2] + N = B if cu_seqlens is None else len(cu_seqlens) - 1 + BK, BV = triton.next_power_of_2(K), min(triton.next_power_of_2(V), 32) + NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV) + assert NK == 1, "NK > 1 is not supported yet" + num_stages = 3 + num_warps = 1 + + o = q.new_empty(NK, *v.shape) + if inplace_final_state: + final_state = initial_state + else: + final_state = q.new_empty(T, HV, V, K, dtype=initial_state.dtype) + + stride_init_state_token = initial_state.stride(0) + stride_final_state_token = final_state.stride(0) + + if ssm_state_indices is None: + stride_indices_seq, stride_indices_tok = 1, 1 + elif ssm_state_indices.ndim == 1: + stride_indices_seq, stride_indices_tok = ssm_state_indices.stride(0), 1 + else: + stride_indices_seq, stride_indices_tok = ssm_state_indices.stride() + + grid = (NK, NV, N * HV) + fused_recurrent_gated_delta_rule_fwd_kernel[grid]( + q=q, + k=k, + v=v, + g=g, + beta=beta, + o=o, + h0=initial_state, + ht=final_state, + cu_seqlens=cu_seqlens, + ssm_state_indices=ssm_state_indices, + num_accepted_tokens=num_accepted_tokens, + scale=scale, + N=N, + T=T, + B=B, + H=H, + HV=HV, + K=K, + V=V, + BK=BK, + BV=BV, + stride_init_state_token=stride_init_state_token, + stride_final_state_token=stride_final_state_token, + stride_indices_seq=stride_indices_seq, + stride_indices_tok=stride_indices_tok, + IS_BETA_HEADWISE=beta.ndim == v.ndim, + USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, + INPLACE_FINAL_STATE=inplace_final_state, + IS_KDA=False, + num_warps=num_warps, + num_stages=num_stages, + ) + o = o.squeeze(0) + return o, final_state + + +class FusedRecurrentFunction(torch.autograd.Function): + @staticmethod + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + inplace_final_state: bool = True, + cu_seqlens: torch.LongTensor | None = None, + ssm_state_indices: torch.Tensor | None = None, + num_accepted_tokens: torch.Tensor | None = None, + use_qk_l2norm_in_kernel: bool = False, + ): + o, final_state = fused_recurrent_gated_delta_rule_fwd( + q=q.contiguous(), + k=k.contiguous(), + v=v.contiguous(), + g=g.contiguous(), + beta=beta.contiguous(), + scale=scale, + initial_state=initial_state, + inplace_final_state=inplace_final_state, + cu_seqlens=cu_seqlens, + ssm_state_indices=ssm_state_indices, + num_accepted_tokens=num_accepted_tokens, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + ) + + return o, final_state + + +def fused_recurrent_gated_delta_rule( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor = None, + scale: float = None, + initial_state: torch.Tensor = None, + inplace_final_state: bool = True, + cu_seqlens: torch.LongTensor | None = None, + ssm_state_indices: torch.Tensor | None = None, + num_accepted_tokens: torch.Tensor | None = None, + use_qk_l2norm_in_kernel: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + r""" + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, HV, V]`. + GVA is applied if `HV > H`. + g (torch.Tensor): + g (decays) of shape `[B, T, HV]`. + beta (torch.Tensor): + betas of shape `[B, T, HV]`. + scale (Optional[int]): + Scale factor for the RetNet attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[N, HV, V, K]` for `N` input sequences. + For equal-length input sequences, `N` equals the batch size `B`. + Default: `None`. + inplace_final_state: bool: + Whether to store the final state in-place to save memory. + Default: `True`. + cu_seqlens (torch.LongTensor): + Cumulative sequence lengths of shape `[N+1]` used for variable-length training, + consistent with the FlashAttention API. + ssm_state_indices (Optional[torch.Tensor]): + Indices to map the input sequences to the initial/final states. + num_accepted_tokens (Optional[torch.Tensor]): + Number of accepted tokens for each sequence during decoding. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, HV, V]`. + final_state (torch.Tensor): + Final state of shape `[N, HV, V, K]`. + + Examples:: + >>> import torch + >>> import torch.nn.functional as F + >>> from einops import rearrange + >>> from fla.ops.gated_delta_rule import fused_recurrent_gated_delta_rule + # inputs with equal lengths + >>> B, T, H, HV, K, V = 4, 2048, 4, 8, 512, 512 + >>> q = torch.randn(B, T, H, K, device='cuda') + >>> k = F.normalize(torch.randn(B, T, H, K, device='cuda'), p=2, dim=-1) + >>> v = torch.randn(B, T, HV, V, device='cuda') + >>> g = F.logsigmoid(torch.rand(B, T, HV, device='cuda')) + >>> beta = torch.rand(B, T, HV, device='cuda').sigmoid() + >>> h0 = torch.randn(B, HV, V, K, device='cuda') + >>> o, ht = fused_gated_recurrent_delta_rule( + q, k, v, g, beta, + initial_state=h0, + ) + # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required + >>> q, k, v, g, beta = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, g, beta)) + # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected + >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long) + >>> o_var, ht_var = fused_gated_recurrent_delta_rule( + q, k, v, g, beta, + initial_state=h0, + cu_seqlens=cu_seqlens + ) + """ + if cu_seqlens is not None and q.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing." + ) + if scale is None: + scale = k.shape[-1] ** -0.5 + else: + assert scale > 0, "scale must be positive" + if beta is None: + beta = torch.ones_like(q[..., 0]) + o, final_state = FusedRecurrentFunction.apply( + q, + k, + v, + g, + beta, + scale, + initial_state, + inplace_final_state, + cu_seqlens, + ssm_state_indices, + num_accepted_tokens, + use_qk_l2norm_in_kernel, + ) + return o, final_state diff --git a/nanovllm/layers/ops/index.py b/nanovllm/layers/ops/index.py new file mode 100644 index 000000000..ff9813f61 --- /dev/null +++ b/nanovllm/layers/ops/index.py @@ -0,0 +1,40 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang +# +# This file contains code copied from the flash-linear-attention project. +# The original source code was licensed under the MIT license and included +# the following copyright notice: +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import torch +import triton + +from nanovllm.layers.ops.utils import tensor_cache + + +@tensor_cache +def prepare_lens(cu_seqlens: torch.LongTensor) -> torch.LongTensor: + return cu_seqlens[1:] - cu_seqlens[:-1] + + +@tensor_cache +def prepare_chunk_indices( + cu_seqlens: torch.LongTensor, chunk_size: int +) -> torch.LongTensor: + indices = torch.cat( + [ + torch.arange(n) + for n in triton.cdiv(prepare_lens(cu_seqlens), chunk_size).tolist() + ] + ) + return torch.stack([indices.eq(0).cumsum(0) - 1, indices], 1).to(cu_seqlens) + + +@tensor_cache +def prepare_chunk_offsets( + cu_seqlens: torch.LongTensor, chunk_size: int +) -> torch.LongTensor: + return torch.cat( + [cu_seqlens.new_tensor([0]), triton.cdiv(prepare_lens(cu_seqlens), chunk_size)] + ).cumsum(-1) diff --git a/nanovllm/layers/ops/l2norm.py b/nanovllm/layers/ops/l2norm.py new file mode 100644 index 000000000..b3611322a --- /dev/null +++ b/nanovllm/layers/ops/l2norm.py @@ -0,0 +1,147 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang +# +# This file contains code copied from the flash-linear-attention project. +# The original source code was licensed under the MIT license and included +# the following copyright notice: +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import os + +import torch + +import triton +import triton.language as tl + +BT_LIST = [8, 16, 32, 64, 128] + +USE_DEFAULT_FLA_NORM = int(os.getenv("USE_DEFAULT_FLA_NORM", "0")) + + +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps) for num_warps in [1, 2, 4, 8, 16, 32] + ], + key=["D"], +) +@triton.jit +def l2norm_fwd_kernel1( + x, + y, + D, + BD: tl.constexpr, + eps, +): + i_t = tl.program_id(0) + x += i_t * D + y += i_t * D + # Compute mean and variance + cols = tl.arange(0, BD) + mask = cols < D + b_x = tl.load(x + cols, mask=mask, other=0.0).to(tl.float32) + b_var = tl.sum(b_x * b_x, axis=0) + b_rstd = 1 / tl.sqrt(b_var + eps) + # tl.store(Rstd + i_t, rstd) + # Normalize and apply linear transformation + b_y = b_x * b_rstd + tl.store(y + cols, b_y, mask=mask) + + +@triton.autotune( + configs=[ + triton.Config({"BT": BT}, num_warps=num_warps) + for num_warps in [1, 2, 4, 8, 16] + for BT in BT_LIST + ], + key=["D"], +) +@triton.jit(do_not_specialize=["NB"]) +def l2norm_fwd_kernel( + x, + y, + eps, + NB, + T, + D: tl.constexpr, + BT: tl.constexpr, + BD: tl.constexpr, +): + i_t = tl.program_id(0) + p_x = tl.make_block_ptr(x, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) + b_x = tl.load(p_x, boundary_check=(0, 1)).to(tl.float32) + b_var = tl.sum(b_x * b_x, axis=1) + b_y = b_x / tl.sqrt(b_var + eps)[:, None] + p_y = tl.make_block_ptr(y, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)) + tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1)) + + +@triton.jit +def l2norm_fwd_kernel2(X, Y, eps, M, N: tl.constexpr, MBLOCK: tl.constexpr): + xoffset = tl.program_id(0) * MBLOCK + row_idx = xoffset + tl.arange(0, MBLOCK)[:, None] + xmask = row_idx < M + rindex = tl.arange(0, N)[None, :] + xs = tl.load(X + (rindex + N * row_idx), xmask).to(tl.float32) + square = tl.broadcast_to(xs * xs, [MBLOCK, N]) + square_sum = tl.sum(tl.where(xmask, square, 0), 1)[:, None] + rsqrt = tl.rsqrt(square_sum + eps) + tl.store(Y + (rindex + N * row_idx), xs * rsqrt, xmask) + + +def l2norm_fwd( + x: torch.Tensor, eps: float = 1e-6, output_dtype: torch.dtype | None = None +): + x_shape_og = x.shape + x = x.view(-1, x.shape[-1]) + # allocate output + if output_dtype is None: + y = torch.empty_like(x) + else: + y = torch.empty_like(x, dtype=output_dtype) + assert y.stride(-1) == 1 + T, D = x.shape[0], x.shape[-1] + # rstd = torch.empty((T,), dtype=torch.float32, device=x.device) + # Less than 64KB per feature: enqueue fused kernel + MAX_FUSED_SIZE = 65536 // x.element_size() + BD = min(MAX_FUSED_SIZE, triton.next_power_of_2(D)) + if D > BD: + raise RuntimeError("This layer doesn't support feature dim >= 64KB.") + + if not USE_DEFAULT_FLA_NORM: + MBLOCK = 32 + # M, N = x.shape + l2norm_fwd_kernel2[(triton.cdiv(T, MBLOCK),)]( + x, + y, + eps, + T, + D, + MBLOCK, + ) + else: + if D <= 512: + NB = triton.cdiv(T, 2048) + + def grid(meta): + return (triton.cdiv(T, meta["BT"]),) + + l2norm_fwd_kernel[grid]( + x, + y, + eps, + NB=NB, + T=T, + D=D, + BD=BD, + ) + else: + l2norm_fwd_kernel1[(T,)]( + x, + y, + eps=eps, + D=D, + BD=BD, + ) + + return y.view(x_shape_og) diff --git a/nanovllm/layers/ops/op.py b/nanovllm/layers/ops/op.py new file mode 100644 index 000000000..58e24799d --- /dev/null +++ b/nanovllm/layers/ops/op.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang +# +# This file contains code copied from the flash-linear-attention project. +# The original source code was licensed under the MIT license and included +# the following copyright notice: +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import os + +import triton +import triton.language as tl + +# Use fast ops if available +if os.environ.get("FLA_USE_FAST_OPS", "0") == "1": + try: + import triton.language.extra.libdevice as tldevice + exp = tldevice.fast_expf + log = tldevice.fast_logf + log2 = tldevice.fast_log2f + except ImportError: + exp = tl.exp + log = tl.log + log2 = tl.log2 +else: + exp = tl.exp + log = tl.log + log2 = tl.log2 + +# Check if gather is supported +try: + gather = tl.gather +except AttributeError: + @triton.jit + def gather(src, index, axis, _builder=None): + """Fallback gather when tl.gather is not supported.""" + return None + +# Check if make_tensor_descriptor is available +if hasattr(triton.language, "_experimental_make_tensor_descriptor"): + make_tensor_descriptor = triton.language._experimental_make_tensor_descriptor +elif hasattr(triton.language, "make_tensor_descriptor"): + make_tensor_descriptor = triton.language.make_tensor_descriptor +else: + @triton.jit + def make_tensor_descriptor( + base, + shape, + strides, + block_shape, + _builder=None, + ): + """Fallback when TMA is not supported.""" + return None diff --git a/nanovllm/layers/ops/solve_tril.py b/nanovllm/layers/ops/solve_tril.py new file mode 100644 index 000000000..3c5051042 --- /dev/null +++ b/nanovllm/layers/ops/solve_tril.py @@ -0,0 +1,557 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang +# +# This file contains code copied from the flash-linear-attention project. +# The original source code was licensed under the MIT license and included +# the following copyright notice: +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang +# ruff: noqa: E501 + +import os + +import torch + +import triton +import triton.language as tl + +from nanovllm.layers.ops.index import prepare_chunk_indices +from nanovllm.layers.ops.op import make_tensor_descriptor +from nanovllm.layers.ops.utils import input_guard, is_amd, is_tma_supported + +FLA_TRIL_PRECISION = os.environ.get("FLA_TRIL_PRECISION", "ieee") +ALLOWED_TRIL_PRECISIONS = ["ieee", "tf32"] if is_amd else ["ieee", "tf32", "tf32x3"] +assert FLA_TRIL_PRECISION in ALLOWED_TRIL_PRECISIONS, ( + f"FLA_TRIL_PRECISION must be one of {ALLOWED_TRIL_PRECISIONS}, but got {FLA_TRIL_PRECISION}" +) + + +@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [1, 2, 4, 8] + for num_stages in [2, 3, 4, 5] + ], + key=["BT"], +) +@triton.jit(do_not_specialize=["T"]) +def solve_tril_16x16_kernel( + A, + Ai, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + BT: tl.constexpr, + USE_TMA: tl.constexpr, + IS_VARLEN: tl.constexpr, + DOT_PRECISION: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + o_i = tl.arange(0, 16) + m_A = o_i[:, None] > o_i[None, :] + m_I = o_i[:, None] == o_i[None, :] + + A = A + (bos * H + i_h) * BT + Ai = Ai + (bos * H + i_h) * 16 + + offset = (i_t * 16) % BT + if not USE_TMA: + p_A = tl.make_block_ptr( + A, (T, BT), (H * BT, 1), (i_t * 16, offset), (16, 16), (1, 0) + ) + # [16, 16] + b_A = tl.load(p_A, boundary_check=(0, 1)).to(tl.float32) + else: + desc = make_tensor_descriptor(A, [T, BT], [H * BT, 1], [16, 16]) + desc_o = make_tensor_descriptor(Ai, [T, 16], [H * 16, 1], [16, 16]) + b_A = desc.load([i_t * 16, offset]).to(tl.float32) + b_A = -tl.where(m_A, b_A, 0) + + for i in range(2, min(16, T - i_t * 16)): + # [16] + b_a = -tl.load(A + (i_t * 16 + i) * H * BT + o_i + offset) + b_a = b_a + tl.sum(b_a[:, None] * b_A, 0) + b_A = tl.where((o_i == i)[:, None], b_a, b_A) + b_A += m_I + if not USE_TMA: + p_Ai = tl.make_block_ptr( + Ai, (T, 16), (H * 16, 1), (i_t * 16, 0), (16, 16), (1, 0) + ) + tl.store( + p_Ai, + b_A.to(p_Ai.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + else: + desc_o.store([i_t * 16, 0], b_A.to(desc_o.dtype, fp_downcast_rounding="rtne")) + + +@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [1, 2, 4, 8] + for num_stages in [2, 3, 4, 5] + ], + key=["H", "BT", "IS_VARLEN"], +) +@triton.jit(do_not_specialize=["T"]) +def merge_16x16_to_32x32_inverse_kernel( + A, + Ai, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + BT: tl.constexpr, + USE_TMA: tl.constexpr, + IS_VARLEN: tl.constexpr, + DOT_PRECISION: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + o_i = tl.arange(0, 16) + m_A = o_i[:, None] > o_i[None, :] + m_I = o_i[:, None] == o_i[None, :] + A += (bos * H + i_h) * BT + Ai += (bos * H + i_h) * BT + + if not USE_TMA: + p_A_11 = tl.make_block_ptr( + A, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0) + ) + p_A_22 = tl.make_block_ptr( + A, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0) + ) + b_Ai_11 = tl.load(p_A_11, boundary_check=(0, 1)).to(tl.float32) + b_Ai_22 = tl.load(p_A_22, boundary_check=(0, 1)).to(tl.float32) + else: + desc = make_tensor_descriptor(A, [T, BT], [H * BT, 1], [16, 16]) + desc_o = make_tensor_descriptor(Ai, [T, BT], [H * BT, 1], [16, 16]) + b_Ai_11 = desc.load([i_t * BT + 0, 0]).to(tl.float32) + b_Ai_22 = desc.load([i_t * BT + 16, 16]).to(tl.float32) + + # [16, 16] + b_Ai_11 = -tl.where(m_A, b_Ai_11, 0) + b_Ai_22 = -tl.where(m_A, b_Ai_22, 0) + + for i in range(2, min(16, T - i_t * BT)): + b_a_11 = -tl.load(A + (i_t * BT + i) * H * BT + o_i) + b_a_11 += tl.sum(b_a_11[:, None] * b_Ai_11, 0) + b_Ai_11 = tl.where((o_i == i)[:, None], b_a_11, b_Ai_11) + for i in range(16 + 2, min(32, T - i_t * BT)): + b_a_22 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 16) + b_a_22 += tl.sum(b_a_22[:, None] * b_Ai_22, 0) + b_Ai_22 = tl.where((o_i == i - 16)[:, None], b_a_22, b_Ai_22) + + b_Ai_11 += m_I + b_Ai_22 += m_I + + if not USE_TMA: + p_A_21 = tl.make_block_ptr( + A, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0) + ) + b_A_21 = tl.load(p_A_21, boundary_check=(0, 1)).to(tl.float32) + else: + b_A_21 = desc.load([i_t * BT + 16, 0]).to(tl.float32) + + b_Ai_21 = -tl.dot( + tl.dot(b_Ai_22, b_A_21, input_precision=DOT_PRECISION), + b_Ai_11, + input_precision=DOT_PRECISION, + ) + + if not USE_TMA: + p_Ai_11 = tl.make_block_ptr( + Ai, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0) + ) + p_Ai_21 = tl.make_block_ptr( + Ai, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0) + ) + p_Ai_22 = tl.make_block_ptr( + Ai, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0) + ) + tl.store( + p_Ai_11, + b_Ai_11.to(p_Ai_11.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_22, + b_Ai_22.to(p_Ai_22.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_21, + b_Ai_21.to(p_Ai_21.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + else: + desc_o.store( + [i_t * BT + 0, 0], b_Ai_11.to(desc_o.dtype, fp_downcast_rounding="rtne") + ) + desc_o.store( + [i_t * BT + 16, 0], b_Ai_21.to(desc_o.dtype, fp_downcast_rounding="rtne") + ) + desc_o.store( + [i_t * BT + 16, 16], b_Ai_22.to(desc_o.dtype, fp_downcast_rounding="rtne") + ) + + +@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4, 5] + ], + key=["H", "BT", "IS_VARLEN"], +) +@triton.jit(do_not_specialize=["T"]) +def merge_16x16_to_64x64_inverse_kernel( + A, + Ai, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + BT: tl.constexpr, + USE_TMA: tl.constexpr, + IS_VARLEN: tl.constexpr, + DOT_PRECISION: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + o_i = tl.arange(0, 16) + m_A = o_i[:, None] > o_i[None, :] + m_I = o_i[:, None] == o_i[None, :] + A += (bos * H + i_h) * BT + Ai += (bos * H + i_h) * BT + + if not USE_TMA: + p_A_11 = tl.make_block_ptr( + A, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0) + ) + p_A_22 = tl.make_block_ptr( + A, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0) + ) + p_A_33 = tl.make_block_ptr( + A, (T, BT), (H * BT, 1), (i_t * BT + 32, 32), (16, 16), (1, 0) + ) + p_A_44 = tl.make_block_ptr( + A, (T, BT), (H * BT, 1), (i_t * BT + 48, 48), (16, 16), (1, 0) + ) + b_Ai_11 = tl.load(p_A_11, boundary_check=(0, 1)).to(tl.float32) + b_Ai_22 = tl.load(p_A_22, boundary_check=(0, 1)).to(tl.float32) + b_Ai_33 = tl.load(p_A_33, boundary_check=(0, 1)).to(tl.float32) + b_Ai_44 = tl.load(p_A_44, boundary_check=(0, 1)).to(tl.float32) + else: + desc = make_tensor_descriptor(A, [T, BT], [H * BT, 1], [16, 16]) + desc_o = make_tensor_descriptor(Ai, [T, BT], [H * BT, 1], [16, 16]) + b_Ai_11 = desc.load([i_t * BT + 0, 0]).to(tl.float32) + b_Ai_22 = desc.load([i_t * BT + 16, 16]).to(tl.float32) + b_Ai_33 = desc.load([i_t * BT + 32, 32]).to(tl.float32) + b_Ai_44 = desc.load([i_t * BT + 48, 48]).to(tl.float32) + + # [16, 16] + b_Ai_11 = -tl.where(m_A, b_Ai_11, 0) + b_Ai_22 = -tl.where(m_A, b_Ai_22, 0) + b_Ai_33 = -tl.where(m_A, b_Ai_33, 0) + b_Ai_44 = -tl.where(m_A, b_Ai_44, 0) + + for i in range(2, min(16, T - i_t * BT)): + b_a_11 = -tl.load(A + (i_t * BT + i) * H * BT + o_i) + b_a_11 += tl.sum(b_a_11[:, None] * b_Ai_11, 0) + b_Ai_11 = tl.where((o_i == i)[:, None], b_a_11, b_Ai_11) + for i in range(16 + 2, min(32, T - i_t * BT)): + b_a_22 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 16) + b_a_22 += tl.sum(b_a_22[:, None] * b_Ai_22, 0) + b_Ai_22 = tl.where((o_i == i - 16)[:, None], b_a_22, b_Ai_22) + for i in range(32 + 2, min(48, T - i_t * BT)): + b_a_33 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 32) + b_a_33 += tl.sum(b_a_33[:, None] * b_Ai_33, 0) + b_Ai_33 = tl.where((o_i == i - 32)[:, None], b_a_33, b_Ai_33) + for i in range(48 + 2, min(64, T - i_t * BT)): + b_a_44 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 48) + b_a_44 += tl.sum(b_a_44[:, None] * b_Ai_44, 0) + b_Ai_44 = tl.where((o_i == i - 48)[:, None], b_a_44, b_Ai_44) + b_Ai_11 += m_I + b_Ai_22 += m_I + b_Ai_33 += m_I + b_Ai_44 += m_I + + if not USE_TMA: + p_A_21 = tl.make_block_ptr( + A, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0) + ) + p_A_31 = tl.make_block_ptr( + A, (T, BT), (H * BT, 1), (i_t * BT + 32, 0), (16, 16), (1, 0) + ) + p_A_32 = tl.make_block_ptr( + A, (T, BT), (H * BT, 1), (i_t * BT + 32, 16), (16, 16), (1, 0) + ) + p_A_41 = tl.make_block_ptr( + A, (T, BT), (H * BT, 1), (i_t * BT + 48, 0), (16, 16), (1, 0) + ) + p_A_42 = tl.make_block_ptr( + A, (T, BT), (H * BT, 1), (i_t * BT + 48, 16), (16, 16), (1, 0) + ) + p_A_43 = tl.make_block_ptr( + A, (T, BT), (H * BT, 1), (i_t * BT + 48, 32), (16, 16), (1, 0) + ) + b_A_21 = tl.load(p_A_21, boundary_check=(0, 1)).to(tl.float32) + b_A_31 = tl.load(p_A_31, boundary_check=(0, 1)).to(tl.float32) + b_A_32 = tl.load(p_A_32, boundary_check=(0, 1)).to(tl.float32) + b_A_41 = tl.load(p_A_41, boundary_check=(0, 1)).to(tl.float32) + b_A_42 = tl.load(p_A_42, boundary_check=(0, 1)).to(tl.float32) + b_A_43 = tl.load(p_A_43, boundary_check=(0, 1)).to(tl.float32) + else: + b_A_21 = desc.load([i_t * BT + 16, 0]).to(tl.float32) + b_A_31 = desc.load([i_t * BT + 32, 0]).to(tl.float32) + b_A_32 = desc.load([i_t * BT + 32, 16]).to(tl.float32) + b_A_41 = desc.load([i_t * BT + 48, 0]).to(tl.float32) + b_A_42 = desc.load([i_t * BT + 48, 16]).to(tl.float32) + b_A_43 = desc.load([i_t * BT + 48, 32]).to(tl.float32) + + b_Ai_21 = -tl.dot( + tl.dot(b_Ai_22, b_A_21, input_precision=DOT_PRECISION), + b_Ai_11, + input_precision=DOT_PRECISION, + ) + b_Ai_32 = -tl.dot( + tl.dot(b_Ai_33, b_A_32, input_precision=DOT_PRECISION), + b_Ai_22, + input_precision=DOT_PRECISION, + ) + b_Ai_43 = -tl.dot( + tl.dot(b_Ai_44, b_A_43, input_precision=DOT_PRECISION), + b_Ai_33, + input_precision=DOT_PRECISION, + ) + + b_Ai_31 = -tl.dot( + b_Ai_33, + tl.dot(b_A_31, b_Ai_11, input_precision=DOT_PRECISION) + + tl.dot(b_A_32, b_Ai_21, input_precision=DOT_PRECISION), + input_precision=DOT_PRECISION, + ) + b_Ai_42 = -tl.dot( + b_Ai_44, + tl.dot(b_A_42, b_Ai_22, input_precision=DOT_PRECISION) + + tl.dot(b_A_43, b_Ai_32, input_precision=DOT_PRECISION), + input_precision=DOT_PRECISION, + ) + b_Ai_41 = -tl.dot( + b_Ai_44, + tl.dot(b_A_41, b_Ai_11, input_precision=DOT_PRECISION) + + tl.dot(b_A_42, b_Ai_21, input_precision=DOT_PRECISION) + + tl.dot(b_A_43, b_Ai_31, input_precision=DOT_PRECISION), + input_precision=DOT_PRECISION, + ) + + if not USE_TMA: + p_Ai_11 = tl.make_block_ptr( + Ai, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0) + ) + p_Ai_22 = tl.make_block_ptr( + Ai, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0) + ) + p_Ai_33 = tl.make_block_ptr( + Ai, (T, BT), (H * BT, 1), (i_t * BT + 32, 32), (16, 16), (1, 0) + ) + p_Ai_44 = tl.make_block_ptr( + Ai, (T, BT), (H * BT, 1), (i_t * BT + 48, 48), (16, 16), (1, 0) + ) + p_Ai_21 = tl.make_block_ptr( + Ai, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0) + ) + p_Ai_31 = tl.make_block_ptr( + Ai, (T, BT), (H * BT, 1), (i_t * BT + 32, 0), (16, 16), (1, 0) + ) + p_Ai_32 = tl.make_block_ptr( + Ai, (T, BT), (H * BT, 1), (i_t * BT + 32, 16), (16, 16), (1, 0) + ) + p_Ai_41 = tl.make_block_ptr( + Ai, (T, BT), (H * BT, 1), (i_t * BT + 48, 0), (16, 16), (1, 0) + ) + p_Ai_42 = tl.make_block_ptr( + Ai, (T, BT), (H * BT, 1), (i_t * BT + 48, 16), (16, 16), (1, 0) + ) + p_Ai_43 = tl.make_block_ptr( + Ai, (T, BT), (H * BT, 1), (i_t * BT + 48, 32), (16, 16), (1, 0) + ) + tl.store( + p_Ai_11, + b_Ai_11.to(p_Ai_11.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_22, + b_Ai_22.to(p_Ai_22.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_33, + b_Ai_33.to(p_Ai_33.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_44, + b_Ai_44.to(p_Ai_44.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_21, + b_Ai_21.to(p_Ai_21.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_31, + b_Ai_31.to(p_Ai_31.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_32, + b_Ai_32.to(p_Ai_32.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_41, + b_Ai_41.to(p_Ai_41.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_42, + b_Ai_42.to(p_Ai_42.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + tl.store( + p_Ai_43, + b_Ai_43.to(p_Ai_43.dtype.element_ty, fp_downcast_rounding="rtne"), + boundary_check=(0, 1), + ) + else: + desc_o.store( + [i_t * BT + 0, 0], b_Ai_11.to(desc_o.dtype, fp_downcast_rounding="rtne") + ) + desc_o.store( + [i_t * BT + 16, 16], b_Ai_22.to(desc_o.dtype, fp_downcast_rounding="rtne") + ) + desc_o.store( + [i_t * BT + 32, 32], b_Ai_33.to(desc_o.dtype, fp_downcast_rounding="rtne") + ) + desc_o.store( + [i_t * BT + 48, 48], b_Ai_44.to(desc_o.dtype, fp_downcast_rounding="rtne") + ) + desc_o.store( + [i_t * BT + 16, 0], b_Ai_21.to(desc_o.dtype, fp_downcast_rounding="rtne") + ) + desc_o.store( + [i_t * BT + 32, 0], b_Ai_31.to(desc_o.dtype, fp_downcast_rounding="rtne") + ) + desc_o.store( + [i_t * BT + 32, 16], b_Ai_32.to(desc_o.dtype, fp_downcast_rounding="rtne") + ) + desc_o.store( + [i_t * BT + 48, 0], b_Ai_41.to(desc_o.dtype, fp_downcast_rounding="rtne") + ) + desc_o.store( + [i_t * BT + 48, 16], b_Ai_42.to(desc_o.dtype, fp_downcast_rounding="rtne") + ) + desc_o.store( + [i_t * BT + 48, 32], b_Ai_43.to(desc_o.dtype, fp_downcast_rounding="rtne") + ) + + +@input_guard +def solve_tril( + A: torch.Tensor, + cu_seqlens: torch.Tensor | None = None, + output_dtype: torch.dtype = torch.float, +) -> torch.Tensor: + """ + Compute the inverse of the matrix I + A + A should be strictly lower triangular, i.e., A.triu() == 0. + + Args: + A (torch.Tensor): + [B, T, H, BT], where BT should only be 16, 32, or 64. + cu_seqlens (torch.Tensor): + The cumulative sequence lengths of the input tensor. Default: `None`. + output_dtype (torch.dtype): + The dtype of the output tensor. Default: `torch.float`. + If `None`, the output dtype will be the same as the input dtype. + + Returns: + (I + A)^-1 with the same shape as A + """ + assert A.shape[-1] in [16, 32, 64] + output_dtype = A.dtype if output_dtype is None else output_dtype + + B, T, H, BT = A.shape + chunk_indices = ( + prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + ) + NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, BT) + + Ai = torch.zeros_like(A, dtype=output_dtype) + if BT == 16: + merge_fn = solve_tril_16x16_kernel + elif BT == 32: + merge_fn = merge_16x16_to_32x32_inverse_kernel + elif BT == 64: + merge_fn = merge_16x16_to_64x64_inverse_kernel + + merge_fn[NT, B * H]( + A=A, + Ai=Ai, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + BT=BT, + USE_TMA=is_tma_supported, + DOT_PRECISION=FLA_TRIL_PRECISION, + ) + return Ai diff --git a/nanovllm/layers/ops/utils.py b/nanovllm/layers/ops/utils.py new file mode 100644 index 000000000..2ee4ddebb --- /dev/null +++ b/nanovllm/layers/ops/utils.py @@ -0,0 +1,157 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang +# +# This file contains code copied from the flash-linear-attention project. +# The original source code was licensed under the MIT license and included +# the following copyright notice: +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +import contextlib +import functools +import os +from collections.abc import Callable +from typing import Any + +import torch + +try: + import triton + HAS_TRITON = True +except ImportError: + HAS_TRITON = False + +SUPPRESS_LEVEL = int(os.getenv("GDN_RECOMPUTE_SUPPRESS_LEVEL", "0")) +FLA_GDN_FIX_BT = os.getenv("FLA_GDN_FIX_BT", "0") == "1" + +# Check if CUDA graphs should be used (simplified check) +use_cuda_graph = os.environ.get("FLA_USE_CUDA_GRAPH", "0") == "1" + +# Check if gather is supported +try: + import triton.language as tl + is_gather_supported = hasattr(tl, "gather") +except ImportError: + is_gather_supported = False + +# Simplified platform checks +try: + is_nvidia = torch.cuda.is_available() + is_nvidia_hopper = is_nvidia and ( + "NVIDIA H" in torch.cuda.get_device_name(0) + or torch.cuda.get_device_capability()[0] >= 9 + ) +except Exception: + is_nvidia = False + is_nvidia_hopper = False + +try: + is_amd = hasattr(torch.version, "hip") and torch.version.hip is not None +except Exception: + is_amd = False + +try: + is_intel = hasattr(torch, "xpu") and torch.xpu.is_available() +except Exception: + is_intel = False + +is_intel_alchemist = False +if is_intel: + try: + is_intel_alchemist = "Intel(R) Arc(TM) A" in torch.xpu.get_device_name(0) + except Exception: + pass + +# Check TMA support +try: + import triton.language as tl + is_tma_supported = ( + is_nvidia + and torch.cuda.get_device_capability(0)[0] >= 9 + and ( + hasattr(tl, "_experimental_make_tensor_descriptor") + or hasattr(tl, "make_tensor_descriptor") + ) + ) +except Exception: + is_tma_supported = False + +# Check shared memory +def check_shared_mem(arch: str = "none", tensor_idx: int = 0) -> bool: + """Simplified shared memory check.""" + # Default to True for most cases + return True + + +def tensor_cache(fn: Callable[..., torch.Tensor]) -> Callable[..., torch.Tensor]: + """ + A decorator that caches the most recent results of a function with tensor inputs. + """ + cache_entries: list[tuple[tuple | None, dict | None, Any]] = [] + cache_size = 8 + + @functools.wraps(fn) + def wrapper(*args: Any, **kwargs: Any) -> Any: + nonlocal cache_entries, cache_size + for i, entry in enumerate(cache_entries): + last_args, last_kwargs, last_result = entry + if ( + len(args) == len(last_args) + and len(kwargs) == len(last_kwargs) + and all(a is b for a, b in zip(args, last_args)) + and all( + k in last_kwargs and v is last_kwargs[k] for k, v in kwargs.items() + ) + ): + cache_entries = ( + cache_entries[:i] + + cache_entries[i + 1 :] + + [(args, kwargs, last_result)] + ) + return last_result + + result = fn(*args, **kwargs) + + if len(cache_entries) >= cache_size: + cache_entries = cache_entries[1:] + cache_entries.append((args, kwargs, result)) + return result + + return wrapper + + +def input_guard(fn: Callable[..., torch.Tensor]) -> Callable[..., torch.Tensor]: + """ + A decorator to make sure all input tensors are contiguous and set the device based on input tensors. + """ + + @functools.wraps(fn) + def wrapper(*args, **kwargs): + contiguous_args = ( + i if not isinstance(i, torch.Tensor) else i.contiguous() for i in args + ) + contiguous_kwargs = { + k: (v if not isinstance(v, torch.Tensor) else v.contiguous()) + for k, v in kwargs.items() + } + + tensor = None + for arg in args: + if isinstance(arg, torch.Tensor): + tensor = arg + break + if tensor is None: + for value in kwargs.values(): + if isinstance(value, torch.Tensor): + tensor = value + break + + if tensor is not None: + ctx = torch.cuda.device(tensor.device.index) + else: + ctx = contextlib.nullcontext() + + with ctx: + return fn(*contiguous_args, **contiguous_kwargs) + + return wrapper diff --git a/nanovllm/layers/ops/wy_fast.py b/nanovllm/layers/ops/wy_fast.py new file mode 100644 index 000000000..b01b89fc7 --- /dev/null +++ b/nanovllm/layers/ops/wy_fast.py @@ -0,0 +1,158 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang +# +# This file contains code copied from the flash-linear-attention project. +# The original source code was licensed under the MIT license and included +# the following copyright notice: +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +# ruff: noqa: E501 + +import torch +import triton +import triton.language as tl + +from nanovllm.layers.ops.index import prepare_chunk_indices + + +@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None}) +@triton.autotune( + configs=[ + triton.Config({}, num_warps=num_warps, num_stages=num_stages) + for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] + ], + key=["H", "K", "V", "BT", "BK", "BV", "IS_VARLEN"], +) +@triton.jit(do_not_specialize=["T"]) +def recompute_w_u_fwd_kernel( + k, + v, + beta, + w, + u, + A, + g, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + Hg: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BT: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_t, i_bh = tl.program_id(0), tl.program_id(1) + i_b, i_h = i_bh // H, i_bh % H + if IS_VARLEN: + i_n, i_t = ( + tl.load(chunk_indices + i_t * 2).to(tl.int32), + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32), + ) + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + p_beta = tl.make_block_ptr( + beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,) + ) + p_g = tl.make_block_ptr(g + (bos * H + i_h), (T,), (H,), (i_t * BT,), (BT,), (0,)) + p_A = tl.make_block_ptr( + A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0) + ) + b_beta = tl.load(p_beta, boundary_check=(0,)) + b_A = tl.load(p_A, boundary_check=(0, 1)) + b_g = tl.exp(tl.load(p_g, boundary_check=(0,))) + + for i_v in range(tl.cdiv(V, BV)): + p_v = tl.make_block_ptr( + v + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + p_u = tl.make_block_ptr( + u + (bos * H + i_h) * V, + (T, V), + (H * V, 1), + (i_t * BT, i_v * BV), + (BT, BV), + (1, 0), + ) + b_v = tl.load(p_v, boundary_check=(0, 1)) + b_vb = (b_v * b_beta[:, None]).to(b_v.dtype) + b_u = tl.dot(b_A, b_vb, allow_tf32=False) + tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1)) + + for i_k in range(tl.cdiv(K, BK)): + p_k = tl.make_block_ptr( + k + (bos * Hg + i_h // (H // Hg)) * K, + (T, K), + (Hg * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + p_w = tl.make_block_ptr( + w + (bos * H + i_h) * K, + (T, K), + (H * K, 1), + (i_t * BT, i_k * BK), + (BT, BK), + (1, 0), + ) + b_k = tl.load(p_k, boundary_check=(0, 1)) + b_kb = (b_k * b_beta[:, None] * b_g[:, None]).to(b_k.dtype) + b_w = tl.dot(b_A, b_kb) + tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1)) + + +def recompute_w_u_fwd( + k: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + g_cumsum: torch.Tensor, + A: torch.Tensor, + cu_seqlens: torch.LongTensor | None, +) -> tuple[torch.Tensor, torch.Tensor]: + B, T, Hg, K, V = *k.shape, v.shape[-1] + H = v.shape[-2] + BT = A.shape[-1] + + chunk_indices = ( + prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None + ) + NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + BK = 64 + BV = 64 + u = torch.empty_like(v) + w = k.new_empty(B, T, H, K) + recompute_w_u_fwd_kernel[(NT, B * H)]( + k=k, + v=v, + beta=beta, + w=w, + u=u, + A=A, + g=g_cumsum, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + Hg=Hg, + K=K, + V=V, + BT=BT, + BK=BK, + BV=BV, + ) + return w, u diff --git a/nanovllm/layers/sampler.py b/nanovllm/layers/sampler.py index 41838aca4..b101018a1 100644 --- a/nanovllm/layers/sampler.py +++ b/nanovllm/layers/sampler.py @@ -4,6 +4,9 @@ class Sampler(nn.Module): + def __init__(self): + super().__init__() + @torch.compile def forward(self, logits: torch.Tensor, temperatures: torch.Tensor): logits = logits.float().div_(temperatures.unsqueeze(dim=1)) diff --git a/nanovllm/models/qwen3_5.py b/nanovllm/models/qwen3_5.py new file mode 100644 index 000000000..34f9ab98f --- /dev/null +++ b/nanovllm/models/qwen3_5.py @@ -0,0 +1,2294 @@ +"""Qwen3.5 multimodal model (vision + language). + +This module implements the Qwen3.5 vision-language model used by +nano-vllm-multimodal-qwen3_5. It does not implement or depend on Qwen3-VL. +""" + +from __future__ import annotations + +import os + +# For diagnose_full_attention.py: when NANOVLLM_DEBUG_FA=1, layer 3 full_attention saves intermediates here +FA_DEBUG_SAVE = {} +import math +from typing import List, Optional, Sequence, Tuple + +import torch +import torch.distributed as dist +import torch.nn.functional as F +from torch import nn + +from nanovllm.layers.activation import SiluAndMul +from nanovllm.layers.attention import Attention +from nanovllm.layers.embed_head import ParallelLMHead, VocabParallelEmbedding +from nanovllm.layers.layernorm import RMSNorm +from nanovllm.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from nanovllm.layers.rotary_embedding import get_rope +from nanovllm.models.qwen3_next import ( + Qwen3NextRMSNorm, + Qwen3NextRMSNormGated, + torch_causal_conv1d_update, + torch_recurrent_gated_delta_rule, + torch_chunk_gated_delta_rule, + mamba_v2_sharded_weight_loader, +) +from nanovllm.utils.loader import sharded_weight_loader + +# Triton FLA ops (same as vLLM fla/ops): prefer for numerical stability +try: + from nanovllm.layers.ops import chunk_gated_delta_rule as fla_chunk_gated_delta_rule + from nanovllm.layers.ops import fused_recurrent_gated_delta_rule as fla_fused_recurrent_gated_delta_rule + _HAS_FLA_TRITON = fla_chunk_gated_delta_rule is not None and fla_fused_recurrent_gated_delta_rule is not None +except Exception: + fla_chunk_gated_delta_rule = None + fla_fused_recurrent_gated_delta_rule = None + _HAS_FLA_TRITON = False + + +# --------------------------------------------------------------------------- +# Text backbone (supports DeepStack injection) +# --------------------------------------------------------------------------- + + +class Qwen3_5TextRotaryEmbedding(nn.Module): + """Rotary Embedding for Qwen3_5 with MRoPE support""" + + def __init__(self, config, device=None): + super().__init__() + self.max_seq_len_cached = config.max_position_embeddings + self.config = config + + rope_parameters = getattr(config, "rope_parameters", {}) + if isinstance(rope_parameters, dict): + self.rope_type = rope_parameters.get("rope_type", "default") + rope_theta = rope_parameters.get("rope_theta", getattr(config, "rope_theta", 1000000)) + self.mrope_section = rope_parameters.get("mrope_section", [11, 11, 10]) + partial_rotary_factor = rope_parameters.get("partial_rotary_factor", 1.0) + else: + # RopeParameters object or similar (e.g. from HF config) + self.rope_type = getattr(rope_parameters, "rope_type", "default") + rope_theta = getattr(rope_parameters, "rope_theta", getattr(config, "rope_theta", 1000000)) + self.mrope_section = getattr(rope_parameters, "mrope_section", [11, 11, 10]) + partial_rotary_factor = getattr(rope_parameters, "partial_rotary_factor", 1.0) + head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads + dim = int(head_dim * partial_rotary_factor) + + # Compute inverse frequencies + inv_freq = 1.0 / ( + rope_theta ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim) + ) + self.register_buffer("inv_freq", inv_freq, persistent=False) + self.attention_scaling = 1.0 + + def apply_interleaved_mrope(self, freqs, mrope_section): + """Apply interleaved MRoPE to 3D rotary embeddings. + Reorganizes frequency layout from chunked [TTT...HHH...WWW] to + interleaved [THWTHWTHW...TT], preserving frequency continuity. + args: + freqs: (3, bs, seq_len, head_dim // 2) + mrope_section: (3,) + returns: + freqs_t: (bs, seq_len, head_dim // 2) + """ + freqs_t = freqs[0] # just overwrite the first dimension T + for dim, offset in enumerate((1, 2), start=1): # H, W + length = mrope_section[dim] * 3 + idx = slice(offset, length, 3) + freqs_t[..., idx] = freqs[dim, ..., idx] + return freqs_t + + def forward(self, x, position_ids): + # Qwen3_5 MRoPE: position_ids can be (3, seq_len) or (3, batch, seq_len) for 3D (T,H,W), + # or (batch, seq_len) for 1D which we expand to (3, batch, seq_len). + if position_ids.ndim == 2: + if position_ids.shape[0] == 3: + # (3, seq_len) -> (3, 1, seq_len) + position_ids = position_ids.unsqueeze(1) + else: + # (batch, seq_len) -> (3, batch, seq_len) + position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1) + # position_ids now (3, batch, seq_len) + inv_freq_expanded = self.inv_freq[None, None, :, None].float().expand(3, position_ids.shape[1], -1, 1) + position_ids_expanded = position_ids[:, :, None, :].float() # shape (3, bs, 1, positions) + + # Force float32 computation + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(2, 3) + freqs = self.apply_interleaved_mrope(freqs, self.mrope_section) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() * self.attention_scaling + sin = emb.sin() * self.attention_scaling + + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +class Qwen3_5TextAttentionMerged(nn.Module): + + def __init__( + self, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + max_position: int, + rms_norm_eps: float, + qkv_bias: bool, + head_dim: int | None, + rope_theta: float, + rope_scaling: tuple | None, + ) -> None: + super().__init__() + tp_size = dist.get_world_size() + self.total_num_heads = num_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = num_kv_heads + assert self.total_num_kv_heads % tp_size == 0 + self.num_kv_heads = self.total_num_kv_heads // tp_size + self.head_dim = head_dim or hidden_size // self.total_num_heads + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim ** -0.5 + self.qkv_bias = qkv_bias + + self.qkv_proj = QKVParallelLinear( + hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=qkv_bias, + ) + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + hidden_size, + bias=False, + ) + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.head_dim, + max_position=max_position, + base=rope_theta, + rope_scaling=rope_scaling, + ) + self.attn = Attention( + self.num_heads, + self.head_dim, + self.scaling, + self.num_kv_heads, + ) + if not self.qkv_bias: + self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps) + self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + hidden_states = hidden_states.to(self.qkv_proj.weight.dtype) + qkv = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q = q.reshape(-1, self.num_heads, self.head_dim) + k = k.reshape(-1, self.num_kv_heads, self.head_dim) + v = v.reshape(-1, self.num_kv_heads, self.head_dim) + if not self.qkv_bias: + q = self.q_norm(q) + k = self.k_norm(k) + q, k = self.rotary_emb(positions, q, k) + o = self.attn(q, k, v) + output = self.o_proj(o.flatten(1, -1)) + return output + + +class Qwen3_5TextAttention(nn.Module): + """Qwen3_5 Text Attention with gate mechanism (similar to Qwen3Next)""" + + def __init__( + self, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + max_position: int, + rms_norm_eps: float, + qkv_bias: bool, + head_dim: int | None, + rope_theta: float, + rope_scaling: tuple | None, + config=None, + ) -> None: + super().__init__() + tp_size = dist.get_world_size() + self.total_num_heads = num_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = num_kv_heads + assert self.total_num_kv_heads % tp_size == 0 + self.num_kv_heads = self.total_num_kv_heads // tp_size + self.head_dim = head_dim or hidden_size // self.total_num_heads + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim ** -0.5 + self.qkv_bias = qkv_bias + + # q_proj outputs num_heads * head_dim * 2 (query + gate) + self.q_proj = ColumnParallelLinear( + hidden_size, + self.total_num_heads * self.head_dim * 2, + bias=qkv_bias, + ) + self.k_proj = ColumnParallelLinear( + hidden_size, + self.total_num_kv_heads * self.head_dim, + bias=qkv_bias, + ) + self.v_proj = ColumnParallelLinear( + hidden_size, + self.total_num_kv_heads * self.head_dim, + bias=qkv_bias, + ) + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + hidden_size, + bias=False, + ) + + # Use Qwen3_5TextRotaryEmbedding if config is provided, otherwise fallback to standard RoPE + if config is not None: + self.rotary_emb = Qwen3_5TextRotaryEmbedding(config, device=None) + else: + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.head_dim, + max_position=max_position, + base=rope_theta, + rope_scaling=rope_scaling, + ) + + self.attn = Attention( + self.num_heads, + self.head_dim, + self.scaling, + self.num_kv_heads, + ) + # Always use q_norm and k_norm for Qwen3_5 + self.q_norm = Qwen3NextRMSNorm(self.head_dim, eps=rms_norm_eps) + self.k_norm = Qwen3NextRMSNorm(self.head_dim, eps=rms_norm_eps) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + hidden_states = hidden_states.to(self.q_proj.weight.dtype) + # q_proj outputs query + gate (concatenated) + q_gate = self.q_proj(hidden_states) + q_gate = q_gate.reshape(-1, self.num_heads, self.head_dim * 2) + query_states, gate = q_gate.split([self.head_dim, self.head_dim], dim=-1) + k = self.k_proj(hidden_states) + k = k.reshape(-1, self.num_kv_heads, self.head_dim) + v = self.v_proj(hidden_states) + v = v.reshape(-1, self.num_kv_heads, self.head_dim) + # Apply normalization (Qwen3NextRMSNorm returns (output, x) when called with one arg) + query_states, _ = self.q_norm(query_states) + k, _ = self.k_norm(k) + + # Apply rotary embedding + # For Qwen3_5TextRotaryEmbedding, we need to handle MRoPE differently + # For now, fallback to standard RoPE if not using MRoPE + if isinstance(self.rotary_emb, Qwen3_5TextRotaryEmbedding): + # For MRoPE, positions should be 3D (temporal, height, width) + # But in nano-vllm, positions is typically 1D, so we use standard RoPE for now + # TODO: Implement proper MRoPE support with 3D position_ids + # For now, convert positions to 2D format expected by Qwen3_5TextRotaryEmbedding + if positions.ndim == 1: + # Convert 1D positions to 2D (batch_size=1, seq_len) + positions_2d = positions.unsqueeze(0) + else: + positions_2d = positions + + cos, sin = self.rotary_emb(hidden_states, positions_2d) + # Apply rotary embedding manually (partial RoPE: cos/sin only cover rotary_dim dims) + rotary_dim = cos.shape[-1] + def rotate_half(x): + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + # cos/sin shape: (bs, seq_len, rotary_dim); query_states/k shape: (seq_len, num_heads, head_dim) + if cos.ndim == 3: + cos = cos.squeeze(0).unsqueeze(1) # (seq_len, 1, rotary_dim) + sin = sin.squeeze(0).unsqueeze(1) + else: + cos = cos.unsqueeze(1) + sin = sin.unsqueeze(1) + q_rot = (query_states[..., :rotary_dim] * cos) + (rotate_half(query_states[..., :rotary_dim]) * sin) + k_rot = (k[..., :rotary_dim] * cos) + (rotate_half(k[..., :rotary_dim]) * sin) + query_states = torch.cat([q_rot, query_states[..., rotary_dim:]], dim=-1) + k = torch.cat([k_rot, k[..., rotary_dim:]], dim=-1) + else: + query_states, k = self.rotary_emb(positions, query_states, k) + + # Attention + o = self.attn(query_states, k, v) + + # Reshape gate and apply sigmoid + gate = gate.flatten(1, -1) # (seq_len, num_heads * head_dim) + o_flat = o.flatten(1, -1) # (seq_len, num_heads * head_dim) + attn_output = o_flat * torch.sigmoid(gate) + + output = self.o_proj(attn_output) + return output + + +class Qwen3_5TextMLPMerged(nn.Module): + + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + ) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + ) + assert hidden_act == "silu" + self.act_fn = SiluAndMul() + + def forward(self, x): + x = x.to(self.gate_up_proj.weight.dtype) + gate_up = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x = self.down_proj(x) + return x + + +class Qwen3_5GatedDeltaNet(nn.Module): + """Gated Delta Net for linear attention in Qwen3_5 + + Qwen3_5 uses separate projections for qkv, z, b, a instead of combined qkvz and ba. + """ + + def __init__( + self, + config, + layer_idx: int = 0, + ) -> None: + super().__init__() + tp_size = dist.get_world_size() + self.hidden_size = config.hidden_size + self.num_v_heads = getattr(config, "linear_num_value_heads", 32) + self.num_k_heads = getattr(config, "linear_num_key_heads", 16) + self.head_k_dim = getattr(config, "linear_key_head_dim", 128) + self.head_v_dim = getattr(config, "linear_value_head_dim", 128) + self.key_dim = self.head_k_dim * self.num_k_heads + self.value_dim = self.head_v_dim * self.num_v_heads + + self.conv_kernel_size = getattr(config, "linear_conv_kernel_dim", 4) + self.layer_idx = layer_idx + self.activation = getattr(config, "hidden_act", "silu") + self.layer_norm_epsilon = config.rms_norm_eps + + # QKV projection + self.conv_dim = self.key_dim * 2 + self.value_dim + + # Conv1d layer (for causal convolution) + # Use ColumnParallelLinear instead of nn.Conv1d for tensor parallelism + self.conv1d = ColumnParallelLinear( + input_size=self.conv_kernel_size, + output_size=self.conv_dim, + bias=False, + ) + # Add dimension for conv1d: (conv_dim, kernel_size) -> (conv_dim, 1, kernel_size) + self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1) + + # Set up custom weight loader for conv1d + query_key_settings = (self.key_dim, 0, False) + value_settings = (self.value_dim, 0, False) + + # Remove default weight_loader and set custom one + if hasattr(self.conv1d.weight, "weight_loader"): + delattr(self.conv1d.weight, "weight_loader") + + # Set custom weight loader + self.conv1d.weight.weight_loader = mamba_v2_sharded_weight_loader( + [ + query_key_settings, # Q projection + query_key_settings, # K projection + value_settings, # V projection + ], + tp_size, + dist.get_rank(), + ) + + # Projection layers - Qwen3_5 uses separate projections + self.in_proj_qkv = MergedColumnParallelLinear( + self.hidden_size, + [self.key_dim, self.key_dim, self.value_dim], + bias=False, + ) + self.in_proj_z = ColumnParallelLinear( + self.hidden_size, + self.value_dim, + bias=False, + ) + self.in_proj_b = ColumnParallelLinear( + self.hidden_size, + self.num_v_heads, + bias=False, + ) + self.in_proj_a = ColumnParallelLinear( + self.hidden_size, + self.num_v_heads, + bias=False, + ) + + # Time step projection parameters (sharded by TP along dim 0, vLLM-style) + self.dt_bias = nn.Parameter(torch.ones(self.num_v_heads // tp_size)) + A = torch.empty(self.num_v_heads // tp_size).uniform_(0, 16) + self.A_log = nn.Parameter(torch.log(A)) + self.A_log.weight_loader = sharded_weight_loader(0) + self.dt_bias.weight_loader = sharded_weight_loader(0) + + # Normalization with gate + self.norm = Qwen3NextRMSNormGated(self.head_v_dim, eps=self.layer_norm_epsilon) + + # Output projection + self.out_proj = RowParallelLinear( + self.value_dim, + self.hidden_size, + bias=False, + ) + + # Cache states (for incremental decoding) + # HF-style: per-sequence isolation via state_cache when sequence_id is provided + self.conv_state = None + self.recurrent_state = None + self.state_cache = {} # sequence_id -> (conv_state, recurrent_state) + self._decode_step_counter = {} # sequence_id -> decode steps (for debug) + + # CUDA graph decode state buffers: persistent tensors so in-place updates + # are captured by the graph and replayed correctly. Size must be >= + # max_num_seqs (we pick 512 as a safe upper bound for nano-vllm). + max_graph_bs = 512 + conv_dim_tp = self.conv_dim // tp_size + num_v_heads_tp = self.num_v_heads // tp_size + self.register_buffer( + "_graph_conv_state", + torch.zeros( + max_graph_bs, + conv_dim_tp, + self.conv_kernel_size - 1, + dtype=torch.bfloat16, + ), + persistent=False, + ) + self.register_buffer( + "_graph_recurrent_state", + torch.zeros( + max_graph_bs, + num_v_heads_tp, + self.head_v_dim, + self.head_k_dim, + dtype=torch.bfloat16, + ), + persistent=False, + ) + # Pool buffer for conv state: indexed by slot, avoids per-seq gather/scatter + self.register_buffer( + "_pool_conv_state", + torch.zeros( + max_graph_bs, + conv_dim_tp, + self.conv_kernel_size - 1, + dtype=torch.bfloat16, + ), + persistent=False, + ) + # Persistent ssm_state_indices: updated in-place before CUDA graph replay + self.register_buffer( + "_graph_ssm_state_indices", + torch.arange(max_graph_bs, dtype=torch.int64), + persistent=False, + ) + + def reset_state(self): + """Reset all cached states. Called after warmup to avoid state pollution.""" + self.conv_state = None + self.recurrent_state = None + self.state_cache.clear() + self._decode_step_counter.clear() + self._graph_conv_state.zero_() + self._graph_recurrent_state.zero_() + self._pool_conv_state.zero_() + self._graph_ssm_state_indices.copy_(torch.arange(self._graph_ssm_state_indices.shape[0], dtype=torch.int64, device=self._graph_ssm_state_indices.device)) + + def _get_states(self, sequence_id: int | None): + """Get conv_state and recurrent_state for this sequence.""" + if sequence_id is not None: + if sequence_id not in self.state_cache: + self.state_cache[sequence_id] = (None, None) + return self.state_cache[sequence_id] + return self.conv_state, self.recurrent_state + + def _set_states(self, conv_state, recurrent_state, sequence_id: int | None): + """Store conv_state and recurrent_state for this sequence.""" + if sequence_id is not None: + self.state_cache[sequence_id] = (conv_state, recurrent_state) + else: + self.conv_state = conv_state + self.recurrent_state = recurrent_state + + def _forward_graph( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + """Graph-mode forward for batch decode using persistent state buffers. + + Args: + hidden_states: (bs, hidden_size) - batch of decode tokens + + Returns: + output: (bs, hidden_size) + """ + bs = hidden_states.size(0) + hidden_states = hidden_states.to(self.in_proj_qkv.weight.dtype) + tp_size = dist.get_world_size() + conv_dim_tp = self.conv_dim // tp_size + + # Persistent graph state buffers (in-place updates are captured by CUDA graph) + conv_state = self._graph_conv_state[:bs] + recurrent_state = self._graph_recurrent_state[:bs] + + # Part 1: Input Projection + mixed_qkv = self.in_proj_qkv(hidden_states) + if isinstance(mixed_qkv, tuple): + mixed_qkv = mixed_qkv[0] + + z = self.in_proj_z(hidden_states) + if isinstance(z, tuple): + z = z[0] + z = z.reshape(bs, -1, self.head_v_dim) + + b = self.in_proj_b(hidden_states) + if isinstance(b, tuple): + b = b[0] + a = self.in_proj_a(hidden_states) + if isinstance(a, tuple): + a = a[0] + + b = b.contiguous() + a = a.contiguous() + + query, key, value = torch.split( + mixed_qkv, + [self.key_dim // tp_size, self.key_dim // tp_size, self.value_dim // tp_size], + dim=-1 + ) + + query = query.reshape(bs, self.num_k_heads // tp_size, self.head_k_dim) + key = key.reshape(bs, self.num_k_heads // tp_size, self.head_k_dim) + value = value.reshape(bs, self.num_v_heads // tp_size, self.head_v_dim) + + query_flat = query.reshape(bs, -1) + key_flat = key.reshape(bs, -1) + value_flat = value.reshape(bs, -1) + + mixed_qkv_conv = torch.cat([query_flat, key_flat, value_flat], dim=-1) + # causal_conv1d_update expects (batch, dim, seq_len); decode seq_len=1 + mixed_qkv_conv = mixed_qkv_conv.unsqueeze(-1) + + conv_weights = self.conv1d.weight.reshape( + self.conv1d.weight.size(0), self.conv1d.weight.size(2) + ) + + mixed_qkv_conv = torch_causal_conv1d_update( + mixed_qkv_conv, + conv_state, + conv_weights, + self.conv1d.bias, + self.activation, + ) + # conv_state is updated in-place by causal_conv1d_update + mixed_qkv_conv = mixed_qkv_conv.squeeze(-1) + + query_conv, key_conv, value_conv = torch.split( + mixed_qkv_conv, + [self.key_dim // tp_size, self.key_dim // tp_size, self.value_dim // tp_size], + dim=-1 + ) + + query_conv = query_conv.reshape(bs, self.num_k_heads // tp_size, self.head_k_dim) + key_conv = key_conv.reshape(bs, self.num_k_heads // tp_size, self.head_k_dim) + value_conv = value_conv.reshape(bs, self.num_v_heads // tp_size, self.head_v_dim) + + if self.num_v_heads // tp_size > self.num_k_heads // tp_size: + repeat_factor = (self.num_v_heads // tp_size) // (self.num_k_heads // tp_size) + query_conv = query_conv.repeat_interleave(repeat_factor, dim=1) + key_conv = key_conv.repeat_interleave(repeat_factor, dim=1) + + # Gating + beta = b.sigmoid() + g = -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias.float()) + g = g.to(hidden_states.dtype) + g = g.clamp(min=-50.0) + g = torch.nan_to_num(g, nan=0.0, posinf=0.0, neginf=-50.0) + + use_float32 = bool(os.environ.get("NANOVLLM_GDN_FLOAT32")) + if use_float32: + query_conv = query_conv.float() + key_conv = key_conv.float() + value_conv = value_conv.float() + g = g.float() + beta = beta.float() + + # Triton recurrent decode path (batch) + q_batch = query_conv.unsqueeze(1).contiguous() + k_batch = key_conv.unsqueeze(1).contiguous() + v_batch = value_conv.unsqueeze(1).contiguous() + g_batch = g.unsqueeze(1).contiguous() + beta_batch = beta.unsqueeze(1).contiguous() + + use_triton = ( + not os.environ.get("NANOVLLM_FORCE_TORCH_GDN") + and not use_float32 + and _HAS_FLA_TRITON + and hidden_states.dtype != torch.float32 + ) + + if use_triton: + # Persistent indices buffer — model_runner updates this in-place before replay + # so CUDA graph reads correct pool slot mappings for each sequence. + ssm_state_indices = self._graph_ssm_state_indices[:bs] + core_attn_out_batch, _ = fla_fused_recurrent_gated_delta_rule( + q=q_batch, + k=k_batch, + v=v_batch, + g=g_batch, + beta=beta_batch, + initial_state=recurrent_state, + inplace_final_state=True, + ssm_state_indices=ssm_state_indices, + use_qk_l2norm_in_kernel=True, + ) + # recurrent_state is updated in-place + else: + # Torch fallback: initial_state expects (N, H, K, V) + recurrent_state_torch = recurrent_state.permute(0, 1, 3, 2).contiguous() + core_attn_out_batch, recurrent_state_torch = torch_recurrent_gated_delta_rule( + query_conv, + key_conv, + value_conv, + g, + beta, + initial_state=recurrent_state_torch, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + ) + # Copy back to graph buffer in-place so CUDA graph captures it + recurrent_state[:, :, :, :] = recurrent_state_torch.permute(0, 1, 3, 2).contiguous() + + core_attn_out = core_attn_out_batch.squeeze(1) + + if use_float32: + core_attn_out = core_attn_out.to(hidden_states.dtype) + + # Output projection with gate + core_attn_out_flat = core_attn_out.reshape(-1, core_attn_out.shape[-1]) + z_flat = z.reshape(-1, z.shape[-1]) + core_attn_out_flat = self.norm(core_attn_out_flat, z_flat) + core_attn_out = core_attn_out_flat.reshape(bs, -1, self.head_v_dim) + core_attn_out = core_attn_out.reshape(bs, -1) + + output = self.out_proj(core_attn_out) + if isinstance(output, tuple): + output = output[0] + + return output + + def forward( + self, + hidden_states: torch.Tensor, + positions: torch.Tensor | None = None, + sequence_id: int | None = None, + use_graph: bool = False, + ) -> torch.Tensor: + """ + Forward pass for GatedDeltaNet (HF-style) + + Args: + hidden_states: (seq_len, hidden_size) - ONE sequence only (no mixed batch) + positions: (3, seq_len) for MRoPE + sequence_id: for batch inference, isolates state per sequence + use_graph: if True, use graph-mode batch path (no sequence_id dict) + + Returns: + output: (seq_len, hidden_size) + """ + if use_graph: + return self._forward_graph(hidden_states) + + seq_len = hidden_states.size(0) + hidden_states = hidden_states.to(self.in_proj_qkv.weight.dtype) + + conv_state, recurrent_state = self._get_states(sequence_id) + + def _gdn_debug_this_layer(): + if not os.environ.get("NANOVLLM_DEBUG_GDN"): + return False + try: + return self.layer_idx == int(os.environ.get("NANOVLLM_DEBUG_GDN_LAYER", "0")) + except ValueError: + return self.layer_idx == 0 + + # Part 1: Input Projection (separate projections for Qwen3_5) + mixed_qkv = self.in_proj_qkv(hidden_states) + # mixed_qkv is a tuple/list from MergedColumnParallelLinear + if isinstance(mixed_qkv, tuple): + mixed_qkv = mixed_qkv[0] + if _gdn_debug_this_layer(): + print(f"[NV GDN L{self.layer_idx}] mixed_qkv: shape={list(mixed_qkv.shape)}, mean={mixed_qkv.mean():.6f}, std={mixed_qkv.std():.6f}, min={mixed_qkv.min():.6f}, max={mixed_qkv.max():.6f}") + + z = self.in_proj_z(hidden_states) + if isinstance(z, tuple): + z = z[0] + z = z.reshape(seq_len, -1, self.head_v_dim) + if _gdn_debug_this_layer(): + print(f"[NV GDN L{self.layer_idx}] z after proj: shape={list(z.shape)}, mean={z.mean():.6f}, std={z.std():.6f}") + + b = self.in_proj_b(hidden_states) + if isinstance(b, tuple): + b = b[0] + a = self.in_proj_a(hidden_states) + if isinstance(a, tuple): + a = a[0] + if _gdn_debug_this_layer(): + print(f"[NV GDN L{self.layer_idx}] b after proj: shape={list(b.shape)}, mean={b.mean():.6f}, std={b.std():.6f}") + print(f"[NV GDN L{self.layer_idx}] a after proj: shape={list(a.shape)}, mean={a.mean():.6f}, std={a.std():.6f}") + + b = b.contiguous() + a = a.contiguous() + + # Split mixed_qkv into query, key, value + query, key, value = torch.split( + mixed_qkv, + [self.key_dim // dist.get_world_size(), + self.key_dim // dist.get_world_size(), + self.value_dim // dist.get_world_size()], + dim=-1 + ) + + # Reshape to (seq_len, num_heads, head_dim) + query = query.reshape(seq_len, self.num_k_heads // dist.get_world_size(), self.head_k_dim) + key = key.reshape(seq_len, self.num_k_heads // dist.get_world_size(), self.head_k_dim) + value = value.reshape(seq_len, self.num_v_heads // dist.get_world_size(), self.head_v_dim) + # Flatten for convolution + query_flat = query.reshape(seq_len, -1) + key_flat = key.reshape(seq_len, -1) + value_flat = value.reshape(seq_len, -1) + + mixed_qkv_conv = torch.cat([query_flat, key_flat, value_flat], dim=-1) + # Transpose for conv1d: (seq_len, conv_dim) -> (conv_dim, seq_len) + mixed_qkv_conv = mixed_qkv_conv.transpose(0, 1).unsqueeze(0) # (1, conv_dim, seq_len) + + # Part 2: Causal Convolution + # HF-style: per-sequence conv_state (isolated by sequence_id) + if conv_state is None: + conv_state = F.pad(mixed_qkv_conv, (self.conv_kernel_size - 1, 0))[:, :, :self.conv_kernel_size - 1].contiguous() + + # Apply causal convolution + conv_weights = self.conv1d.weight.reshape( + self.conv1d.weight.size(0), self.conv1d.weight.size(2) + ) + mixed_qkv_conv = torch_causal_conv1d_update( + mixed_qkv_conv, + conv_state, + conv_weights, # (conv_dim // tp_size, kernel_size) + self.conv1d.bias, + self.activation, + ) + + # Transpose back: (1, conv_dim, seq_len) -> (seq_len, conv_dim) + mixed_qkv_conv = mixed_qkv_conv.squeeze(0).transpose(0, 1) + if _gdn_debug_this_layer(): + print(f"[NV GDN L{self.layer_idx}] after conv1d: shape={list(mixed_qkv_conv.shape)}, mean={mixed_qkv_conv.mean():.6f}, std={mixed_qkv_conv.std():.6f}, min={mixed_qkv_conv.min():.6f}, max={mixed_qkv_conv.max():.6f}") + # Scheme C: numerical stability check after conv1d + if os.environ.get("NANOVLLM_DEBUG_STABILITY"): + m_max = mixed_qkv_conv.abs().max().item() + has_nan = torch.isnan(mixed_qkv_conv).any().item() + has_inf = torch.isinf(mixed_qkv_conv).any().item() + if has_nan or has_inf or m_max > 1e4: + print(f"[GDN layer{self.layer_idx}] post-conv1d: max_abs={m_max:.4f} has_nan={has_nan} has_inf={has_inf}") + + # Split back into query, key, value + query_conv, key_conv, value_conv = torch.split( + mixed_qkv_conv, + [self.key_dim // dist.get_world_size(), + self.key_dim // dist.get_world_size(), + self.value_dim // dist.get_world_size()], + dim=-1 + ) + + # Reshape back to (seq_len, num_heads, head_dim) + query_conv = query_conv.reshape(seq_len, self.num_k_heads // dist.get_world_size(), self.head_k_dim) + key_conv = key_conv.reshape(seq_len, self.num_k_heads // dist.get_world_size(), self.head_k_dim) + value_conv = value_conv.reshape(seq_len, self.num_v_heads // dist.get_world_size(), self.head_v_dim) + if _gdn_debug_this_layer(): + print(f"[NV GDN L{self.layer_idx}] query: shape={list(query_conv.shape)}, mean={query_conv.mean():.6f}, std={query_conv.std():.6f}") + print(f"[NV GDN L{self.layer_idx}] key: shape={list(key_conv.shape)}, mean={key_conv.mean():.6f}, std={key_conv.std():.6f}") + print(f"[NV GDN L{self.layer_idx}] value: shape={list(value_conv.shape)}, mean={value_conv.mean():.6f}, std={value_conv.std():.6f}") + # Repeat query/key if num_v_heads > num_k_heads + if self.num_v_heads // dist.get_world_size() > self.num_k_heads // dist.get_world_size(): + repeat_factor = (self.num_v_heads // dist.get_world_size()) // (self.num_k_heads // dist.get_world_size()) + query_conv = query_conv.repeat_interleave(repeat_factor, dim=1) + key_conv = key_conv.repeat_interleave(repeat_factor, dim=1) + + # Compute gating parameters + beta = b.sigmoid() # (seq_len, num_v_heads) + if _gdn_debug_this_layer(): + print(f"[NV GDN L{self.layer_idx}] beta (after sigmoid): shape={list(beta.shape)}, mean={beta.mean():.6f}, std={beta.std():.6f}") + # g = -exp(A_log) * softplus(a + dt_bias) + g = -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias.float()) + g = g.to(hidden_states.dtype) + # Numerical stability: avoid -inf/NaN + g = g.clamp(min=-50.0) + g = torch.nan_to_num(g, nan=0.0, posinf=0.0, neginf=-50.0) + if _gdn_debug_this_layer(): + print(f"[NV GDN L{self.layer_idx}] g: shape={list(g.shape)}, mean={g.mean():.6f}, std={g.std():.6f}, min={g.min():.6f}, max={g.max():.6f}") + # Scheme C: stability check for g/beta + if os.environ.get("NANOVLLM_DEBUG_STABILITY"): + g_min, g_max = g.min().item(), g.max().item() + if torch.isnan(g).any() or torch.isinf(g).any() or g_min < -100 or g_max > 100: + print(f"[GDN layer{self.layer_idx}] gating: g_range=[{g_min:.4f},{g_max:.4f}]") + + # Diagnostic: check for NaN before chunk (all layers, when NANOVLLM_DEBUG_LOGITS) + if os.environ.get("NANOVLLM_DEBUG_LOGITS"): + qn, kn, vn = query_conv.norm().item(), key_conv.norm().item(), value_conv.norm().item() + has_nan = torch.isnan(query_conv).any() or torch.isnan(key_conv).any() or torch.isnan(value_conv).any() + if has_nan: + print(f"[Qwen3_5 GatedDeltaNet] layer{self.layer_idx} pre-chunk: q_norm={qn:.4f} k_norm={kn:.4f} v_norm={vn:.4f} has_nan=True") + + # Part 3: Gated Delta Rule (Recurrent Attention) + # Prefer Triton FLA ops when available; use NANOVLLM_FORCE_TORCH_GDN=1 to force torch (for debugging NaN). + # NANOVLLM_GDN_FLOAT32=1: run in float32 for numerical stability (skips Triton). + use_float32 = bool(os.environ.get("NANOVLLM_GDN_FLOAT32")) + debug_state = bool(os.environ.get("NANOVLLM_DEBUG_GDN_STATE")) + is_capturing = ( + torch.cuda.is_available() + and hidden_states.is_cuda + and torch.cuda.is_current_stream_capturing() + ) + if use_float32: + query_conv = query_conv.float() + key_conv = key_conv.float() + value_conv = value_conv.float() + g = g.float() + beta = beta.float() + use_recurrent = (seq_len == 1 and recurrent_state is not None) + # Force first decode-state initialization with canonical layout (N, H, V, K) + # so triton recurrent path can do in-place updates on a stable storage. + if seq_len == 1 and recurrent_state is None: + recurrent_state = torch.zeros( + 1, + self.num_v_heads // dist.get_world_size(), + self.head_v_dim, + self.head_k_dim, + dtype=value_conv.dtype, + device=value_conv.device, + ) + use_recurrent = True + + decode_step = None + if seq_len == 1: + seq_key = sequence_id if sequence_id is not None else -1 + decode_step = self._decode_step_counter.get(seq_key, 0) + 1 + self._decode_step_counter[seq_key] = decode_step + if debug_state and not is_capturing: + print( + "[GDN state] layer=%d seq_id=%s use_recurrent=%s recurrent_none=%s" + % (self.layer_idx, str(sequence_id), str(use_recurrent), str(recurrent_state is None)) + ) + if decode_step is not None and decode_step <= 3 and recurrent_state is not None: + before_decode_sum = float(recurrent_state.float().sum().item()) + print( + "[GDN state] decode step=%d layer=%d seq_id=%s before_sum=%.6f" + % (decode_step, self.layer_idx, str(sequence_id), before_decode_sum) + ) + # Canonical cache layout is (N, H, V, K). + # NOTE: for Qwen3.5-0.8B, head_k_dim == head_v_dim (both 128), so shape-based + # auto-detection of (K,V) vs (V,K) is ambiguous and can corrupt state. + # Keep layout fixed and only convert explicitly at torch fallback boundaries. + if debug_state and not is_capturing and recurrent_state is not None: + print( + "[GDN state] cache layout (N,H,V,K) shape=%s stride=%s contig=%s" + % ( + str(tuple(recurrent_state.shape)), + str(tuple(recurrent_state.stride())), + str(recurrent_state.is_contiguous()), + ) + ) + use_triton = ( + not os.environ.get("NANOVLLM_FORCE_TORCH_GDN") + and not use_float32 + and _HAS_FLA_TRITON + and hidden_states.dtype != torch.float32 + ) + core_attn_out = None + triton_succeeded = False + if use_triton: + try: + # FLA Triton expects (B, T, H, K) / (B, T, H, V) / (B, T, H); initial_state (N, H, V, K) + q_batch = query_conv.unsqueeze(0).contiguous() + k_batch = key_conv.unsqueeze(0).contiguous() + v_batch = value_conv.unsqueeze(0).contiguous() + g_batch = g.unsqueeze(0).contiguous() + beta_batch = beta.unsqueeze(0).contiguous() + initial_state_batch = None + if recurrent_state is not None: + # Cache layout is canonicalized to (N, H, V, K). + initial_state_batch = recurrent_state + if debug_state and not is_capturing: + before_sum = float(initial_state_batch.float().sum().item()) + print( + "[GDN state] triton init shape=%s stride=%s contig=%s sum=%.6f" + % ( + str(tuple(initial_state_batch.shape)), + str(tuple(initial_state_batch.stride())), + str(initial_state_batch.is_contiguous()), + before_sum, + ) + ) + + # HF-style: one sequence per call, no cu_seqlens + cu_seqlens = None + + if use_recurrent: + core_attn_out_batch, final_state_batch = fla_fused_recurrent_gated_delta_rule( + q=q_batch, + k=k_batch, + v=v_batch, + g=g_batch, + beta=beta_batch, + initial_state=initial_state_batch, + inplace_final_state=True, + cu_seqlens=cu_seqlens, + use_qk_l2norm_in_kernel=True, + ) + else: + core_attn_out_batch, final_state_batch = fla_chunk_gated_delta_rule( + q=q_batch, + k=k_batch, + v=v_batch, + g=g_batch, + beta=beta_batch, + initial_state=initial_state_batch, + output_final_state=True, + cu_seqlens=cu_seqlens, + use_qk_l2norm_in_kernel=True, + ) + core_attn_out = core_attn_out_batch.squeeze(0) + # NaN fallback: if Triton produced NaN, retry with torch (do not use corrupted state) + if torch.isnan(core_attn_out).any(): + if not getattr(Qwen3_5GatedDeltaNet, "_nan_fallback_logged", False): + print("[Qwen3_5 GatedDeltaNet] Triton produced NaN, falling back to torch (set NANOVLLM_FORCE_TORCH_GDN=1 to skip Triton)") + Qwen3_5GatedDeltaNet._nan_fallback_logged = True + core_attn_out = None + else: + triton_succeeded = True + # For recurrent mode with inplace_final_state=True, state is + # already written back into `initial_state_batch` storage. + if (not use_recurrent) and final_state_batch is not None: + # Chunk path returns explicit final state in (N,H,V,K). + recurrent_state = final_state_batch.contiguous() + if debug_state and not is_capturing and recurrent_state is not None: + after_sum = float(recurrent_state.float().sum().item()) + print( + "[GDN state] triton updated shape=%s stride=%s contig=%s sum=%.6f" + % ( + str(tuple(recurrent_state.shape)), + str(tuple(recurrent_state.stride())), + str(recurrent_state.is_contiguous()), + after_sum, + ) + ) + except Exception: + use_triton = False + + if core_attn_out is None: + if use_recurrent: + # Torch fallback expects initial_state layout (N, H, K, V). + recurrent_state_torch = ( + recurrent_state.permute(0, 1, 3, 2).contiguous() + if recurrent_state is not None + else None + ) + core_attn_out, recurrent_state = torch_recurrent_gated_delta_rule( + query_conv, + key_conv, + value_conv, + g, + beta, + initial_state=recurrent_state_torch, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + ) + # Convert back to canonical cache layout (N, H, V, K). + if recurrent_state is not None: + recurrent_state = recurrent_state.permute(0, 1, 3, 2).contiguous() + else: + recurrent_state_torch = ( + recurrent_state.permute(0, 1, 3, 2).contiguous() + if recurrent_state is not None + else None + ) + core_attn_out, recurrent_state = torch_chunk_gated_delta_rule( + query_conv, + key_conv, + value_conv, + g, + beta, + chunk_size=64, + initial_state=recurrent_state_torch, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + ) + if recurrent_state is not None: + recurrent_state = recurrent_state.permute(0, 1, 3, 2).contiguous() + if use_float32 and core_attn_out is not None: + core_attn_out = core_attn_out.to(hidden_states.dtype) + if recurrent_state is not None: + recurrent_state = recurrent_state.to(hidden_states.dtype) + if ( + debug_state + and not is_capturing + and decode_step is not None + and decode_step <= 3 + and recurrent_state is not None + ): + after_decode_sum = float(recurrent_state.float().sum().item()) + print( + "[GDN state] decode step=%d layer=%d seq_id=%s after_sum=%.6f" + % (decode_step, self.layer_idx, str(sequence_id), after_decode_sum) + ) + if _gdn_debug_this_layer() and core_attn_out is not None: + print(f"[NV GDN L{self.layer_idx}] core_attn_out: shape={list(core_attn_out.shape)}, mean={core_attn_out.mean():.6f}, std={core_attn_out.std():.6f}, min={core_attn_out.min():.6f}, max={core_attn_out.max():.6f}") + if os.environ.get("NANOVLLM_DEBUG_LOGITS") and core_attn_out is not None and torch.isnan(core_attn_out).any(): + backend = "triton" if triton_succeeded else "torch" + print(f"[Qwen3_5 GatedDeltaNet] layer{self.layer_idx} post-chunk ({backend}): has_nan=True") + # Scheme C: stability check after chunk/recurrent + if os.environ.get("NANOVLLM_DEBUG_STABILITY") and core_attn_out is not None: + c_max = core_attn_out.abs().max().item() + has_nan = torch.isnan(core_attn_out).any().item() + if has_nan or c_max > 1e4: + backend = "triton" if triton_succeeded else "torch" + print(f"[GDN layer{self.layer_idx}] post-chunk ({backend}): max_abs={c_max:.4f} has_nan={has_nan}") + # Log backend once (Triton vs torch) and dtype + if not getattr(Qwen3_5GatedDeltaNet, "_gdn_backend_logged", False): + backend = "triton" if triton_succeeded else "torch" + # print(f"[Qwen3_5 GatedDeltaNet] backend={backend} dtype={hidden_states.dtype}") + Qwen3_5GatedDeltaNet._gdn_backend_logged = True + + # Persist states for this sequence (HF-style per-seq isolation) + self._set_states(conv_state, recurrent_state, sequence_id) + + # Part 4: Output Projection with Gate + # CRITICAL: Gate application must match vLLM implementation + # core_attn_out: (seq_len, num_v_heads, head_v_dim) + # z: (seq_len, num_v_heads, head_v_dim) + z_shape_og = z.shape + core_attn_out_flat = core_attn_out.reshape(-1, core_attn_out.shape[-1]) + z_flat = z.reshape(-1, z.shape[-1]) + + # Apply normalization with gate: norm(core_attn_out) * silu(gate) + # The gate (z) is applied via SiLU activation in Qwen3NextRMSNormGated + core_attn_out_flat = self.norm(core_attn_out_flat, z_flat) + core_attn_out = core_attn_out_flat.reshape(z_shape_og) + if _gdn_debug_this_layer(): + print(f"[NV GDN L{self.layer_idx}] after norm: shape={list(core_attn_out.shape)}, mean={core_attn_out.mean():.6f}, std={core_attn_out.std():.6f}") + if os.environ.get("NANOVLLM_DEBUG_LOGITS") and torch.isnan(core_attn_out).any(): + print(f"[Qwen3_5 GatedDeltaNet] layer{self.layer_idx} post-norm: has_nan=True") + + # Flatten for output projection + core_attn_out = core_attn_out.reshape(seq_len, -1) + + # Output projection + # RowParallelLinear may return tuple in some implementations + output = self.out_proj(core_attn_out) + if isinstance(output, tuple): + output = output[0] + if _gdn_debug_this_layer(): + print(f"[NV GDN L{self.layer_idx}] output: shape={list(output.shape)}, mean={output.mean():.6f}, std={output.std():.6f}, min={output.min():.6f}, max={output.max():.6f}") + if os.environ.get("NANOVLLM_DEBUG_LOGITS") and torch.isnan(output).any(): + print(f"[Qwen3_5 GatedDeltaNet] layer{self.layer_idx} post-out_proj: has_nan=True") + + return output + + +class Qwen3_5TextDecoderLayerMerged(nn.Module): + + def __init__( + self, + config, + ) -> None: + super().__init__() + rope_scaling = getattr(config, "rope_scaling", None) + if isinstance(rope_scaling, dict): + rope_scaling = None + + self.self_attn = Qwen3_5TextAttentionMerged( + hidden_size=config.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=config.num_key_value_heads, + max_position=config.max_position_embeddings, + rms_norm_eps=config.rms_norm_eps, + qkv_bias=getattr(config, "attention_bias", True), + head_dim=getattr(config, "head_dim", None), + rope_theta=getattr(config, "rope_theta", 1000000), + rope_scaling=rope_scaling, + ) + self.mlp = Qwen3_5TextMLPMerged( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + ) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + +class Qwen3_5TextDecoderLayer(nn.Module): + """Qwen3_5 Text Decoder Layer with layer_types support""" + + def __init__( + self, + config, + layer_idx: int = 0, + ) -> None: + super().__init__() + rope_scaling = getattr(config, "rope_scaling", None) + if isinstance(rope_scaling, dict): + rope_scaling = None + + self.layer_idx = layer_idx + # Support layer_types (full_attention or linear_attention) + # Must match transformers default when layer_types is None: (i+1) % 4 -> linear else full + layer_types = getattr(config, "layer_types", None) + if layer_types is not None: + self.layer_type = layer_types[layer_idx] if layer_idx < len(layer_types) else "full_attention" + else: + interval = getattr(config, "full_attention_interval", 4) + self.layer_type = ( + "linear_attention" if (layer_idx + 1) % interval else "full_attention" + ) + + if self.layer_type == "full_attention": + self.self_attn = Qwen3_5TextAttention( + hidden_size=config.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=config.num_key_value_heads, + max_position=config.max_position_embeddings, + rms_norm_eps=config.rms_norm_eps, + qkv_bias=getattr(config, "attention_bias", False), + head_dim=getattr(config, "head_dim", None), + rope_theta=getattr(config, "rope_theta", 1000000), + rope_scaling=rope_scaling, + config=config, + ) + self.linear_attn = None + elif self.layer_type == "linear_attention": + self.self_attn = None + self.linear_attn = Qwen3_5GatedDeltaNet( + config=config, + layer_idx=layer_idx, + ) + else: + raise NotImplementedError(f"Layer type {self.layer_type} not implemented") + + self.mlp = Qwen3_5TextMLPMerged( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + ) + self.input_layernorm = Qwen3NextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = Qwen3NextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + sequence_ids: list[int] | None = None, + sequence_lengths: list[int] | None = None, + use_graph: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor]: + # Optional debug: check residual propagation at layer 1. + if os.environ.get("NANOVLLM_DEBUG_RESIDUAL") and self.layer_idx == 1 and residual is not None: + print( + f"[DEBUG Layer 1] hidden_states.norm()={hidden_states.norm().item():.4f} " + f"residual.norm()={residual.norm().item():.4f}" + ) + if residual is None: + hidden_states, residual = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + # Scheme A: layer-wise diagnostics (NANOVLLM_DEBUG_LAYERS=1 prints every layer) + if os.environ.get("NANOVLLM_DEBUG_LAYERS"): + norm_val = hidden_states.norm().item() if hidden_states.numel() > 0 else float("nan") + has_nan = torch.isnan(hidden_states).any().item() + print(f"[Layer {self.layer_idx}] pre-attn norm={norm_val:.4f} has_nan={has_nan} type={self.layer_type}") + elif os.environ.get("NANOVLLM_DEBUG_LOGITS") and torch.isnan(hidden_states).any(): + print(f"[Qwen3_5 DecoderLayer] layer{self.layer_idx} pre-attn (after input_layernorm): has_nan=True") + + if self.layer_type == "full_attention": + debug_fa = os.environ.get("NANOVLLM_DEBUG_FA") + try: + debug_fa_layer = int(os.environ.get("NANOVLLM_DEBUG_FA_LAYER", "3")) + except ValueError: + debug_fa_layer = 3 + if debug_fa and self.layer_idx == debug_fa_layer: + FA_DEBUG_SAVE["pre_attn_norm"] = hidden_states.detach().cpu().float() + attn_out = self.self_attn(positions, hidden_states) + if debug_fa and self.layer_idx == debug_fa_layer: + FA_DEBUG_SAVE["attn_output"] = attn_out.detach().cpu().float() + hidden_states = attn_out + elif self.layer_type == "linear_attention": + if use_graph: + # CUDA graph mode: process entire batch at once using persistent state buffers + hidden_states = self.linear_attn(hidden_states, positions, use_graph=True) + elif sequence_lengths is not None and len(sequence_lengths) > 1 and sequence_ids is not None: + # HF-style: per-sequence isolation - process each sequence separately to avoid state cross-contamination + outputs = [] + offset = 0 + pos_2d = positions.dim() == 2 + for seq_len, seq_id in zip(sequence_lengths, sequence_ids): + h = hidden_states[offset : offset + seq_len] + p = positions[:, offset : offset + seq_len] if pos_2d else positions[offset : offset + seq_len] + out = self.linear_attn(h, p, sequence_id=seq_id) + outputs.append(out) + offset += seq_len + hidden_states = torch.cat(outputs, dim=0) + else: + seq_id = sequence_ids[0] if (sequence_ids is not None and len(sequence_ids) == 1) else None + hidden_states = self.linear_attn(hidden_states, positions, sequence_id=seq_id) + else: + raise NotImplementedError(f"Layer type {self.layer_type} not implemented") + + hidden_states = hidden_states + residual + if os.environ.get("NANOVLLM_DEBUG_LAYERS"): + norm_val = hidden_states.norm().item() if hidden_states.numel() > 0 else float("nan") + has_nan = torch.isnan(hidden_states).any().item() + print(f"[Layer {self.layer_idx}] post-attn norm={norm_val:.4f} has_nan={has_nan}") + elif os.environ.get("NANOVLLM_DEBUG_LOGITS") and torch.isnan(hidden_states).any(): + print(f"[Qwen3_5 DecoderLayer] layer{self.layer_idx} post-attn ({self.layer_type}): has_nan=True") + + # HF: residual = hidden_states, then norm(hidden_states); add residual after mlp + residual = hidden_states + hidden_states, _ = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = hidden_states + residual + if os.environ.get("NANOVLLM_DEBUG_LAYERS"): + norm_val = hidden_states.norm().item() if hidden_states.numel() > 0 else float("nan") + has_nan = torch.isnan(hidden_states).any().item() + print(f"[Layer {self.layer_idx}] post-mlp norm={norm_val:.4f} has_nan={has_nan}") + elif os.environ.get("NANOVLLM_DEBUG_LOGITS") and torch.isnan(hidden_states).any(): + print(f"[Qwen3_5 DecoderLayer] layer{self.layer_idx} post-mlp: has_nan=True") + # Return full layer output as residual for next layer (align with vLLM/HF semantics) + return hidden_states, hidden_states + + +class Qwen3_5TextModelMerged(nn.Module): + + def __init__( + self, + config, + ) -> None: + super().__init__() + self.embed_tokens = VocabParallelEmbedding(config.vocab_size, config.hidden_size) + self.layers = nn.ModuleList( + [Qwen3_5TextDecoderLayerMerged(config) for _ in range(config.num_hidden_layers)] + ) + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + input_ids: torch.Tensor | None = None, + positions: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + vision_token_count: int | None = None, + visual_pos_mask: torch.Tensor | None = None, + deepstack_visual_embeds: list[torch.Tensor] | None = None, + ) -> torch.Tensor: + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = self.embed_tokens(input_ids) + # vLLM-style: keep activation dtype consistent with weights + hidden_states = hidden_states.to(self.embed_tokens.weight.dtype) + + residual = None + for layer in self.layers: + hidden_states, residual = layer(positions, hidden_states, residual) + # HF: final norm only normalizes hidden_states (no residual merge) + hidden_states = self.norm(hidden_states) + return hidden_states + + +class Qwen3_5TextForCausalLMMerged(nn.Module): + packed_modules_mapping = { + "q_proj": ("qkv_proj", "q"), + "k_proj": ("qkv_proj", "k"), + "v_proj": ("qkv_proj", "v"), + "gate_proj": ("gate_up_proj", 0), + "up_proj": ("gate_up_proj", 1), + } + + def __init__(self, config) -> None: + super().__init__() + self.model = Qwen3_5TextModelMerged(config) + self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size) + if config.tie_word_embeddings: + self.lm_head.weight.data = self.model.embed_tokens.weight.data + + def forward( + self, + input_ids: torch.Tensor | None = None, + positions: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + vision_token_count: int | None = None, + visual_pos_mask: torch.Tensor | None = None, + deepstack_visual_embeds: list[torch.Tensor] | None = None, + **_: dict, + ) -> torch.Tensor: + return self.model( + input_ids, + positions, + inputs_embeds=inputs_embeds, + vision_token_count=vision_token_count, + visual_pos_mask=visual_pos_mask, + deepstack_visual_embeds=deepstack_visual_embeds, + ) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = hidden_states.to(self.lm_head.weight.dtype) + return self.lm_head(hidden_states) + + +class Qwen3_5TextModel(nn.Module): + """Qwen3_5 Text Model""" + + def __init__( + self, + config, + ) -> None: + super().__init__() + self.embed_tokens = VocabParallelEmbedding(config.vocab_size, config.hidden_size) + self.layers = nn.ModuleList( + [Qwen3_5TextDecoderLayer(config, layer_idx=i) for i in range(config.num_hidden_layers)] + ) + self.norm = Qwen3NextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + input_ids: torch.Tensor | None = None, + positions: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + vision_token_count: int | None = None, + visual_pos_mask: torch.Tensor | None = None, + deepstack_visual_embeds: list[torch.Tensor] | None = None, + sequence_ids: list[int] | None = None, + sequence_lengths: list[int] | None = None, + use_graph: bool = False, + ) -> torch.Tensor: + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = self.embed_tokens(input_ids) + # vLLM-style: keep activation dtype consistent with weights + hidden_states = hidden_states.to(self.embed_tokens.weight.dtype) + + residual = None + for layer in self.layers: + hidden_states, residual = layer( + positions, hidden_states, residual, + sequence_ids=sequence_ids, + sequence_lengths=sequence_lengths, + use_graph=use_graph, + ) + hidden_states, _ = self.norm(hidden_states, residual) + return hidden_states + + +class Qwen3_5TextForCausalLM(nn.Module): + """Qwen3_5 Text For Causal LM""" + + packed_modules_mapping = { + "q_proj": ("q_proj", None), + "k_proj": ("k_proj", None), + "v_proj": ("v_proj", None), + "gate_proj": ("gate_up_proj", 0), + "up_proj": ("gate_up_proj", 1), + } + + def __init__(self, config) -> None: + super().__init__() + self.model = Qwen3_5TextModel(config) + self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size) + if config.tie_word_embeddings: + self.lm_head.weight.data = self.model.embed_tokens.weight.data + + def forward( + self, + input_ids: torch.Tensor | None = None, + positions: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + vision_token_count: int | None = None, + visual_pos_mask: torch.Tensor | None = None, + deepstack_visual_embeds: list[torch.Tensor] | None = None, + sequence_ids: list[int] | None = None, + sequence_lengths: list[int] | None = None, + use_graph: bool = False, + **_: dict, + ) -> torch.Tensor: + return self.model( + input_ids, + positions, + inputs_embeds=inputs_embeds, + vision_token_count=vision_token_count, + visual_pos_mask=visual_pos_mask, + deepstack_visual_embeds=deepstack_visual_embeds, + sequence_ids=sequence_ids, + sequence_lengths=sequence_lengths, + use_graph=use_graph, + ) + + def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = hidden_states.to(self.lm_head.weight.dtype) + return self.lm_head(hidden_states) + + +# --------------------------------------------------------------------------- +# Vision encoder (Qwen3.5: vision tower only; no DeepStack) +# --------------------------------------------------------------------------- + + +def gelu_pytorch_tanh(x: torch.Tensor) -> torch.Tensor: + inner = math.sqrt(2 / math.pi) * (x + 0.044715 * x * x * x) + return 0.5 * x * (1 + torch.tanh(inner)) + + +def rotate_half(x: torch.Tensor) -> torch.Tensor: + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary_pos_emb_vision( + q: torch.Tensor, + k: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + cos = cos.unsqueeze(0).to(q.dtype) + sin = sin.unsqueeze(0).to(q.dtype) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +class Qwen3_5VisionPatchEmbed(nn.Module): + def __init__(self, config) -> None: + super().__init__() + self.patch_size = config.patch_size + self.temporal_patch_size = getattr(config, "temporal_patch_size", 1) + self.in_channels = config.in_channels + self.embed_dim = config.hidden_size + + kernel_size = [ + self.temporal_patch_size, + self.patch_size, + self.patch_size, + ] + stride = kernel_size + self.proj = nn.Conv3d( + self.in_channels, + self.embed_dim, + kernel_size=kernel_size, + stride=stride, + bias=True, + ) + + def forward(self, inputs: torch.Tensor) -> torch.Tensor: + if inputs.dim() == 4: + inputs = inputs.unsqueeze(2) + hidden_states = self.proj(inputs) + hidden_states = hidden_states.flatten(2).transpose(1, 2) + return hidden_states + + +class Qwen3_5VisionRotaryEmbedding(nn.Module): + inv_freq: torch.Tensor + + def __init__(self, dim: int, theta: float = 10000.0) -> None: + super().__init__() + inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float) / dim)) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + def forward(self, seqlen: int) -> torch.Tensor: + seq = torch.arange(seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype) + freqs = torch.outer(seq, self.inv_freq) + return freqs + + +class Qwen3_5VisionMLP(nn.Module): + def __init__(self, config) -> None: + super().__init__() + self.linear_fc1 = nn.Linear(config.hidden_size, config.intermediate_size, bias=True) + self.linear_fc2 = nn.Linear(config.intermediate_size, config.hidden_size, bias=True) + if getattr(config, "hidden_act", "gelu") == "gelu_pytorch_tanh": + self.act_fn = gelu_pytorch_tanh + else: + self.act_fn = lambda x: F.gelu(x, approximate="tanh") + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.linear_fc1(hidden_states) + hidden_states = self.act_fn(hidden_states) + hidden_states = self.linear_fc2(hidden_states) + return hidden_states + + +class Qwen3_5VisionAttention(nn.Module): + def __init__(self, config) -> None: + super().__init__() + self.hidden_size = config.hidden_size + self.num_heads = config.num_heads + self.head_dim = self.hidden_size // self.num_heads + self.scale = self.head_dim**-0.5 + + self.qkv = nn.Linear(self.hidden_size, self.hidden_size * 3, bias=True) + self.proj = nn.Linear(self.hidden_size, self.hidden_size, bias=True) + + def forward( + self, + hidden_states: torch.Tensor, + seq_lengths: Sequence[int], + position_embeddings: Tuple[torch.Tensor, torch.Tensor], + ) -> torch.Tensor: + outputs = [] + offset = 0 + cos, sin = position_embeddings + for length in seq_lengths: + chunk = hidden_states[offset : offset + length] + cos_chunk = cos[offset : offset + length] + sin_chunk = sin[offset : offset + length] + + qkv = self.qkv(chunk) + q, k, v = qkv.chunk(3, dim=-1) + + q = q.reshape(length, self.num_heads, self.head_dim).transpose(0, 1) + k = k.reshape(length, self.num_heads, self.head_dim).transpose(0, 1) + v = v.reshape(length, self.num_heads, self.head_dim).transpose(0, 1) + + q, k = apply_rotary_pos_emb_vision(q, k, cos_chunk, sin_chunk) + if q.dtype != v.dtype: + q = q.to(v.dtype) + k = k.to(v.dtype) + + attn_scores = torch.matmul(q, k.transpose(-1, -2)) * self.scale + attn_weights = torch.softmax(attn_scores, dim=-1, dtype=torch.float32).to(v.dtype) + attn_output = torch.matmul(attn_weights, v) + + attn_output = attn_output.transpose(0, 1).reshape(length, self.hidden_size) + attn_output = self.proj(attn_output) + outputs.append(attn_output) + offset += length + + return torch.cat(outputs, dim=0) + + +class Qwen3_5VisionPatchMerger(nn.Module): + def __init__(self, config, use_postshuffle_norm: bool = False) -> None: + super().__init__() + self.hidden_size = config.hidden_size * (config.spatial_merge_size**2) + self.use_postshuffle_norm = use_postshuffle_norm + norm_dim = self.hidden_size if use_postshuffle_norm else config.hidden_size + self.norm = nn.LayerNorm(norm_dim, eps=1e-6) + self.linear_fc1 = nn.Linear(self.hidden_size, self.hidden_size, bias=True) + self.linear_fc2 = nn.Linear(self.hidden_size, config.out_hidden_size, bias=True) + self.act_fn = nn.GELU() + self.merge_size = config.spatial_merge_size + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + if self.use_postshuffle_norm: + hidden_states = hidden_states.reshape(-1, self.hidden_size) + hidden_states = self.norm(hidden_states) + hidden_states = hidden_states.reshape(-1, self.hidden_size) + hidden_states = self.linear_fc1(hidden_states) + hidden_states = self.act_fn(hidden_states) + hidden_states = self.linear_fc2(hidden_states) + return hidden_states + + +class Qwen3_5VisionBlock(nn.Module): + def __init__(self, config) -> None: + super().__init__() + self.norm1 = nn.LayerNorm(config.hidden_size, eps=1e-6) + self.attn = Qwen3_5VisionAttention(config) + self.norm2 = nn.LayerNorm(config.hidden_size, eps=1e-6) + self.mlp = Qwen3_5VisionMLP(config) + + def forward( + self, + hidden_states: torch.Tensor, + seq_lengths: Sequence[int], + position_embeddings: Tuple[torch.Tensor, torch.Tensor], + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.norm1(hidden_states) + hidden_states = residual + self.attn(hidden_states, seq_lengths, position_embeddings) + residual = hidden_states + hidden_states = self.norm2(hidden_states) + hidden_states = residual + self.mlp(hidden_states) + return hidden_states + + +class Qwen3_5VisionModel(nn.Module): + def __init__(self, config) -> None: + super().__init__() + self.config = config + self.spatial_merge_size = config.spatial_merge_size + self.patch_embed = Qwen3_5VisionPatchEmbed(config) + self.hidden_size = config.hidden_size + 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_5VisionRotaryEmbedding(head_dim // 2) + + self.blocks = nn.ModuleList([Qwen3_5VisionBlock(config) for _ in range(config.depth)]) + self.merger = Qwen3_5VisionPatchMerger(config=config, use_postshuffle_norm=False) + # Qwen3.5 has no DeepStack + + def rot_pos_emb(self, grid_thw: torch.Tensor) -> torch.Tensor: + merge_size = self.spatial_merge_size + max_hw = int(grid_thw[:, 1:].max().item()) + freq_table = self.rotary_pos_emb(max_hw) + device = freq_table.device + + total_tokens = int(torch.prod(grid_thw, dim=1).sum().item()) + pos_ids = torch.empty((total_tokens, 2), dtype=torch.long, device=device) + + offset = 0 + for num_frames, height, width in grid_thw: + merged_h = height // merge_size + merged_w = 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.view(1, 1, -1, 1) + ) + col_idx = ( + block_cols[None, :, None, None] * merge_size + + intra_col.view(1, 1, 1, -1) + ) + + 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: torch.Tensor) -> torch.Tensor: + grid_ts, grid_hs, grid_ws = grid_thw[:, 0], grid_thw[:, 1], grid_thw[:, 2] + device = grid_thw.device + + idx_list = [[] for _ in range(4)] + weight_list = [[] for _ in range(4)] + + for t, h, w in zip(grid_ts, grid_hs, grid_ws): + h_idxs = torch.linspace(0, self.num_grid_per_side - 1, h, device=device) + w_idxs = torch.linspace(0, self.num_grid_per_side - 1, w, device=device) + + h_floor = h_idxs.long() + w_floor = w_idxs.long() + h_ceil = (h_floor + 1).clip(max=self.num_grid_per_side - 1) + w_ceil = (w_floor + 1).clip(max=self.num_grid_per_side - 1) + + dh = h_idxs - h_floor + dw = w_idxs - w_floor + + base_h = h_floor * self.num_grid_per_side + base_h_ceil = h_ceil * self.num_grid_per_side + + indices = [ + (base_h[None].T + w_floor[None]).flatten(), + (base_h[None].T + w_ceil[None]).flatten(), + (base_h_ceil[None].T + w_floor[None]).flatten(), + (base_h_ceil[None].T + w_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) * weight_tensor[:, :, None] + patch_pos_embeds = pos_embeds.sum(dim=0) + + patch_pos_embeds = patch_pos_embeds.split([h * w for h, w in zip(grid_hs, grid_ws)]) + + patch_pos_embeds_permute = [] + merge_size = self.config.spatial_merge_size + for pos_embed, t, h, w in zip(patch_pos_embeds, grid_ts, grid_hs, grid_ws): + pos_embed = pos_embed.repeat(int(t.item()), 1) + pos_embed = ( + pos_embed.view(int(t.item()), 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) + patch_pos_embeds = torch.cat(patch_pos_embeds_permute) + return patch_pos_embeds + + def forward(self, pixel_values: torch.Tensor, grid_thw: torch.Tensor) -> Tuple[torch.Tensor, List[torch.Tensor]]: + pixel_values = pixel_values.to(self.pos_embed.weight.dtype) + seq_tokens = self.patch_embed(pixel_values) + hidden_states = seq_tokens.reshape(-1, self.hidden_size) + + pos_embeds = self.fast_pos_embed_interpolate(grid_thw) + hidden_states = hidden_states + pos_embeds + + rotary_pos_emb = self.rot_pos_emb(grid_thw) + rotary_pos_emb = rotary_pos_emb.reshape(hidden_states.size(0), -1) + emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1) + position_embeddings = (emb.cos(), emb.sin()) + + seq_lengths = (grid_thw[:, 0] * grid_thw[:, 1] * grid_thw[:, 2]).tolist() + + for block in self.blocks: + hidden_states = block(hidden_states, seq_lengths, position_embeddings) + + hidden_states = self.merger(hidden_states) + return hidden_states + + +class Qwen3VisionEncoder(nn.Module): + def __init__(self, vision_config) -> None: + super().__init__() + self.config = vision_config + self.vision = Qwen3_5VisionModel(vision_config) + + def _linear_patch_embed(self, patch_tokens: torch.Tensor) -> torch.Tensor: + proj = self.vision.patch_embed.proj + weight = proj.weight.view(proj.out_channels, -1) + bias = proj.bias + return torch.nn.functional.linear(patch_tokens, weight, bias) + + def _run_vision_from_tokens( + self, + token_list: list[torch.Tensor], + grid_thw: torch.Tensor, + ) -> list[torch.Tensor]: + proj = self.vision.patch_embed.proj + device = proj.weight.device + dtype = proj.weight.dtype + + tokens = torch.cat([t.to(device=device, dtype=dtype) for t in token_list], dim=0) + grids = grid_thw.to(device=device, dtype=torch.int32) + + hidden_states = tokens + + pos_embeds = self.vision.fast_pos_embed_interpolate(grids).to(hidden_states.dtype) + hidden_states = hidden_states + pos_embeds + + rotary_pos_emb = self.vision.rot_pos_emb(grids).to(hidden_states.dtype) + rotary_pos_emb = rotary_pos_emb.reshape(hidden_states.size(0), -1) + emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1) + position_embeddings = (emb.cos(), emb.sin()) + + seq_lengths = (grids[:, 0] * grids[:, 1] * grids[:, 2]).tolist() + + for block in self.vision.blocks: + hidden_states = block(hidden_states, seq_lengths, position_embeddings) + + hidden_states = self.vision.merger(hidden_states) + + split_sizes = ( + grids.prod(-1) // (self.config.spatial_merge_size**2) + ).tolist() + + image_chunks = list(torch.split(hidden_states, split_sizes)) + return image_chunks + + def _normalize_pixel_inputs( + self, + pixel_values: torch.Tensor, + ) -> Tuple[torch.Tensor, int, int, int, int]: + in_channels = getattr(self.config, "in_channels", 3) + num_dims = pixel_values.dim() + + channel_axis = None + for axis in range(1, num_dims - 2): + if pixel_values.shape[axis] == in_channels: + channel_axis = axis + break + if channel_axis is None: + channel_axis = 1 + + permute_order = [0, channel_axis] + temporal_axes = [ + axis for axis in range(1, num_dims - 2) if axis != channel_axis + ] + permute_order.extend(temporal_axes) + permute_order.extend([num_dims - 2, num_dims - 1]) + + pixel_values = pixel_values.permute(*permute_order).contiguous() + + batch = pixel_values.shape[0] + channels = pixel_values.shape[1] + height = pixel_values.shape[-2] + width = pixel_values.shape[-1] + + temporal = int(math.prod(pixel_values.shape[2:-2])) + pixel_values = pixel_values.reshape( + batch, + channels, + temporal, + height, + width, + ) + + return pixel_values, batch, temporal, height, width + + def forward( + self, + pixel_values: torch.Tensor, + image_grid_thw: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if pixel_values.dim() <= 3: + if image_grid_thw is None: + raise ValueError("image_grid_thw is required for flattened inputs") + grids = image_grid_thw.to(pixel_values.device).to(torch.int64) + tokens_per_image = grids.prod(-1).tolist() + if pixel_values.dim() == 3: + batch, tokens, feature = pixel_values.shape + flat = pixel_values.reshape(batch * tokens, feature) + else: + flat = pixel_values + + splits = torch.split(flat, tokens_per_image, dim=0) + token_list = [self._linear_patch_embed(chunk) for chunk in splits] + + image_chunks = self._run_vision_from_tokens(token_list, grids) + return torch.stack(list(image_chunks), dim=0) + + pixel_values, batch, temporal, height, width = self._normalize_pixel_inputs( + pixel_values + ) + + if image_grid_thw is None: + grid = torch.tensor( + [ + [ + temporal, + height // self.config.patch_size, + width // self.config.patch_size, + ] + ] + * batch, + device=pixel_values.device, + dtype=torch.int32, + ) + image_grid_thw = grid + else: + if image_grid_thw.dim() == 1: + image_grid_thw = image_grid_thw.unsqueeze(0) + image_grid_thw = image_grid_thw.to(device=pixel_values.device, dtype=torch.int32) + + image_embeds = self.vision(pixel_values, image_grid_thw) + split_sizes = ( + image_grid_thw.prod(-1) // (self.config.spatial_merge_size**2) + ).tolist() + + image_chunks = torch.split(image_embeds, split_sizes) + image_tokens = torch.stack(list(image_chunks), dim=0) + return image_tokens + + +def create_vision_model(config, **kwargs) -> nn.Module: + del kwargs + return Qwen3VisionEncoder(config) + + +def get_vision_model(config, **kwargs) -> nn.Module: + return create_vision_model(config, **kwargs) + + +# --------------------------------------------------------------------------- +# Multimodal wrapper +# --------------------------------------------------------------------------- + + +class Qwen3_5ForConditionalGeneration(nn.Module): + """Qwen3.5 conditional generation model (language + vision).""" + + def __init__(self, config): + super().__init__() + self.config = config + self.text_config = getattr(config, "text_config", config) + self.vision_config = getattr(config, "vision_config", None) + + if self.vision_config is None: + raise ValueError("vision_config is missing; cannot build a multimodal model") + + self.visual = create_vision_model(self.vision_config) + self.language_model = Qwen3_5TextForCausalLM(self.text_config) + + print("[Qwen3_5ForConditionalGeneration] Initialization complete") + print(f" - Vision encoder: {type(self.visual).__name__}") + print(f" - Language model: {type(self.language_model).__name__}") + + self.packed_modules_mapping = { + "mlp.gate_proj": ("mlp.gate_up_proj", 0), + "mlp.up_proj": ("mlp.gate_up_proj", 1), + "q_proj": ("q_proj", None), + "k_proj": ("k_proj", None), + "v_proj": ("v_proj", None), + } + + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.language_model.model.embed_tokens(input_ids) + + def get_vision_position_ids( + self, + start_position: int, + grid_thw: List[int] | torch.Tensor, + temp_merge_size: int = 1, + spatial_merge_size: int = 1, + time_interval: int = 1, + device: str | torch.device | None = None, + ) -> torch.Tensor: + """Compute 3D positional indices for vision tokens (PR 43972).""" + if isinstance(grid_thw, torch.Tensor): + g0, g1, g2 = grid_thw[0].item(), grid_thw[1].item(), grid_thw[2].item() + else: + g0, g1, g2 = grid_thw[0], grid_thw[1], grid_thw[2] + llm_grid_t = g0 // temp_merge_size + llm_grid_h = g1 // spatial_merge_size + llm_grid_w = g2 // spatial_merge_size + image_seq_length = llm_grid_h * llm_grid_w * llm_grid_t + if device is None: + device = next(self.parameters()).device + position_width = torch.arange( + start_position, start_position + llm_grid_w, device=device, dtype=torch.long + ).repeat(llm_grid_h * llm_grid_t) + position_height = torch.arange( + start_position, start_position + llm_grid_h, device=device, dtype=torch.long + ).repeat_interleave(llm_grid_w * llm_grid_t) + position_temporal = torch.full( + (image_seq_length,), start_position, device=device, dtype=torch.long + ) + position_temporal = position_temporal * time_interval + return torch.stack([position_temporal, position_height, position_width], dim=0) + + def _build_3d_position_ids_for_sequence( + self, + seq_len: int, + sorted_slices: list[dict], + image_grid_thw: torch.Tensor, + device: torch.device, + ) -> torch.Tensor: + """Build (3, seq_len) position ids for one sequence from text/vision segments.""" + spatial_merge_size = getattr(self.vision_config, "spatial_merge_size", 1) + chunks: list[torch.Tensor] = [] + current_pos = 0 + for seg_start, seg_end, is_vision, placeholder_idx in self._segment_ranges( + seq_len, sorted_slices + ): + seg_len = seg_end - seg_start + if seg_len <= 0: + continue + if is_vision and image_grid_thw is not None and placeholder_idx >= 0 and placeholder_idx < image_grid_thw.size(0): + grid = image_grid_thw[placeholder_idx] + pos_3d = self.get_vision_position_ids( + current_pos, grid, temp_merge_size=1, + spatial_merge_size=spatial_merge_size, device=device, + ) + if pos_3d.size(1) == seg_len: + chunks.append(pos_3d) + else: + ar = torch.arange(seg_len, device=device, dtype=torch.long) + current_pos + chunks.append(ar.view(1, -1).expand(3, -1)) + else: + ar = torch.arange(seg_len, device=device, dtype=torch.long) + current_pos + chunks.append(ar.view(1, -1).expand(3, -1)) + current_pos += seg_len + if not chunks: + ar = torch.arange(seq_len, device=device, dtype=torch.long) + return ar.view(1, -1).expand(3, -1) + return torch.cat(chunks, dim=1) + + @staticmethod + def _segment_ranges( + seq_len: int, sorted_slices: list[dict] + ) -> list[tuple[int, int, bool, int]]: + """Yield (start, end, is_vision, placeholder_idx) in order.""" + out: list[tuple[int, int, bool, int]] = [] + current = 0 + for s in sorted_slices: + off = s["target_offset"] + length = s["length"] + pid = s.get("placeholder_idx", 0) + if current < off: + out.append((current, off, False, -1)) + out.append((off, off + length, True, pid)) + current = off + length + if current < seq_len: + out.append((current, seq_len, False, -1)) + return out + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor | None, + inputs_embeds: torch.Tensor | None = None, + pixel_values: torch.Tensor | None = None, + image_grid_thw: torch.Tensor | None = None, + sequence_lengths: list[int] | None = None, + sequence_ids: list[int] | None = None, + vision_slices_per_seq: list[list[dict]] | None = None, + image_grid_thw_per_seq: list[torch.Tensor | None] | None = None, + use_graph: bool = False, + ) -> torch.Tensor: + if inputs_embeds is None: + if input_ids is None: + raise ValueError("input_ids and inputs_embeds cannot be None simultaneously") + inputs_embeds = self.get_input_embeddings(input_ids) + + total_tokens = inputs_embeds.size(0) + inputs_embeds = inputs_embeds.clone() + + visual_pos_mask = torch.zeros( + total_tokens, dtype=torch.bool, device=inputs_embeds.device + ) + vision_token_count = 0 + + if vision_slices_per_seq: + if sequence_lengths is None: + raise ValueError("sequence_lengths must be provided to align visual features") + if len(sequence_lengths) != len(vision_slices_per_seq): + raise ValueError("sequence_lengths and vision_slices_per_seq have different lengths") + if sum(sequence_lengths) != total_tokens: + raise ValueError("sum of sequence_lengths does not match total input tokens") + + offsets = [0] + for length in sequence_lengths: + offsets.append(offsets[-1] + length) + + for seq_idx, (start, end) in enumerate(zip(offsets[:-1], offsets[1:])): + seq_slices = vision_slices_per_seq[seq_idx] + if not seq_slices: + continue + + for slice_info in seq_slices: + token_slice = slice_info["tokens"].to( + device=inputs_embeds.device, + dtype=inputs_embeds.dtype, + ) + length = slice_info["length"] + target_offset = slice_info["target_offset"] + + if token_slice.size(0) != length: + raise ValueError("Visual token slice length does not match the declared length") + + target_start = start + target_offset + target_end = target_start + length + if target_end > end: + raise ValueError("Visual token target range is out of bounds") + + inputs_embeds[target_start:target_end] = token_slice + visual_pos_mask[target_start:target_end] = True + vision_token_count += length + + elif pixel_values is not None: + # Fallback path: process raw images when slices are not provided (legacy compatibility) + if input_ids is None: + raise ValueError("input_ids are required to locate visual placeholders") + if sequence_lengths is None: + raise ValueError("sequence_lengths are required to align visual features") + if sum(sequence_lengths) != total_tokens: + raise ValueError("sum of sequence_lengths does not match total input tokens") + + image_tokens = self.visual(pixel_values, image_grid_thw) + if image_tokens is None or image_tokens.numel() == 0: + raise ValueError("The vision encoder did not return valid image features") + + offsets = [0] + for length in sequence_lengths: + offsets.append(offsets[-1] + length) + + total_replaced = 0 + image_chunks = [image_tokens[i] for i in range(image_tokens.size(0))] + image_iter = iter(image_chunks) + + for start, end in zip(offsets[:-1], offsets[1:]): + seq_length = end - start + if seq_length <= 0: + continue + + try: + token_slice = next(image_iter) + except StopIteration: + break + + token_slice = token_slice.to(inputs_embeds.device, inputs_embeds.dtype) + slice_len = token_slice.size(0) + if slice_len > seq_length: + raise ValueError("Visual tokens exceed the available sequence length") + + inputs_embeds[start : start + slice_len] = token_slice + visual_pos_mask[start : start + slice_len] = True + total_replaced += slice_len + + vision_token_count = total_replaced + + if vision_token_count == 0: + visual_pos_mask = None + + # Build 3D position_ids (PR 43972) when we have vision segments and grid per sequence + if ( + vision_slices_per_seq is not None + and image_grid_thw_per_seq is not None + and len(image_grid_thw_per_seq) == len(sequence_lengths or []) + and sum(sequence_lengths or []) == total_tokens + ): + device = inputs_embeds.device + pos_chunks = [] + for seq_idx, seq_len in enumerate(sequence_lengths or []): + slices = vision_slices_per_seq[seq_idx] + sorted_slices = sorted(slices, key=lambda s: s["target_offset"]) + grid = image_grid_thw_per_seq[seq_idx] + pos_3d = self._build_3d_position_ids_for_sequence( + seq_len, sorted_slices, grid, device + ) + pos_chunks.append(pos_3d) + if pos_chunks: + positions = torch.cat(pos_chunks, dim=1) + + # For Qwen3.5, even pure text input needs 3D position_ids (T, H, W) for MRoPE + if positions is None: + positions = torch.arange(total_tokens, device=inputs_embeds.device) + positions = positions.view(1, 1, -1).expand(3, 1, -1).reshape(3, total_tokens) + elif positions.ndim == 1: + if total_tokens != positions.size(0): + positions = torch.arange(total_tokens, device=inputs_embeds.device) + positions = positions.view(1, -1).expand(3, -1) + + if visual_pos_mask is not None and vision_token_count: + visual_pos_mask = visual_pos_mask.to(inputs_embeds.device) + else: + visual_pos_mask = None + + # HF-style: per-sequence GDN state isolation + # Use actual sequence_ids from the scheduler (seq.seq_id) so that + # GDN state_cache keys remain stable even when batch composition changes. + seq_ids = sequence_ids if sequence_ids is not None else (list(range(len(sequence_lengths))) if sequence_lengths else None) + seq_lens = sequence_lengths + + hidden_states = self.language_model( + input_ids=None, + positions=positions, + inputs_embeds=inputs_embeds, + vision_token_count=vision_token_count, + visual_pos_mask=visual_pos_mask, + sequence_ids=seq_ids, + sequence_lengths=seq_lens, + use_graph=use_graph, + ) + + return hidden_states + + def compute_logits(self, hidden_states): + """Compute logits (delegate to language model)""" + return self.language_model.compute_logits(hidden_states) + + +def _qwen_multimodal_name_mapping(weight_name: str) -> str | None: + """Weight name mapping for Qwen3.5 multimodal checkpoints.""" + # lm_head.* (top-level, some checkpoints save as "lm_head" without "model." prefix) + if weight_name.startswith("lm_head."): + return "language_model." + weight_name + # model.lm_head.* (top-level lm_head under Qwen3_5ForConditionalGeneration) + if weight_name.startswith("model.lm_head."): + return "language_model." + weight_name[len("model.") :] + if weight_name.startswith("model.language_model."): + sub_name = weight_name[len("model.language_model.") :] + text_model_prefixes = ( + "model.", + "embed_tokens.", + "layers.", + "norm.", + "rotary_emb.", + ) + if sub_name.startswith(text_model_prefixes): + if sub_name.startswith("model."): + sub_name = sub_name[len("model.") :] + sub_name = "language_model.model." + sub_name + elif sub_name.startswith("lm_head."): + sub_name = "language_model.lm_head." + sub_name[len("lm_head.") :] + else: + sub_name = "language_model." + sub_name + return sub_name + if weight_name.startswith("model.visual."): + sub_name = weight_name[len("model.visual.") :] + return "visual.vision." + sub_name + return None + + +def load_qwen3_5_model(model_path, config): + """ + Load Qwen3.5 (or Qwen3.5-MoE) multimodal model. + + Args: + model_path: Model path + config: Configuration object + + Returns: + model: Qwen3_5ForConditionalGeneration instance + """ + hf_config = config.hf_config + model = Qwen3_5ForConditionalGeneration(hf_config) + from nanovllm.utils.loader import load_model + print("[load_qwen3_5_model] Loading Qwen3.5 multimodal weights...") + load_model(model, model_path, name_mapping=_qwen_multimodal_name_mapping) + print("[load_qwen3_5_model] Model ready.") + return model diff --git a/nanovllm/models/qwen3_next.py b/nanovllm/models/qwen3_next.py new file mode 100755 index 000000000..84471c516 --- /dev/null +++ b/nanovllm/models/qwen3_next.py @@ -0,0 +1,1142 @@ +import torch +from torch import nn +import torch.distributed as dist +import torch.nn.functional as F +from transformers import Qwen3NextConfig +from typing import Callable + +from nanovllm.layers.activation import SiluAndMul +from nanovllm.layers.attention import Attention +from nanovllm.layers.layernorm import RMSNorm +from nanovllm.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from nanovllm.layers.rotary_embedding import get_rope +from nanovllm.layers.embed_head import VocabParallelEmbedding, ParallelLMHead +from nanovllm.utils.loader import sharded_weight_loader + +# Import optimized GatedDeltaNet kernels +print("[Qwen3Next] 正在导入优化的 GatedDeltaNet kernels...") +try: + from nanovllm.layers.ops import fused_recurrent_gated_delta_rule, chunk_gated_delta_rule + HAS_FUSED_RECURRENT = fused_recurrent_gated_delta_rule is not None + HAS_CHUNK = chunk_gated_delta_rule is not None + + if fused_recurrent_gated_delta_rule is not None: + print("[Qwen3Next] ✓ 成功导入 fused_recurrent_gated_delta_rule (Triton kernel)") + else: + print("[Qwen3Next] ✗ fused_recurrent_gated_delta_rule 导入失败，将使用 Python fallback") + + if chunk_gated_delta_rule is not None: + print("[Qwen3Next] ✓ 成功导入 chunk_gated_delta_rule (Triton kernel)") + else: + print("[Qwen3Next] ✗ chunk_gated_delta_rule 导入失败，将使用 Python fallback") +except ImportError as e: + HAS_FUSED_RECURRENT = False + HAS_CHUNK = False + fused_recurrent_gated_delta_rule = None + chunk_gated_delta_rule = None + print("[Qwen3Next] ✗ 导入 GatedDeltaNet kernels 失败: {}".format(e)) + print("[Qwen3Next] 将使用 Python fallback 实现") + + +def mamba_v2_sharded_weight_loader( + shard_spec: list[tuple[int, int, bool]], + tp_size: int, + tp_rank: int, +) -> Callable: + """ + Create a weight loader for mamba v2 / GatedDeltaNet conv1d weights. + This ensures that the projections are correctly sharded so that they can be + split into query, key, value. It also ensures that all the groups corresponding + to a head shard is placed together with it. + + Args: + shard_spec: List of (full_dim, extra, duplicate_groups) tuples + - full_dim: The model dimension (before TP) + - extra: Expected overall increase of dimensions due to replication + - duplicate_groups: Whether groups are duplicated across TP ranks + tp_size: Tensor parallel size + tp_rank: Current tensor parallel rank + + Returns: + A weight loader function that takes (param, loaded_weight) and copies data + """ + def loader(param: nn.Parameter, loaded_weight: torch.Tensor) -> None: + # Track boundary of (sharded) param, and loaded_weight, respectively + boundary, loaded_boundary = 0, 0 + + # Iterate over the shard specs + for full_dim, extra, duplicate_groups in shard_spec: + # full_dim is the model dim (before TP). + # extra > 0 means there is expected overall increase of dimensions. + # This is so because of replication. + # duplicate_groups is used to map the tp_rank to the actual shard rank. + # This is useful when there is replication of groups to accompany head shards. + + # Size of the loaded shard + shard_size = full_dim // tp_size + + # Compute the rank into the loaded shard. + # If there is replication, different TP shards will take from the same rank. + # NOTE: currently we only support duplication in the case where num_groups == 1 + rank = 0 if duplicate_groups else tp_rank + + # Leftmost boundary index into loaded weight. + loaded_skip = rank * shard_size + loaded_start_idx = loaded_boundary + loaded_skip + + # Take these many dims from the loaded weight. + take = min(shard_size, full_dim - extra - loaded_skip) + + # Always shard on dim 0 + param.data[ + boundary : (boundary + take), ... # type: ignore[misc] + ] = loaded_weight[ + loaded_start_idx : ( + loaded_start_idx + take + ) # type: ignore[misc] + ] # type: ignore[misc] + + # Move indexing boundaries + boundary += shard_size + loaded_boundary += full_dim - extra + + return loader + + +class Qwen3NextRMSNorm(nn.Module): + """RMSNorm implementation for Qwen3Next (uses 1.0 + weight instead of weight)""" + + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.zeros(dim)) + + def _norm(self, x): + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + + def forward(self, x, residual: torch.Tensor | None = None): + # Align with HF Qwen3_5: norm only normalizes x; residual is NOT added into norm input. + # When residual is None: return (norm(x), x) so caller gets residual = pre-norm x for the add. + # When residual is not None: return (norm(x), residual) so caller adds the same residual after attn/mlp. + orig_dtype = x.dtype + output = self._norm(x.float()) + # Llama does x.to(float16) * w whilst Qwen3Next is (x * w).to(float16) + # See https://github.com/huggingface/transformers/pull/29402 + output = output * (1.0 + self.weight.float()) + output = output.type_as(x) + + if residual is not None: + return output, residual + return output, x + + def extra_repr(self): + return f"{tuple(self.weight.shape)}, eps={self.eps}" + + +class Qwen3NextRMSNormGated(nn.Module): + """RMSNorm with gate mechanism for GatedDeltaNet""" + + def __init__(self, hidden_size, eps=1e-6): + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states, gate=None): + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + # Norm before gate + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + hidden_states = self.weight * hidden_states.to(input_dtype) + if gate is not None: + hidden_states = hidden_states * F.silu(gate.to(torch.float32)) + return hidden_states.to(input_dtype) + + +def torch_causal_conv1d_update( + hidden_states, + conv_state, + weight, + bias=None, + activation=None, +): + """Update causal conv1d state and compute output""" + _, hidden_size, seq_len = hidden_states.shape + state_len = conv_state.shape[-1] + + hidden_states_new = torch.cat([conv_state, hidden_states], dim=-1).to(weight.dtype) + conv_state.copy_(hidden_states_new[:, :, -state_len:]) + out = F.conv1d(hidden_states_new, weight.unsqueeze(1), bias, padding=0, groups=hidden_size) + if activation == "silu": + out = F.silu(out[:, :, -seq_len:]) + else: + out = out[:, :, -seq_len:] + out = out.to(hidden_states.dtype) + return out + + +def l2norm(x: torch.Tensor, dim: int = -1, eps: float = 1e-6): + """L2 normalization""" + inv_norm = torch.rsqrt((x * x).sum(dim=dim, keepdim=True) + eps) + return x * inv_norm + + +def torch_recurrent_gated_delta_rule( + query, key, value, g, beta, initial_state=None, output_final_state=False, use_qk_l2norm_in_kernel=False +): + """Recurrent gated delta rule implementation""" + initial_dtype = query.dtype + if use_qk_l2norm_in_kernel: + query = l2norm(query, dim=-1, eps=1e-6) + key = l2norm(key, dim=-1, eps=1e-6) + + # Reshape: (seq_len, num_heads, head_dim) -> (batch=1, num_heads, seq_len, head_dim) + if query.dim() == 3: + query = query.unsqueeze(0) + key = key.unsqueeze(0) + value = value.unsqueeze(0) + g = g.unsqueeze(0) if g.dim() == 2 else g + beta = beta.unsqueeze(0) if beta.dim() == 2 else beta + + query, key, value, beta, g = [ + x.transpose(1, 2).contiguous().to(torch.float32) for x in (query, key, value, beta, g) + ] + + batch_size, num_heads, sequence_length, k_head_dim = key.shape + v_head_dim = value.shape[-1] + # HF-style: apply scale even with L2 norm (aligns with HF implementation) + scale = 1 / (query.shape[-1] ** 0.5) + query = query * scale + g = g.clamp(min=-50.0) + g = torch.nan_to_num(g, nan=0.0, posinf=0.0, neginf=-50.0) + + core_attn_out = torch.zeros(batch_size, num_heads, sequence_length, v_head_dim).to(value) + + # Important for CUDA graph replay: + # We must update the *original* `initial_state` storage in-place (copy back), + # because Python-side state cache updates do not run during replay. + initial_state_orig = initial_state + last_recurrent_state = ( + torch.zeros(batch_size, num_heads, k_head_dim, v_head_dim).to(value) + if initial_state_orig is None + else initial_state_orig.to(value) # working copy (typically float32) + ) + + for i in range(sequence_length): + q_t = query[:, :, i] + k_t = key[:, :, i] + v_t = value[:, :, i] + g_t = g[:, :, i].exp().unsqueeze(-1).unsqueeze(-1) + beta_t = beta[:, :, i].unsqueeze(-1) + + last_recurrent_state = last_recurrent_state * g_t + kv_mem = (last_recurrent_state * k_t.unsqueeze(-1)).sum(dim=-2) + delta = (v_t - kv_mem) * beta_t + last_recurrent_state = last_recurrent_state + k_t.unsqueeze(-1) * delta.unsqueeze(-2) + core_attn_out[:, :, i] = (last_recurrent_state * q_t.unsqueeze(-1)).sum(dim=-2) + + # Persist updated recurrent state back into the original tensor object. + # This enables state to survive across CUDA graph replay steps. + if output_final_state and initial_state_orig is not None: + initial_state_orig.copy_(last_recurrent_state.to(initial_state_orig.dtype)) + last_recurrent_state = initial_state_orig + elif not output_final_state: + last_recurrent_state = None + + core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype) + # Remove batch dimension if it was added + if core_attn_out.shape[0] == 1: + core_attn_out = core_attn_out.squeeze(0) + + return core_attn_out, last_recurrent_state + + +def torch_chunk_gated_delta_rule( + query, key, value, g, beta, chunk_size=64, initial_state=None, output_final_state=False, use_qk_l2norm_in_kernel=False +): + """Chunk-based gated delta rule implementation""" + initial_dtype = query.dtype + if use_qk_l2norm_in_kernel: + query = l2norm(query, dim=-1, eps=1e-6) + key = l2norm(key, dim=-1, eps=1e-6) + + # Reshape: (seq_len, num_heads, head_dim) -> (batch=1, num_heads, seq_len, head_dim) + if query.dim() == 3: + query = query.unsqueeze(0) + key = key.unsqueeze(0) + value = value.unsqueeze(0) + g = g.unsqueeze(0) if g.dim() == 2 else g + beta = beta.unsqueeze(0) if beta.dim() == 2 else beta + + query, key, value, beta, g = [ + x.transpose(1, 2).contiguous().to(torch.float32) for x in (query, key, value, beta, g) + ] + + batch_size, num_heads, sequence_length, k_head_dim = key.shape + v_head_dim = value.shape[-1] + pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size + 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_sequence_length = sequence_length + pad_size + # HF-style: apply scale even with L2 norm (aligns with HF implementation) + scale = 1 / (query.shape[-1] ** 0.5) + query = query * scale + + # Numerical stability: avoid -inf in g (e.g. from uninitialized or extreme weights) + # which would make cumsum -inf and (g[i]-g[j]).exp() -> NaN. Clamp g to finite range. + g = g.clamp(min=-50.0) + g = torch.nan_to_num(g, nan=0.0, posinf=0.0, neginf=-50.0) + + v_beta = value * beta.unsqueeze(-1) + k_beta = key * beta.unsqueeze(-1) + # reshape to chunks + query, key, value, k_beta, v_beta = [ + x.reshape(x.shape[0], x.shape[1], -1, chunk_size, x.shape[-1]) for x in (query, key, value, k_beta, v_beta) + ] + g = g.reshape(g.shape[0], g.shape[1], -1, chunk_size) + mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), diagonal=0) + + # chunk decay + g = g.cumsum(dim=-1) + decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).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)) + last_recurrent_state = ( + torch.zeros(batch_size, num_heads, k_head_dim, v_head_dim).to(value) + if initial_state is None + else initial_state.to(value) + ) + core_attn_out = torch.zeros_like(value) + mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), diagonal=1) + + # for each chunk + for i in range(0, total_sequence_length // chunk_size): + q_i, k_i, v_i = query[:, :, i], key[:, :, i], value[:, :, i] + attn = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask, 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 @ 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 + ) + + if not output_final_state: + last_recurrent_state = None + core_attn_out = core_attn_out.reshape(core_attn_out.shape[0], core_attn_out.shape[1], -1, core_attn_out.shape[-1]) + core_attn_out = core_attn_out[:, :, :sequence_length] + core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype) + # Remove batch dimension if it was added + if core_attn_out.shape[0] == 1: + core_attn_out = core_attn_out.squeeze(0) + + return core_attn_out, last_recurrent_state + + +class Qwen3NextAttention(nn.Module): + + def __init__( + self, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + max_position: int = 4096 * 32, + head_dim: int | None = None, + rms_norm_eps: float = 1e-06, + qkv_bias: bool = False, + rope_theta: float = 10000, + rope_scaling: tuple | None = None, + ) -> None: + super().__init__() + tp_size = dist.get_world_size() + self.total_num_heads = num_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = num_kv_heads + if self.total_num_kv_heads >= tp_size: + # Number of KV heads is greater than TP size, so we partition + # the KV heads across multiple tensor parallel GPUs. + assert self.total_num_kv_heads % tp_size == 0 + else: + # Number of KV heads is less than TP size, so we replicate + # the KV heads across multiple tensor parallel GPUs. + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + self.head_dim = head_dim or hidden_size // self.total_num_heads + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim ** -0.5 + self.qkv_bias = qkv_bias + + # q_proj outputs num_heads * head_dim * 2 (query + gate) + self.q_proj = ColumnParallelLinear( + hidden_size, + self.total_num_heads * self.head_dim * 2, + bias=qkv_bias, + ) + self.k_proj = ColumnParallelLinear( + hidden_size, + self.total_num_kv_heads * self.head_dim, + bias=qkv_bias, + ) + self.v_proj = ColumnParallelLinear( + hidden_size, + self.total_num_kv_heads * self.head_dim, + bias=qkv_bias, + ) + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + hidden_size, + bias=False, + ) + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.head_dim, + max_position=max_position, + base=rope_theta, + rope_scaling=rope_scaling, + ) + self.attn = Attention( + self.num_heads, + self.head_dim, + self.scaling, + self.num_kv_heads, + ) + # Always use q_norm and k_norm for Qwen3Next + self.q_norm = Qwen3NextRMSNorm(self.head_dim, eps=rms_norm_eps) + self.k_norm = Qwen3NextRMSNorm(self.head_dim, eps=rms_norm_eps) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + # vLLM-style: cast to weight dtype + hidden_states = hidden_states.to(self.q_proj.weight.dtype) + # q_proj outputs query + gate (concatenated) + # Shape: (seq_len, num_heads * head_dim * 2) + q_gate = self.q_proj(hidden_states) + # Reshape and split into query and gate (reshape for stride safety) + q_gate = q_gate.reshape(-1, self.num_heads, self.head_dim * 2) + query_states, gate = q_gate.split([self.head_dim, self.head_dim], dim=-1) + k = self.k_proj(hidden_states) + k = k.reshape(-1, self.num_kv_heads, self.head_dim) + v = self.v_proj(hidden_states) + v = v.reshape(-1, self.num_kv_heads, self.head_dim) + + # Apply normalization + # Qwen3NextRMSNorm returns (output, x) when called with one arg + query_states, _ = self.q_norm(query_states) + k, _ = self.k_norm(k) + + # Apply rotary embedding + query_states, k = self.rotary_emb(positions, query_states, k) + + # Attention + o = self.attn(query_states, k, v) + # o shape: (seq_len, num_heads, head_dim) + + # Reshape gate and apply sigmoid + # Flatten gate to match attention output shape + gate = gate.flatten(1, -1) # (seq_len, num_heads * head_dim) + o_flat = o.flatten(1, -1) # (seq_len, num_heads * head_dim) + attn_output = o_flat * torch.sigmoid(gate) + + output = self.o_proj(attn_output) + return output + + +class Qwen3NextMLP(nn.Module): + + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + ) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + ) + assert hidden_act == "silu" + self.act_fn = SiluAndMul() + + def forward(self, x): + gate_up = self.gate_up_proj(x) + if isinstance(gate_up, tuple): + gate_up = gate_up[0] + x = self.act_fn(gate_up) + x = self.down_proj(x) + if isinstance(x, tuple): + x = x[0] + return x + + +class Qwen3NextSparseMoeBlock(nn.Module): + """ + Sparse MoE block: gate (router) + experts (ModuleList) + optional shared expert. + Param paths: mlp.gate.weight, mlp.experts.0.*, ... to match checkpoint. + """ + + def __init__(self, config: Qwen3NextConfig) -> None: + super().__init__() + num_experts = getattr(config, "num_experts", 0) + if num_experts <= 0: + raise ValueError("num_experts must be > 0 for SparseMoeBlock") + self.num_experts = num_experts + top_k = getattr(config, "num_experts_per_tok", 1) + self.top_k = min(top_k, num_experts) + self.norm_topk_prob = getattr(config, "norm_topk_prob", True) + shared_size = getattr(config, "shared_expert_intermediate_size", 0) + hidden_act = getattr(config, "hidden_act", "silu") + moe_intermediate_size = getattr( + config, "moe_intermediate_size", config.intermediate_size + ) + + # Gate linear directly under mlp so path is mlp.gate.weight (match checkpoint) + self.gate = ReplicatedLinear( + config.hidden_size, num_experts, bias=False + ) + # Direct ModuleList so param path is mlp.experts.0, experts.1, ... (match checkpoint) + self.experts = nn.ModuleList([ + Qwen3NextMLP( + hidden_size=config.hidden_size, + intermediate_size=moe_intermediate_size, + hidden_act=hidden_act, + ) + for _ in range(num_experts) + ]) + if shared_size > 0: + self.shared_expert = Qwen3NextMLP( + hidden_size=config.hidden_size, + intermediate_size=shared_size, + hidden_act=hidden_act, + ) + self.shared_expert_gate = ReplicatedLinear(config.hidden_size, 1, bias=False) + else: + self.shared_expert = None + self.shared_expert_gate = None + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + orig_shape = hidden_states.shape + # vLLM-style: cast to weight dtype to avoid float vs bfloat16 matmul + hidden_states = hidden_states.to(self.gate.weight.dtype) + hidden_states = hidden_states.reshape(-1, hidden_states.size(-1)) + # Router: gate -> softmax -> top-k (was Qwen3NextTopKRouter) + router_logits = self.gate(hidden_states) + if isinstance(router_logits, tuple): + router_logits = router_logits[0] + probs = F.softmax(router_logits.float(), dim=-1).to(hidden_states.dtype) + top_k_weights, top_k_indices = torch.topk(probs, self.top_k, dim=-1) + if self.norm_topk_prob: + top_k_weights = top_k_weights / ( + top_k_weights.sum(dim=-1, keepdim=True) + 1e-10 + ) + # Expert dispatch + T, H = hidden_states.shape + top_k = top_k_indices.size(-1) + final = torch.zeros_like(hidden_states) + for k in range(top_k): + expert_idx = top_k_indices[:, k] + w = top_k_weights[:, k] + for e in range(self.num_experts): + mask = (expert_idx == e) + if not mask.any(): + continue + tokens_e = mask.nonzero(as_tuple=True)[0] + x_e = hidden_states[tokens_e] + out_e = self.experts[e](x_e) + final[tokens_e] = ( + final[tokens_e] + + out_e * w[tokens_e].unsqueeze(-1).to(out_e.dtype) + ) + expert_output = final + if self.shared_expert is not None: + shared_out = self.shared_expert(hidden_states) + gate_out = self.shared_expert_gate(hidden_states) + if isinstance(gate_out, tuple): + gate_out = gate_out[0] + shared_out = shared_out * torch.sigmoid(gate_out) + expert_output = expert_output + shared_out + return expert_output.reshape(orig_shape) + + +class Qwen3NextGatedDeltaNet(nn.Module): + """Gated Delta Net for linear attention in Qwen3Next""" + + def __init__( + self, + config: Qwen3NextConfig, + layer_idx: int = 0, + ) -> None: + super().__init__() + tp_size = dist.get_world_size() + self.hidden_size = config.hidden_size + self.num_v_heads = getattr(config, "linear_num_value_heads", 32) + self.num_k_heads = getattr(config, "linear_num_key_heads", 16) + self.head_k_dim = getattr(config, "linear_key_head_dim", 128) + self.head_v_dim = getattr(config, "linear_value_head_dim", 128) + self.key_dim = self.head_k_dim * self.num_k_heads + self.value_dim = self.head_v_dim * self.num_v_heads + + self.conv_kernel_size = getattr(config, "linear_conv_kernel_dim", 4) + self.layer_idx = layer_idx + self.activation = getattr(config, "hidden_act", "silu") + self.layer_norm_epsilon = config.rms_norm_eps + + # QKV projection + self.conv_dim = self.key_dim * 2 + self.value_dim + + # Conv1d layer (for causal convolution) + # Use ColumnParallelLinear instead of nn.Conv1d for tensor parallelism + # The weight shape will be (conv_dim // tp_size, conv_kernel_size) + # After unsqueeze(1), it becomes (conv_dim // tp_size, 1, conv_kernel_size) + self.conv1d = ColumnParallelLinear( + input_size=self.conv_kernel_size, + output_size=self.conv_dim, + bias=False, + ) + # Add dimension for conv1d: (conv_dim, kernel_size) -> (conv_dim, 1, kernel_size) + self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1) + + # Set up custom weight loader for conv1d + # The conv1d weight needs special handling because it contains Q, K, V projections + query_key_settings = (self.key_dim, 0, False) + value_settings = (self.value_dim, 0, False) + + # Remove default weight_loader and set custom one + if hasattr(self.conv1d.weight, "weight_loader"): + delattr(self.conv1d.weight, "weight_loader") + + # Set custom weight loader + self.conv1d.weight.weight_loader = mamba_v2_sharded_weight_loader( + [ + query_key_settings, # Q projection + query_key_settings, # K projection + value_settings, # V projection + ], + tp_size, + dist.get_rank(), + ) + + # Projection layers - Qwen3Next uses combined qkvz projection + projection_size_qkvz = self.key_dim * 2 + self.value_dim * 2 + projection_size_ba = self.num_v_heads * 2 + + self.in_proj_qkvz = ColumnParallelLinear( + self.hidden_size, + projection_size_qkvz, + bias=False, + ) + self.in_proj_ba = ColumnParallelLinear( + self.hidden_size, + projection_size_ba, + bias=False, + ) + + # Time step projection parameters (sharded by TP along dim 0, vLLM-style) + self.dt_bias = nn.Parameter(torch.ones(self.num_v_heads // tp_size)) + A = torch.empty(self.num_v_heads // tp_size).uniform_(0, 16) + self.A_log = nn.Parameter(torch.log(A)) + self.A_log.weight_loader = sharded_weight_loader(0) + self.dt_bias.weight_loader = sharded_weight_loader(0) + + # Normalization with gate + self.norm = Qwen3NextRMSNormGated(self.head_v_dim, eps=self.layer_norm_epsilon) + + # Output projection + self.out_proj = RowParallelLinear( + self.value_dim, + self.hidden_size, + bias=False, + ) + + # Cache states (for incremental decoding) + # For single sequence: use instance variables + # For batch inference: use state_cache dict with sequence indices + self.conv_state = None + self.recurrent_state = None + self.state_cache = {} # Dict mapping sequence_id -> (conv_state, recurrent_state) + + def fix_query_key_value_ordering(self, mixed_qkvz, mixed_ba): + """ + Derives `query`, `key`, `value`, `z`, `b`, `a` tensors from `mixed_qkvz` and `mixed_ba`. + This follows the vllm implementation structure. + """ + tp_size = dist.get_world_size() + seq_len = mixed_qkvz.size(0) + + # Reshape mixed_qkvz: (seq_len, ...) -> (seq_len, num_k_heads/tp, ...) + # The structure is: [q_per_k_head, k_per_k_head, v_per_k_head, z_per_k_head] * (num_v_heads / num_k_heads) + num_v_per_k = self.num_v_heads // self.num_k_heads + new_tensor_shape_qkvz = (seq_len, self.num_k_heads // tp_size, + self.head_k_dim + self.head_k_dim + + num_v_per_k * self.head_v_dim + num_v_per_k * self.head_v_dim) + new_tensor_shape_ba = (seq_len, self.num_k_heads // tp_size, + 2 * num_v_per_k) + + mixed_qkvz = mixed_qkvz.reshape(*new_tensor_shape_qkvz) + mixed_ba = mixed_ba.reshape(*new_tensor_shape_ba) + + # Split along the last dimension + split_arg_list_qkvz = [ + self.head_k_dim, # q + self.head_k_dim, # k + num_v_per_k * self.head_v_dim, # v (grouped by num_v_per_k) + num_v_per_k * self.head_v_dim, # z (grouped by num_v_per_k) + ] + split_arg_list_ba = [ + num_v_per_k, # b + num_v_per_k, # a + ] + + (query, key, value, z) = torch.split(mixed_qkvz, split_arg_list_qkvz, dim=-1) + (b, a) = torch.split(mixed_ba, split_arg_list_ba, dim=-1) + + # Reshape value and z: (seq_len, num_k_heads/tp, num_v_per_k * head_v_dim) + # -> (seq_len, num_v_heads/tp, head_v_dim) + value = value.reshape(seq_len, -1, self.head_v_dim) + z = z.reshape(seq_len, -1, self.head_v_dim) + b = b.reshape(seq_len, -1) + a = a.reshape(seq_len, -1) + + return query, key, value, z, b, a + + def forward( + self, + hidden_states: torch.Tensor, + positions: torch.Tensor | None = None, + sequence_id: int | None = None, + ) -> torch.Tensor: + """ + Forward pass for GatedDeltaNet + + Args: + hidden_states: (seq_len, hidden_size) + positions: (seq_len,) - not used in linear attention but kept for compatibility + sequence_id: Optional sequence ID for batch inference state management + + Returns: + output: (seq_len, hidden_size) + """ + seq_len = hidden_states.size(0) + # vLLM-style: cast to weight dtype for linear layers + hidden_states = hidden_states.to(self.in_proj_qkvz.weight.dtype) + + # Get or initialize states for this sequence + # CRITICAL: Each sequence must have independent states to prevent cross-contamination + if sequence_id is not None: + # Batch inference: use state_cache with per-sequence isolation + if sequence_id not in self.state_cache: + # Initialize fresh states for new sequence + self.state_cache[sequence_id] = (None, None) + conv_state, recurrent_state = self.state_cache[sequence_id] + # Ensure we get references, not copies + if conv_state is None: + conv_state = None # Will be initialized later + if recurrent_state is None: + recurrent_state = None # Will be initialized later + else: + # Single sequence: use instance variables + conv_state = self.conv_state + recurrent_state = self.recurrent_state + + # Part 1: Input Projection + # ColumnParallelLinear returns (output, _) tuple in some implementations + projected_states_qkvz = self.in_proj_qkvz(hidden_states) + if isinstance(projected_states_qkvz, tuple): + projected_states_qkvz = projected_states_qkvz[0] + + projected_states_ba = self.in_proj_ba(hidden_states) + if isinstance(projected_states_ba, tuple): + projected_states_ba = projected_states_ba[0] + + query, key, value, z, b, a = self.fix_query_key_value_ordering( + projected_states_qkvz, projected_states_ba + ) + + # Flatten for convolution: (seq_len, num_heads, head_dim) -> (seq_len, num_heads * head_dim) + query_flat = query.reshape(seq_len, -1) + key_flat = key.reshape(seq_len, -1) + value_flat = value.reshape(seq_len, -1) + + mixed_qkv = torch.cat([query_flat, key_flat, value_flat], dim=-1) + # mixed_qkv shape: (seq_len, conv_dim // tp_size) + # Transpose for conv1d: (seq_len, conv_dim // tp_size) -> (conv_dim // tp_size, seq_len) + mixed_qkv = mixed_qkv.transpose(0, 1).unsqueeze(0) # (1, conv_dim // tp_size, seq_len) + + # Part 2: Causal Convolution + if conv_state is None: + # First pass: initialize conv_state for this specific sequence + # CRITICAL: Create new tensor for each sequence to ensure isolation + conv_state = torch.zeros( + 1, self.conv_dim // dist.get_world_size(), self.conv_kernel_size - 1, + dtype=mixed_qkv.dtype, device=mixed_qkv.device + ) + # Update state cache immediately to ensure isolation + if sequence_id is not None: + self.state_cache[sequence_id] = (conv_state, recurrent_state) + else: + self.conv_state = conv_state + + # Apply causal convolution + # Extract conv weights: (conv_dim // tp_size, 1, kernel_size) -> (conv_dim // tp_size, kernel_size) + conv_weights = self.conv1d.weight.reshape( + self.conv1d.weight.size(0), self.conv1d.weight.size(2) + ) + mixed_qkv = torch_causal_conv1d_update( + mixed_qkv, + conv_state, + conv_weights, # (conv_dim // tp_size, kernel_size) + self.conv1d.bias, + self.activation, + ) + + # Update conv_state (it's modified in-place by torch_causal_conv1d_update) + # CRITICAL: Always update state cache to ensure state persistence + if sequence_id is not None: + # Update state cache with current states (conv_state is modified in-place) + self.state_cache[sequence_id] = (conv_state, recurrent_state) + else: + self.conv_state = conv_state + + # Transpose back: (1, conv_dim, seq_len) -> (seq_len, conv_dim) + mixed_qkv = mixed_qkv.squeeze(0).transpose(0, 1) + + # Split back into query, key, value + query_conv, key_conv, value_conv = torch.split( + mixed_qkv, + [self.key_dim // dist.get_world_size(), + self.key_dim // dist.get_world_size(), + self.value_dim // dist.get_world_size()], + dim=-1 + ) + + # Reshape back to (seq_len, num_heads, head_dim) + query_conv = query_conv.reshape(seq_len, self.num_k_heads // dist.get_world_size(), self.head_k_dim) + key_conv = key_conv.reshape(seq_len, self.num_k_heads // dist.get_world_size(), self.head_k_dim) + value_conv = value_conv.reshape(seq_len, self.num_v_heads // dist.get_world_size(), self.head_v_dim) + + # Repeat query/key if num_v_heads > num_k_heads + if self.num_v_heads // dist.get_world_size() > self.num_k_heads // dist.get_world_size(): + repeat_factor = (self.num_v_heads // dist.get_world_size()) // (self.num_k_heads // dist.get_world_size()) + query_conv = query_conv.repeat_interleave(repeat_factor, dim=1) + key_conv = key_conv.repeat_interleave(repeat_factor, dim=1) + + # Compute gating parameters + beta = b.sigmoid() # (seq_len, num_v_heads) + # g = -exp(A_log) * softplus(a + dt_bias) + g = -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias.float()) + g = g.to(hidden_states.dtype) + # Numerical stability for both Triton and torch paths + g = g.clamp(min=-50.0) + g = torch.nan_to_num(g, nan=0.0, posinf=0.0, neginf=-50.0) + + # Part 3: Gated Delta Rule (Recurrent Attention) + # CRITICAL: Determine mode based on sequence length and state availability + use_recurrent = (seq_len == 1 and recurrent_state is not None) + + # Ensure recurrent_state is properly initialized for chunk mode if needed + if not use_recurrent and recurrent_state is None: + # Initialize recurrent_state for chunk mode (prefill) + # Shape: (num_v_heads, head_v_dim, head_k_dim) per sequence + tp_size = dist.get_world_size() + num_v_heads_shard = self.num_v_heads // tp_size + recurrent_state = torch.zeros( + num_v_heads_shard, + self.head_v_dim, + self.head_k_dim, + dtype=value_conv.dtype, + device=value_conv.device + ) + + # Convert from (seq_len, num_heads, head_dim) to (B=1, T=seq_len, H, K) format + # Add batch dimension + q_batch = query_conv.unsqueeze(0) # (1, seq_len, num_k_heads, head_k_dim) + k_batch = key_conv.unsqueeze(0) # (1, seq_len, num_k_heads, head_k_dim) + v_batch = value_conv.unsqueeze(0) # (1, seq_len, num_v_heads, head_v_dim) + g_batch = g.unsqueeze(0) # (1, seq_len, num_v_heads) + beta_batch = beta.unsqueeze(0) # (1, seq_len, num_v_heads) + + # Prepare initial_state if needed: (num_v_heads, head_v_dim, head_k_dim) -> (1, num_v_heads, head_v_dim, head_k_dim) + initial_state_batch = None + if recurrent_state is not None: + if recurrent_state.dim() == 3: + initial_state_batch = recurrent_state.unsqueeze(0) # Add batch dim + else: + initial_state_batch = recurrent_state + + if use_recurrent and HAS_FUSED_RECURRENT and fused_recurrent_gated_delta_rule is not None: + # Use optimized Triton kernel for recurrent mode + # Note: inplace_final_state=False to avoid requiring ssm_state_indices + # When inplace_final_state=True, kernel expects ssm_state_indices which we don't provide + core_attn_out_batch, final_state_batch = fused_recurrent_gated_delta_rule( + q=q_batch, + k=k_batch, + v=v_batch, + g=g_batch, + beta=beta_batch, + initial_state=initial_state_batch, + inplace_final_state=False, + use_qk_l2norm_in_kernel=True, + ) + # Remove batch dimension: (1, seq_len, num_v_heads, head_v_dim) -> (seq_len, num_v_heads, head_v_dim) + core_attn_out = core_attn_out_batch.squeeze(0) + if final_state_batch is not None: + recurrent_state = final_state_batch.squeeze(0) if final_state_batch.dim() == 4 else final_state_batch + elif not use_recurrent and HAS_CHUNK and chunk_gated_delta_rule is not None: + # Use optimized Triton kernel for chunk mode + core_attn_out_batch, final_state_batch = chunk_gated_delta_rule( + q=q_batch, + k=k_batch, + v=v_batch, + g=g_batch, + beta=beta_batch, + initial_state=initial_state_batch, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + ) + # Remove batch dimension + core_attn_out = core_attn_out_batch.squeeze(0) + if final_state_batch is not None: + recurrent_state = final_state_batch.squeeze(0) if final_state_batch.dim() == 4 else final_state_batch + else: + # Fallback to Python implementation + if use_recurrent: + core_attn_out, recurrent_state = torch_recurrent_gated_delta_rule( + query_conv, + key_conv, + value_conv, + g, + beta, + initial_state=recurrent_state, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + ) + else: + core_attn_out, recurrent_state = torch_chunk_gated_delta_rule( + query_conv, + key_conv, + value_conv, + g, + beta, + chunk_size=64, + initial_state=recurrent_state, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + ) + + # CRITICAL: Always update state cache to ensure state persistence per sequence + if sequence_id is not None: + # Store updated states in cache for this specific sequence + self.state_cache[sequence_id] = (conv_state, recurrent_state) + else: + self.recurrent_state = recurrent_state + + # Part 4: Output Projection with Gate + # CRITICAL: Gate application must match vLLM implementation + # core_attn_out: (seq_len, num_v_heads, head_v_dim) + # z: (seq_len, num_v_heads, head_v_dim) + z_shape_og = z.shape + core_attn_out_flat = core_attn_out.reshape(-1, core_attn_out.shape[-1]) + z_flat = z.reshape(-1, z.shape[-1]) + + # Apply normalization with gate: norm(core_attn_out) * silu(gate) + # The gate (z) is applied via SiLU activation in Qwen3NextRMSNormGated + core_attn_out_flat = self.norm(core_attn_out_flat, z_flat) + core_attn_out = core_attn_out_flat.reshape(z_shape_og) + + # Flatten for output projection + core_attn_out = core_attn_out.reshape(seq_len, -1) + + # Output projection + # RowParallelLinear may return tuple in some implementations + output = self.out_proj(core_attn_out) + if isinstance(output, tuple): + output = output[0] + + return output + + +class Qwen3NextDecoderLayer(nn.Module): + + def __init__( + self, + config: Qwen3NextConfig, + layer_idx: int = 0, + ) -> None: + super().__init__() + rope_scaling = getattr(config, "rope_scaling", None) + if isinstance(rope_scaling, dict): + rope_scaling = None # TODO: MRoPE support + + # layer_types: "full_attention" or "linear_attention" per layer (vllm/transformers) + layer_types = getattr(config, "layer_types", None) + if layer_types is not None: + self.layer_type = layer_types[layer_idx] if layer_idx < len(layer_types) else "full_attention" + else: + self.layer_type = "full_attention" + + if self.layer_type == "full_attention": + self.self_attn = Qwen3NextAttention( + hidden_size=config.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=config.num_key_value_heads, + max_position=config.max_position_embeddings, + rms_norm_eps=config.rms_norm_eps, + qkv_bias=getattr(config, 'attention_bias', False), + head_dim=getattr(config, 'head_dim', None), + rope_theta=getattr(config, "rope_theta", 1000000), + rope_scaling=rope_scaling, + ) + self.linear_attn = None + elif self.layer_type == "linear_attention": + self.self_attn = None + self.linear_attn = Qwen3NextGatedDeltaNet( + config=config, + layer_idx=layer_idx, + ) + else: + raise ValueError(f"Invalid layer_type: {self.layer_type}") + + # MLP: standard MLP or Sparse MoE (vllm/transformers logic) + mlp_only_layers = getattr(config, "mlp_only_layers", []) + num_experts = getattr(config, "num_experts", 0) + decoder_sparse_step = getattr(config, "decoder_sparse_step", 1) + use_moe = ( + (layer_idx not in mlp_only_layers) + and (num_experts > 0) + and ((layer_idx + 1) % decoder_sparse_step == 0) + ) + if use_moe: + self.mlp = Qwen3NextSparseMoeBlock(config) + else: + self.mlp = Qwen3NextMLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + ) + + self.input_layernorm = Qwen3NextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = Qwen3NextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if residual is None: + hidden_states, residual = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + + if self.layer_type == "full_attention": + hidden_states = self.self_attn(positions, hidden_states) + elif self.layer_type == "linear_attention": + hidden_states = self.linear_attn(hidden_states, positions) + else: + raise ValueError(f"Invalid layer_type: {self.layer_type}") + + hidden_states = hidden_states + residual + # HF: residual = hidden_states, then norm(hidden_states); add residual after mlp + residual = hidden_states + hidden_states, residual = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = hidden_states + residual + return hidden_states, residual + + +class Qwen3NextModel(nn.Module): + + def __init__( + self, + config: Qwen3NextConfig, + ) -> None: + super().__init__() + self.embed_tokens = VocabParallelEmbedding(config.vocab_size, config.hidden_size) + self.layers = nn.ModuleList([ + Qwen3NextDecoderLayer(config, layer_idx=i) + for i in range(config.num_hidden_layers) + ]) + self.norm = Qwen3NextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + input_ids: torch.Tensor | None = None, + positions: torch.Tensor = None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor: + # 支持直接传入 embeddings（用于多模态） + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = self.embed_tokens(input_ids) + # vLLM-style: keep activation dtype consistent with weights + hidden_states = hidden_states.to(self.embed_tokens.weight.dtype) + + residual = None + for layer in self.layers: + hidden_states, residual = layer(positions, hidden_states, residual) + hidden_states, _ = self.norm(hidden_states, residual) + return hidden_states + + +class Qwen3NextForCausalLM(nn.Module): + packed_modules_mapping = { + "q_proj": ("q_proj", None), + "k_proj": ("k_proj", None), + "v_proj": ("v_proj", None), + "gate_proj": ("gate_up_proj", 0), + "up_proj": ("gate_up_proj", 1), + } + + def __init__( + self, + config: Qwen3NextConfig + ) -> None: + super().__init__() + self.model = Qwen3NextModel(config) + self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size) + if config.tie_word_embeddings: + self.lm_head.weight.data = self.model.embed_tokens.weight.data + + def forward( + self, + input_ids: torch.Tensor | None = None, + positions: torch.Tensor = None, + inputs_embeds: torch.Tensor | None = None, + pixel_values=None, # 多模态参数（向后兼容） + image_grid_thw=None, # 多模态参数（向后兼容） + **kwargs # 其他参数 + ) -> torch.Tensor: + return self.model(input_ids, positions, inputs_embeds=inputs_embeds) + + def compute_logits( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + hidden_states = hidden_states.to(self.lm_head.weight.dtype) + return self.lm_head(hidden_states) diff --git a/nanovllm/utils/loader.py b/nanovllm/utils/loader.py index 4ef804006..ae2df2793 100644 --- a/nanovllm/utils/loader.py +++ b/nanovllm/utils/loader.py @@ -1,6 +1,8 @@ +import inspect import os from glob import glob import torch +import torch.distributed as dist from torch import nn from safetensors import safe_open @@ -9,20 +11,79 @@ def default_weight_loader(param: nn.Parameter, loaded_weight: torch.Tensor): param.data.copy_(loaded_weight) -def load_model(model: nn.Module, path: str): +def sharded_weight_loader(shard_axis: int): + """ + Return a weight loader that shards loaded_weight along shard_axis by TP rank + (vLLM-style). Use for A_log, dt_bias when tensor_parallel_size > 1. + """ + + def loader(param: nn.Parameter, loaded_weight: torch.Tensor) -> None: + tp_rank = dist.get_rank() + shard_size = param.data.shape[shard_axis] + start_idx = tp_rank * shard_size + loaded_shard = loaded_weight.narrow(shard_axis, start_idx, shard_size) + param.data.copy_(loaded_shard) + + return loader + + +def load_model(model: nn.Module, path: str, name_mapping=None): packed_modules_mapping = getattr(model, "packed_modules_mapping", {}) for file in glob(os.path.join(path, "*.safetensors")): with safe_open(file, "pt", "cpu") as f: for weight_name in f.keys(): - for k in packed_modules_mapping: + # Skip mtp.* (e.g. speculative / MTP adapter), same as vLLM + if weight_name.startswith("mtp."): + continue + target_name = weight_name + if name_mapping is not None: + target_name = name_mapping(target_name) + if target_name is None: + continue + + for k, (v, shard_id) in packed_modules_mapping.items(): if k in weight_name: - v, shard_id = packed_modules_mapping[k] - param_name = weight_name.replace(k, v) + param_name = target_name + if k in param_name: + param_name = param_name.replace(k, v) + elif k == "gate_proj" and "gate_up_proj" in param_name: + param_name = param_name.replace("gate_up_proj", v) + elif k == "up_proj" and "gate_up_proj" in param_name: + param_name = param_name.replace("gate_up_proj", v) param = model.get_parameter(param_name) weight_loader = getattr(param, "weight_loader") - weight_loader(param, f.get_tensor(weight_name), shard_id) + tensor = f.get_tensor(weight_name) + if tensor.dtype != param.dtype: + tensor = tensor.to(param.dtype) + if shard_id is not None: + weight_loader(param, tensor, shard_id) + else: + weight_loader(param, tensor) break else: - param = model.get_parameter(weight_name) + try: + param = model.get_parameter(target_name) + except AttributeError as e: + raise AttributeError( + f"Failed to locate parameter '{target_name}' " + f"mapped from '{weight_name}'" + ) from e weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, f.get_tensor(weight_name)) + tensor = f.get_tensor(weight_name) + if tensor.dtype != param.dtype: + tensor = tensor.to(param.dtype) + sig = inspect.signature(weight_loader) + if len(sig.parameters) >= 3 and "loaded_shard_id" in sig.parameters: + module = getattr(weight_loader, "__self__", None) + if module is not None and hasattr(module, "output_sizes"): + # Single merged tensor (e.g. gate_up_proj): split and load each shard + output_sizes = module.output_sizes + offset = 0 + for shard_id, size in enumerate(output_sizes): + part = tensor.narrow(0, offset, size) + weight_loader(param, part, shard_id) + offset += size + else: + weight_loader(param, tensor, 0) + else: + weight_loader(param, tensor)