From 21303d0f0abc169ca3d648a59d06f5aaea76432e Mon Sep 17 00:00:00 2001 From: Lucas Liebenwein <11156568+lucaslie@users.noreply.github.com> Date: Thu, 5 Feb 2026 09:11:26 -0800 Subject: [PATCH 01/17] [#11032][feat] MLA revisited and GLM 4.7 Flash support Signed-off-by: Lucas Liebenwein <11156568+lucaslie@users.noreply.github.com> Co-authored-by: Balamurugan Marimuthu <246387390+bmarimuthu-nv@users.noreply.github.com> Co-authored-by: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com> --- .../glm_4.7_flash_trtllm_cookbook.ipynb | 348 +++ .../model_registry/configs/glm-4.7-flash.yaml | 5 + .../auto_deploy/model_registry/models.yaml | 2 + .../_torch/auto_deploy/config/default.yaml | 7 +- .../_torch/auto_deploy/custom_ops/README.md | 35 - .../custom_ops/attention/torch_attention.py | 297 --- .../attention/torch_backend_attention.py | 28 +- .../auto_deploy/custom_ops/mla/__init__.py | 35 +- .../custom_ops/mla/flashinfer_mla.py | 1037 ++++++++ .../_torch/auto_deploy/custom_ops/mla/mla.py | 280 --- .../custom_ops/mla/torch_backend_mla.py | 518 ++++ .../auto_deploy/custom_ops/mla/torch_mla.py | 157 ++ .../auto_deploy/models/custom/__init__.py | 6 +- .../models/custom/modeling_deepseek.py | 655 +++++ .../models/custom/modeling_glm4_moe_lite.py | 830 +++++++ .../auto_deploy/models/patches/deepseek.py | 185 -- .../auto_deploy/models/quant_config_reader.py | 2 +- .../transform/library/fused_moe.py | 52 +- .../defs/accuracy/test_llm_api_autodeploy.py | 67 + .../_utils_test/_model_test_utils.py | 3 +- .../custom_ops/test_flashinfer_mla_op.py | 2173 +++++++++++++++++ .../singlegpu/custom_ops/test_torch_mla_op.py | 780 ++++++ .../custom_ops/test_update_kv_cache.py | 4 +- .../singlegpu/models/test_deepseek_custom.py | 494 ++++ .../singlegpu/models/test_deepseek_patches.py | 110 - .../models/test_glm4_moe_lite_modeling.py | 652 +++++ .../unit/singlegpu/models/test_sdpa_mla.py | 97 - 27 files changed, 7819 insertions(+), 1040 deletions(-) create mode 100644 examples/auto_deploy/cookbooks/glm_4.7_flash_trtllm_cookbook.ipynb create mode 100644 examples/auto_deploy/model_registry/configs/glm-4.7-flash.yaml create mode 100644 tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashinfer_mla.py delete mode 100644 tensorrt_llm/_torch/auto_deploy/custom_ops/mla/mla.py create mode 100644 tensorrt_llm/_torch/auto_deploy/custom_ops/mla/torch_backend_mla.py create mode 100644 tensorrt_llm/_torch/auto_deploy/custom_ops/mla/torch_mla.py create mode 100644 tensorrt_llm/_torch/auto_deploy/models/custom/modeling_deepseek.py create mode 100644 tensorrt_llm/_torch/auto_deploy/models/custom/modeling_glm4_moe_lite.py delete mode 100644 tensorrt_llm/_torch/auto_deploy/models/patches/deepseek.py create mode 100644 tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_flashinfer_mla_op.py create mode 100644 tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_torch_mla_op.py create mode 100644 tests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_deepseek_custom.py delete mode 100644 tests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_deepseek_patches.py create mode 100644 tests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_glm4_moe_lite_modeling.py delete mode 100644 tests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_sdpa_mla.py diff --git a/examples/auto_deploy/cookbooks/glm_4.7_flash_trtllm_cookbook.ipynb b/examples/auto_deploy/cookbooks/glm_4.7_flash_trtllm_cookbook.ipynb new file mode 100644 index 000000000000..e4733b24bea6 --- /dev/null +++ b/examples/auto_deploy/cookbooks/glm_4.7_flash_trtllm_cookbook.ipynb @@ -0,0 +1,348 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Deploying GLM-4.7-Flash with TensorRT-LLM\n", + "\n", + "This notebook walks you through deploying the `zai-org/GLM-4.7-Flash` model using TensorRT-LLM.\n", + "\n", + "[TensorRT-LLM](https://nvidia.github.io/TensorRT-LLM/) is NVIDIA's open-source library for accelerating and optimizing LLM inference on NVIDIA GPUs. Support for GLM-4.7-Flash is enabled through the AutoDeploy workflow. More details about AutoDeploy can be found [here](https://nvidia.github.io/TensorRT-LLM/torch/auto_deploy/auto-deploy.html).\n", + "\n", + "**Model Resources:**\n", + "- [HuggingFace Model Card](https://huggingface.co/zai-org/GLM-4.7-Flash)\n", + "- [Technical Blog](https://z.ai/blog/glm-4.7)\n", + "- [Technical Report (GLM-4.5)](https://arxiv.org/abs/2508.06471)\n", + "- [Z.ai API Platform](https://docs.z.ai/guides/llm/glm-4.7)\n", + "\n", + "**Model Highlights:**\n", + "- 30B-A3B Mixture of Experts (MoE) architecture\n", + "- 131,072 token context length\n", + "- Tool calling support\n", + "- MIT License\n", + "\n", + "**Prerequisites:**\n", + "- NVIDIA GPU with recent drivers (≥ 64 GB VRAM for BF16) and CUDA 12.x\n", + "- Python 3.10+\n", + "- TensorRT-LLM ([container](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/tensorrt-llm/containers/release) or pip install)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Prerequisites & Environment\n", + "\n", + "Set up a containerized environment for TensorRT-LLM by running the following command in a terminal:\n", + "\n", + "```shell\n", + "docker run --rm -it --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 --gpus=all -p 8000:8000 nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc1\n", + "```\n", + "\n", + "You now have TensorRT-LLM set up!" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# If pip not found\n", + "!python -m ensurepip --default-pip" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "%pip install torch openai" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Verify GPU\n", + "\n", + "Check that CUDA is available and the GPU is detected correctly." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Python: 3.12.3 (main, Nov 6 2025, 13:44:16) [GCC 13.3.0]\n", + "CUDA available: True\n", + "Num GPUs: 8\n", + "GPU[0]: NVIDIA H100 80GB HBM3\n", + "GPU[1]: NVIDIA H100 80GB HBM3\n", + "GPU[2]: NVIDIA H100 80GB HBM3\n", + "GPU[3]: NVIDIA H100 80GB HBM3\n", + "GPU[4]: NVIDIA H100 80GB HBM3\n", + "GPU[5]: NVIDIA H100 80GB HBM3\n", + "GPU[6]: NVIDIA H100 80GB HBM3\n", + "GPU[7]: NVIDIA H100 80GB HBM3\n" + ] + } + ], + "source": [ + "# Environment check\n", + "import sys\n", + "\n", + "import torch\n", + "\n", + "print(f\"Python: {sys.version}\")\n", + "print(f\"CUDA available: {torch.cuda.is_available()}\")\n", + "print(f\"Num GPUs: {torch.cuda.device_count()}\")\n", + "\n", + "if torch.cuda.is_available():\n", + " for i in range(torch.cuda.device_count()):\n", + " print(f\"GPU[{i}]: {torch.cuda.get_device_name(i)}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## OpenAI-Compatible Server\n", + "\n", + "Start a local OpenAI-compatible server with TensorRT-LLM via the terminal, within the running docker container.\n", + "\n", + "Ensure that the following commands are executed from the docker terminal.\n", + "\n", + "Start with the GLM 4.7 Flash Yaml here: `examples/auto_deploy/model_registry/configs/glm-4.7-flash.yaml`" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Load the Model\n", + "\n", + "Launch the TensorRT-LLM server with GLM-4.7-Flash:\n", + "\n", + "```shell\n", + "trtllm-serve \"zai-org/GLM-4.7-Flash\" \\\n", + " --host 0.0.0.0 \\\n", + " --port 8000 \\\n", + " --backend _autodeploy \\\n", + " --trust_remote_code \\\n", + " --extra_llm_api_options examples/auto_deploy/model_registry/configs/glm-4.7-flash.yaml\n", + "```" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Your server is now running!" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Use the API\n", + "\n", + "Use the OpenAI-compatible client to send requests to the TensorRT-LLM server." + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "metadata": {}, + "outputs": [], + "source": [ + "from openai import OpenAI\n", + "\n", + "# Setup client\n", + "BASE_URL = \"http://0.0.0.0:8000/v1\"\n", + "API_KEY = \"null\"\n", + "client = OpenAI(base_url=BASE_URL, api_key=API_KEY)\n", + "\n", + "MODEL_ID = \"zai-org/GLM-4.7-Flash\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Chat Completion Example\n", + "==================================================\n", + "Response:\n", + "1. **Analyze the Request:** The user wants to know 15% of 85 and wants to see the reasoning behind the calculation.\n", + "\n", + "2. **Identify the Core Task:** Calculate $15\\% \\times 85$.\n", + "\n", + "3. **Determine the Mathematical Approach:** There are several ways to solve this:\n", + " * *Method 1: Fraction multiplication.* Convert 15% to a fraction ($\\frac{15}{100}$), then multiply by 85.\n", + " * *Method 2: Decimal multiplication.* Convert 15% to a decimal ($0.15$), then multiply by 85.\n", + " * *Method 3: Decomposition (Breaking it down).* $15\\% = 10\\% + 5\\%$.\n", + " * $10\\%$ of $85 = 8.5$\n", + " * $5\\%$ of $85 = \\frac{8.5}{2} = 4.25$\n", + " * Sum: $8.5 + 4.25 = 12.75$\n", + "\n", + "4. **Select the Best Approach for Explanation:** Method 3 is often easiest for a general audience to follow step-by-step because it avoids dealing with decimals until the end or simplifies large multiplications. Method 2 is the most direct standard school method. I will use Method 3 (Splitting 15% into 10% and 5%) as the primary reasoning because it is intuitive, but I might briefly mention the standard formula ($\\frac{\\text{percent}}{100} \\times \\text{number}$).\n", + "\n", + "5. **Execute the Calculation (Method 3):**\n", + " * Step 1: Find 10% of 85.\n", + " * Moving the decimal point one place to the left: 8.5.\n", + " * Step 2: Find 5% of 85.\n", + " * Since 5% is half of 10%, take half of 8.5.\n", + " * $8.5 / 2 = 4.25$.\n", + " * Step 3: Add them together.\n", + " * $8.5 + 4.25$.\n", + " * $8.5 + 4 = 12.5$.\n", + " * $12.5 + 0\n" + ] + } + ], + "source": [ + "# Basic chat completion\n", + "print(\"Chat Completion Example\")\n", + "print(\"=\" * 50)\n", + "\n", + "response = client.chat.completions.create(\n", + " model=MODEL_ID,\n", + " messages=[\n", + " {\"role\": \"system\", \"content\": \"You are a helpful assistant.\"},\n", + " {\"role\": \"user\", \"content\": \"What is 15% of 85? Show your reasoning.\"},\n", + " ],\n", + " temperature=1,\n", + " top_p=0.95,\n", + " max_tokens=512,\n", + ")\n", + "\n", + "print(\"Response:\")\n", + "print(response.choices[0].message.content)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Streaming response:\n", + "==================================================\n", + "1. **Analyze the Request:** The user is asking for the \"first 5 prime numbers\".\n", + "\n", + "2. **Define \"Prime Number\":** A prime number is a natural number greater than 1 that is not a product of two smaller natural numbers. In other words, it has exactly two distinct positive divisors: 1 and itself.\n", + "\n", + "3. **Identify the First Numbers:**\n", + " * Start checking from 1 (exclusive).\n", + " * Check 2: Divisors are 1 and 2. Prime. (1st)\n", + " * Check 3: Divisors are 1 and 3. Prime. (2nd)\n", + " * Check 4: Divisors are 1, 2, 4. Not prime (2 * 2).\n", + " * Check 5: Divisors are 1 and 5. Prime. (3rd)\n", + " * Check 6: Divisors are 1, 2, 3, 6. Not prime.\n", + " * Check 7: Divisors are 1 and 7. Prime. (4th)\n", + " * Check 8: Divisors are 1, 2, 4, 8. Not prime.\n", + " * Check 9: Divisors are 1, 3, 9. Not prime.\n", + " * Check 10: Divisors are 1, 2, 5, 10. Not prime.\n", + " * Check 11: Divisors are 1 and 11. Prime. (5th)\n", + "\n", + "4. **Compile the List:** 2, 3, 5, 7, 11.\n", + "\n", + "5. **Formulate the Output:** Present the list clearly.\n", + "\n", + "6. **Final Review:** Does this answer the user's prompt accurately? Yes.The first 5 prime numbers are **2, 3, 5, 7, and 11**." + ] + } + ], + "source": [ + "# Streaming chat completion\n", + "print(\"Streaming response:\")\n", + "print(\"=\" * 50)\n", + "\n", + "stream = client.chat.completions.create(\n", + " model=MODEL_ID,\n", + " messages=[\n", + " {\"role\": \"system\", \"content\": \"You are a helpful assistant.\"},\n", + " {\"role\": \"user\", \"content\": \"What are the first 5 prime numbers?\"},\n", + " ],\n", + " temperature=0.7,\n", + " max_tokens=1024,\n", + " stream=True,\n", + ")\n", + "\n", + "for chunk in stream:\n", + " if chunk.choices[0].delta.content:\n", + " print(chunk.choices[0].delta.content, end=\"\", flush=True)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Evaluation Parameters\n", + "\n", + "For optimal results, use the following parameters based on your task:\n", + "\n", + "**Default Settings (Most Tasks)**\n", + "- `temperature`: 1.0\n", + "- `top_p`: 0.95\n", + "- `max_tokens`: 131072\n", + "\n", + "**Agentic Tasks (SWE-bench, Terminal Bench)**\n", + "- `temperature`: 0.7\n", + "- `top_p`: 1.0\n", + "- `max_tokens`: 16384\n", + "\n", + "**Deterministic Tasks**\n", + "- `temperature`: 0\n", + "- `max_tokens`: 16384" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Additional Resources\n", + "\n", + "- [TensorRT-LLM Documentation](https://nvidia.github.io/TensorRT-LLM/)\n", + "- [AutoDeploy Guide](https://nvidia.github.io/TensorRT-LLM/torch/auto_deploy/auto-deploy.html)\n", + "- [GLM-4.7-Flash on HuggingFace](https://huggingface.co/zai-org/GLM-4.7-Flash)\n", + "- [Z.ai Discord Community](https://discord.gg/QR7SARHRxK)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.3" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/examples/auto_deploy/model_registry/configs/glm-4.7-flash.yaml b/examples/auto_deploy/model_registry/configs/glm-4.7-flash.yaml new file mode 100644 index 000000000000..fb434e8bb7c0 --- /dev/null +++ b/examples/auto_deploy/model_registry/configs/glm-4.7-flash.yaml @@ -0,0 +1,5 @@ +compile_backend: torch-cudagraph +max_batch_size: 64 +max_seq_len: 4096 +enable_chunked_prefill: true +cuda_graph_batch_sizes: [1, 2, 4, 8, 16, 32, 64] diff --git a/examples/auto_deploy/model_registry/models.yaml b/examples/auto_deploy/model_registry/models.yaml index 879a5eb0c6bb..8ddf66c805fc 100644 --- a/examples/auto_deploy/model_registry/models.yaml +++ b/examples/auto_deploy/model_registry/models.yaml @@ -221,3 +221,5 @@ models: yaml_extra: ['dashboard_default.yaml', 'world_size_8.yaml', 'multimodal.yaml', 'llama4_maverick_lite.yaml'] - name: nvidia/NVIDIA-Nemotron-3-Super-120B-BF16-BF16KV-010726 yaml_extra: ['dashboard_default.yaml', 'world_size_4.yaml','super_v3.yaml'] +- name: zai-org/GLM-4.7-Flash + yaml_extra: ['glm-4.7-flash.yaml'] diff --git a/tensorrt_llm/_torch/auto_deploy/config/default.yaml b/tensorrt_llm/_torch/auto_deploy/config/default.yaml index 948e2f813972..6dd3e617e1d7 100644 --- a/tensorrt_llm/_torch/auto_deploy/config/default.yaml +++ b/tensorrt_llm/_torch/auto_deploy/config/default.yaml @@ -160,15 +160,16 @@ transforms: visualize_namespace: stage: visualize enabled: false - ############################################################################################ + ########################################################################################### # SWITCH TO CACHED+FLATTENED ATTENTION + INITIALIZE CACHES - ############################################################################################ + ########################################################################################### insert_cached_attention: stage: cache_init backend: flashinfer insert_cached_mla_attention: stage: cache_init - backend: MultiHeadLatentAttention + requires_shape_prop: true + backend: flashinfer_mla insert_cached_ssm_attention: stage: cache_init backend: triton_ssm diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/README.md b/tensorrt_llm/_torch/auto_deploy/custom_ops/README.md index 27f18bed1b92..addc6cc22288 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/README.md +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/README.md @@ -5,38 +5,3 @@ All AutoDeploy custom operators follow the following naming convention: `torch.ops.auto_deploy.__` The table below lists the operators ordered by their backend. - -### Available Custom Operators - -| Operator Name | Description | -|--------------|-------------| -| `torch.ops.auto_deploy.flashinfer_attention_mha_with_cache` | FlashInfer attention with KV cache support | -| `torch.ops.auto_deploy.flashinfer_rope` | FlashInfer RoPE (Rotary Position Embedding) implementation | -| `torch.ops.auto_deploy.torch_attention_deepseek_fused_mla` | DeepSeek fused MLA (Multi-head Linear Attention) | -| `torch.ops.auto_deploy.torch_attention_deepseek_mla` | DeepSeek MLA implementation | -| `torch.ops.auto_deploy.torch_attention` | Grouped SDPA implementation with `bsnd` and `bnsd` layout supported | -| `torch.ops.auto_deploy.torch_attention_repeat_kv` | KV repetition for attention | -| `torch.ops.auto_deploy.torch_attention_sdpa` | Standard SDPA implementation | -| `torch.ops.auto_deploy.torch_dist_all_gather` | Distributed all-gather operation (PyTorch backend, demollm mode) | -| `torch.ops.auto_deploy.torch_dist_all_reduce` | Distributed all-reduce operation (PyTorch backend, demollm mode) | -| `torch.ops.auto_deploy.torch_linear_simple` | Simple linear layer implementation | -| `torch.ops.auto_deploy.torch_moe` | Mixture of Experts implementation | -| `torch.ops.auto_deploy.torch_moe_fused` | Fused Mixture of Experts implementation | -| `torch.ops.auto_deploy.torch_quant_fn` | Generic quantization function that scales, rounds, and clamps input values | -| `torch.ops.auto_deploy.torch_quant_nvfp4_linear` | FP4 quantized linear layer | -| `torch.ops.auto_deploy.torch_quant_fp8_linear` | FP8 quantized linear layer | -| `torch.ops.auto_deploy.torch_rope_with_complex_freqs` | RoPE with complex frequencies | -| `torch.ops.auto_deploy.torch_rope_with_explicit_cos_sin` | RoPE with explicit cosine/sine | -| `torch.ops.auto_deploy.torch_rope_with_qk_interleaving` | RoPE with QK interleaving | -| `torch.ops.auto_deploy.triton_attention_fused_flattened_mha_with_cache` | Triton fused flattened MHA with cache | -| `torch.ops.auto_deploy.triton_attention_fused_flattened_mha_with_cache_rope_fusion` | Triton fused flattened MHA with cache and RoPE fusion | -| `torch.ops.auto_deploy.triton_attention_fused_mha_with_cache` | Triton fused MHA with cache | -| `torch.ops.auto_deploy.triton_attention_fused_mha_with_paged_cache` | Triton fused MHA with paged cache | -| `torch.ops.auto_deploy.triton_attention_flattened_mha_with_cache` | Triton flattened MHA with cache | -| `torch.ops.auto_deploy.triton_attention_fused_flattened_mla_with_cache` | Triton fused flattened Multi-head Latent Attention with cache support | -| `torch.ops.auto_deploy.triton_rope_on_flattened_inputs` | Triton RoPE on flattened inputs | -| `torch.ops.auto_deploy.triton_rope_with_input_pos` | Triton RoPE with input positions | -| `torch.ops.auto_deploy.trtllm_moe_fused` | TensorRT LLM fused MoE implementation | -| `torch.ops.auto_deploy.trtllm_dist_all_gather` | Distributed all-gather operation (TRT-LLM backend, MPI mode) | -| `torch.ops.auto_deploy.trtllm_dist_all_reduce` | Distributed all-reduce operation (TRT-LLM backend, MPI mode) | -| `torch.ops.dist.trtllm_fused_allreduce_residual_rmsnorm` | Fused all-reduce + residual add + RMSNorm (TRT-LLM backend, MPI mode) | diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/torch_attention.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/torch_attention.py index da76b1e52e97..8d0d819300e9 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/torch_attention.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/torch_attention.py @@ -19,7 +19,6 @@ from typing import Optional import torch -import torch.nn as nn import torch.nn.functional as F @@ -242,299 +241,3 @@ def torch_attention_fake( layout: str = "bnsd", ): return query.new_empty(*query.shape[:-1], value.shape[-1]).contiguous() - - -def update_kv_cache( - key_states: torch.Tensor, - value_states: torch.Tensor, - k_cache: torch.Tensor, - v_cache: torch.Tensor, - seq_len: torch.Tensor, # metadata - input_pos: torch.Tensor, # metadata - slot_idx: torch.Tensor, - seq_start: torch.Tensor, -) -> torch.Tensor: - """ - Reference implementation for update kv cache function. Assumes KV cache layout to be [B,S,N,D]. - This function can be used to build reference attention implementations that use KV cache. - """ - - for idx in range(seq_len.shape[0]): - k_cache[slot_idx[idx], input_pos[idx] : input_pos[idx] + seq_len[idx], :, :] = key_states[ - seq_start[idx] : seq_start[idx] + seq_len[idx], ... - ] - v_cache[slot_idx[idx], input_pos[idx] : input_pos[idx] + seq_len[idx], :, :] = value_states[ - seq_start[idx] : seq_start[idx] + seq_len[idx], ... - ] - - -@torch.library.custom_op("auto_deploy::torch_attention_fused_mla_ref", mutates_args=()) -def fused_mla_ref( - q_nope: torch.Tensor, - q_pe: torch.Tensor, - kv: torch.Tensor, - k_pe: torch.Tensor, - seq_len: torch.Tensor, # metadata - input_pos: torch.Tensor, # metadata - cache_loc: torch.Tensor, - seq_start: torch.Tensor, - k_cache: torch.Tensor, # caches - v_cache: torch.Tensor, # caches - freqs_cis: torch.Tensor, - softmax_scale: Optional[float] = None, -) -> torch.Tensor: - """ - Reference implementation for Fused MLA with KV cache support. - This implementation flattens the inputs and can be used as a reference to debug the triton kernels. - """ - # Compute parameters - bs, num_heads, q_len, qk_nope_head_dim = q_nope.shape - qk_rope_head_dim = q_pe.shape[-1] - v_head_dim = kv.shape[-1] - qk_nope_head_dim - q_head_dim = qk_nope_head_dim + qk_rope_head_dim - - # Flatten inputs - bs_view = (bs * q_len,) - - k_nope, value_states = torch.split(kv, [qk_nope_head_dim, v_head_dim], dim=-1) - - q_nope = q_nope.transpose(1, 2).view(*bs_view, num_heads, qk_nope_head_dim).contiguous() - q_pe = q_pe.transpose(1, 2).clone().view(*bs_view, num_heads, qk_rope_head_dim).contiguous() - k_nope = k_nope.clone().transpose(1, 2).view(*bs_view, num_heads, qk_nope_head_dim).contiguous() - k_pe = k_pe.clone().transpose(1, 2).view(*bs_view, -1, qk_rope_head_dim).contiguous() - value_states = value_states.transpose(1, 2).view(*bs_view, -1, v_head_dim).contiguous() - - if freqs_cis is not None: - cos_base = freqs_cis[0, ...] - sin_base = freqs_cis[1, ...] - for i in range(seq_len.shape[0]): - start = seq_start[i] - length = seq_len[i] - if q_len == 1: - idx = (input_pos[i] + length - 1).item() - pos_ids = torch.tensor(idx, device=cos_base.device) - else: - pos_ids = torch.arange(input_pos[i], input_pos[i] + length, device=cos_base.device) - - cos = cos_base[pos_ids] # [..., 1, head_dim] - sin = sin_base[pos_ids] - q_slice = q_pe[start : start + length] - k_slice = k_pe[start : start + length] - - q_rot, k_rot = torch.ops.auto_deploy.torch_rope_with_qk_interleaving( - q_slice, - k_slice, - cos, - sin, - -2, - ) - - q_pe[start : start + length] = q_rot - k_pe[start : start + length] = k_rot - - query_states = k_pe.new_empty(*bs_view, num_heads, q_head_dim) # [b*s,n,d] - query_states[..., :qk_nope_head_dim] = q_nope - query_states[..., qk_nope_head_dim:] = q_pe - - key_states = k_pe.new_empty(*bs_view, num_heads, q_head_dim) - key_states[..., :qk_nope_head_dim] = k_nope - key_states[..., qk_nope_head_dim:] = k_pe - - # Update KV cache - update_kv_cache( - key_states, value_states, k_cache, v_cache, seq_len, input_pos, cache_loc, seq_start - ) - - # Compute attention - attn_outputs = [] - for idx in range(seq_len.shape[0]): - # Get inputs from KV cache - k = k_cache[cache_loc[idx], : input_pos[idx] + seq_len[idx], :, :] # [kv_seq_len, n, d] - v = v_cache[cache_loc[idx], : input_pos[idx] + seq_len[idx], :, :] # [kv_seq_len, n, d] - # Generate attention mask - if q_len == 1: - # Generate phase - single token attention mask - attn_mask = torch.zeros( - 1, input_pos[idx] + 1, device=query_states.device, dtype=query_states.dtype - ) - else: - # Context phase - causal attention mask - temp_mask = torch.ones( - seq_len[idx], - input_pos[idx] + seq_len[idx], - dtype=torch.bool, - device=query_states.device, - ).tril(diagonal=0) - attn_bias = torch.zeros( - seq_len[idx], - input_pos[idx] + seq_len[idx], - device=query_states.device, - dtype=query_states.dtype, - ) - attn_mask = attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf")).to( - query_states.device - ) - - # Compute attention weights - attn_weights = ( - torch.matmul( - query_states[seq_start[idx] : seq_start[idx] + seq_len[idx], :, :].transpose(0, 1), - k.transpose(0, 1).transpose(1, 2), - ) - * 1 - / math.sqrt(query_states.size(-1)) - ) - attn_weights = attn_weights + attn_mask - # upcast attention to fp32 - attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to( - query_states.dtype - ) - attn_weights = torch.nn.functional.dropout(attn_weights, p=0.0, training=False) - attn_output = torch.matmul(attn_weights, v.transpose(0, 1)) - attn_outputs.append(attn_output) - - if q_len == 1: - attn_output = torch.stack(attn_outputs) - else: - attn_output = torch.cat(attn_outputs, dim=-2).unsqueeze(0) - - return attn_output - - -@fused_mla_ref.register_fake -def fused_mla_ref_fake( - q_nope: torch.Tensor, - q_pe: torch.Tensor, - kv: torch.Tensor, - k_pe: torch.Tensor, - seq_len: torch.Tensor, # metadata - input_pos: torch.Tensor, # metadata - cache_loc: torch.Tensor, - seq_start: torch.Tensor, - k_cache: torch.Tensor, # caches - v_cache: torch.Tensor, # caches - freqs_cis: torch.Tensor, - softmax_scale: Optional[float] = None, -): - """Fake Fused MLA+Rope with KV cache support.""" - v_head_dim = kv.shape[-1] - q_nope.shape[-1] - return torch.empty_like(kv[..., -v_head_dim:]) - - -@torch.library.custom_op("auto_deploy::torch_attention_deepseek_fused_mla", mutates_args=()) -def fused_mla( - q_nope: torch.Tensor, - q_pe: torch.Tensor, - kv: torch.Tensor, - k_pe: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - position_ids: torch.Tensor, - attention_mask: torch.Tensor, - softmax_scale: float, -) -> torch.Tensor: - """MultiHeadLatentAttention as implemented in DeepSeekV3Attention. This does not capture KV cache use/update.""" - # Did not implement KV cache logic, since we would be writing our own custom op and inserting KV caches later - # Compute parameters - bs, num_heads, q_len, qk_nope_head_dim = q_nope.shape - qk_rope_head_dim = q_pe.shape[-1] - v_head_dim = kv.shape[-1] - qk_nope_head_dim - q_head_dim = qk_nope_head_dim + qk_rope_head_dim - - k_nope, value_states = torch.split(kv, [qk_nope_head_dim, v_head_dim], dim=-1) - kv_seq_len = value_states.shape[-2] - - cos = cos[position_ids] - sin = sin[position_ids] - q_pe, k_pe = torch.ops.auto_deploy.torch_rope_with_qk_interleaving(q_pe, k_pe, cos, sin) - - query_states = k_pe.new_empty(bs, num_heads, q_len, q_head_dim) - query_states[:, :, :, :qk_nope_head_dim] = q_nope - query_states[:, :, :, qk_nope_head_dim:] = q_pe - - key_states = k_pe.new_empty(bs, num_heads, q_len, q_head_dim) - key_states[:, :, :, :qk_nope_head_dim] = k_nope - key_states[:, :, :, qk_nope_head_dim:] = k_pe - - # Use old logic - attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) * softmax_scale - - if attn_weights.size() != (bs, num_heads, q_len, kv_seq_len): - raise ValueError( - f"Attention weights should be of size {(bs, num_heads, q_len, kv_seq_len)}, but is" - f" {attn_weights.size()}" - ) - assert attention_mask is not None - if attention_mask is not None: - if attention_mask.size() != (bs, 1, q_len, kv_seq_len): - raise ValueError( - f"Attention mask should be of size {(bs, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}" - ) - attn_weights = attn_weights + attention_mask - - # upcast attention to fp32 - attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to( - query_states.dtype - ) - attn_weights = nn.functional.dropout(attn_weights, p=0.0, training=False) - attn_output = torch.matmul(attn_weights, value_states) - - if attn_output.size() != (bs, num_heads, q_len, v_head_dim): - raise ValueError( - f"`attn_output` should be of size {(bs, num_heads, q_len, v_head_dim)}, but is" - f" {attn_output.size()}" - ) - - # We do not return attn_weights along with attn_output - return attn_output - - -@fused_mla.register_fake -def fused_mla( - q_nope: torch.Tensor, - q_pe: torch.Tensor, - kv: torch.Tensor, - k_pe: torch.Tensor, - cos: torch.Tensor, - sin: torch.Tensor, - position_ids: torch.Tensor, - attention_mask: torch.Tensor, - softmax_scale: float, -) -> torch.Tensor: - v_head_dim = kv.shape[-1] - q_nope.shape[-1] - return torch.empty_like(kv[..., -v_head_dim:]) - - -@torch.library.custom_op("auto_deploy::torch_attention_deepseek_mla", mutates_args=()) -def mla( - q_nope: torch.Tensor, # Down projected q_nope - q_pe: torch.Tensor, # q_pe after applying rope - kv: torch.Tensor, # compressed kv after passing through layernorm - pe: torch.Tensor, # k_pe after applying rope - attention_mask: torch.Tensor, # attention mask - softmax_scale: float, # softmax scale -) -> torch.Tensor: - """ - Reference implementation for MLA style attention that handles compressed kv. - """ - scores = ( - torch.einsum("bhsc,btc->bsht", q_nope, kv) + torch.einsum("bhsr,btr->bsht", q_pe, pe) - ) * softmax_scale - if attention_mask is not None: - scores += attention_mask.unsqueeze(1) - scores = scores.softmax(dim=-1, dtype=torch.float32).type_as(q_nope) - attn_output = torch.einsum("bsht,btc->bshc", scores, kv) - return attn_output - - -@mla.register_fake -def mla( - q_nope: torch.Tensor, # Down projected q_nope - q_pe: torch.Tensor, # q_pe after applying rope - kv: torch.Tensor, # compressed kv after passing through layernorm - k_pe: torch.Tensor, # k_pe after applying rope - attention_mask: torch.Tensor, # attention mask - softmax_scale: float, # softmax scale -) -> torch.Tensor: - """MLA style attention that handles compressed kv.""" - return torch.empty_like(q_nope) diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/torch_backend_attention.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/torch_backend_attention.py index a8f68574c5f4..36c8d54d4e4d 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/torch_backend_attention.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/torch_backend_attention.py @@ -35,7 +35,31 @@ ResourceHandlerDict, UnpagedResourceHandler, ) -from .torch_attention import repeat_kv, update_kv_cache +from .torch_attention import repeat_kv + + +def _update_kv_cache( + key_states: torch.Tensor, + value_states: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + seq_len: torch.Tensor, # metadata + input_pos: torch.Tensor, # metadata + slot_idx: torch.Tensor, + seq_start: torch.Tensor, +) -> torch.Tensor: + """ + Reference implementation for update kv cache function. Assumes KV cache layout to be [B,S,N,D]. + This function can be used to build reference attention implementations that use KV cache. + """ + + for idx in range(seq_len.shape[0]): + k_cache[slot_idx[idx], input_pos[idx] : input_pos[idx] + seq_len[idx], :, :] = key_states[ + seq_start[idx] : seq_start[idx] + seq_len[idx], ... + ] + v_cache[slot_idx[idx], input_pos[idx] : input_pos[idx] + seq_len[idx], :, :] = value_states[ + seq_start[idx] : seq_start[idx] + seq_len[idx], ... + ] def _apply_logit_softcapping(attn_scores: torch.Tensor, logit_cap: Optional[float]) -> torch.Tensor: @@ -149,7 +173,7 @@ def _torch_context_mha( ) -> None: """Context attention (multiple tokens, potentially multiple sequences) using existing torch functions.""" # Update KV cache first using existing function - update_kv_cache(k, v, k_cache, v_cache, seq_len, input_pos, slot_idx, seq_start) + _update_kv_cache(k, v, k_cache, v_cache, seq_len, input_pos, slot_idx, seq_start) # Compute attention for each sequence attn_outputs = [] diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/__init__.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/__init__.py index b2c4737b6786..261617de1b71 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/__init__.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/__init__.py @@ -1,24 +1,21 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. +"""MLA (Multi-head Latent Attention) custom ops. -"""Multi-head Latent Attention operations. - -This module provides Multi-head Latent Attention (MLA) implementations: -- mla: MLA operations and attention descriptor +Exports: +- TorchBackendMLAAttention: Attention descriptor for MLA (registered as "torch_mla") +- FlashInferMLAAttention: Attention descriptor for FlashInfer MLA (registered as "flashinfer_mla") +- torch_mla: Source op for MLA attention +- torch_backend_mla_with_cache: Cached backend op with FlashInfer-compatible cache +- flashinfer_mla_with_cache: Cached backend op using FlashInfer MLA kernels """ +from .flashinfer_mla import FlashInferMLAAttention, flashinfer_mla_with_cache +from .torch_backend_mla import TorchBackendMLAAttention, torch_backend_mla_with_cache +from .torch_mla import torch_mla + __all__ = [ - "mla", + "TorchBackendMLAAttention", + "FlashInferMLAAttention", + "torch_mla", + "torch_backend_mla_with_cache", + "flashinfer_mla_with_cache", ] diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashinfer_mla.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashinfer_mla.py new file mode 100644 index 000000000000..3f9b353187bc --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashinfer_mla.py @@ -0,0 +1,1037 @@ +"""FlashInfer-based MLA (Multi-head Latent Attention) backend with paged caching. + +This module provides: +- FlashInferMLAAttention: attention descriptor using FlashInfer MLA kernels +- flashinfer_mla_with_cache: cached backend op with paged KV cache + +FlashInfer MLA uses: +- Regular prefill (input_pos == 0): BatchPrefillWithRaggedKVCacheWrapper with expanded K, V +- Chunked prefill (input_pos > 0): BatchMLAPagedAttentionWrapper with matrix absorption +- Decode: BatchMLAPagedAttentionWrapper with paged compressed KV cache + +FlashInfer MLA Cache Layout (two separate caches): + ckv_cache: [num_pages, page_size, kv_lora_rank] + kpe_cache: [num_pages, page_size, qk_rope_head_dim] + - No num_heads dimension (MLA-specific optimization) + +Reference: https://docs.flashinfer.ai/api/mla.html +""" + +import math +from dataclasses import dataclass, fields +from math import prod +from typing import Dict, List, Literal, Optional, Tuple + +import flashinfer +import torch +from torch._ops import OpOverloadPacket +from torch._subclasses import FakeTensor +from torch.fx import Node + +from .....llmapi.llm_args import KvCacheConfig +from ...utils.cuda_graph import cuda_graph_state +from ..attention_interface import ( + AttentionDescriptor, + AttentionLayout, + AttentionRegistry, + Constant, + MHACallable, + PrepareMetadataCallable, + PrepareMetadataHostCallable, + ResourceHandler, + ResourceHandlerDict, + SequenceInfo, +) + + +@dataclass +class MLADecodePlanParams: + """Parameters that affect the FlashInfer MLA decode execution plan.""" + + num_heads: int + kv_lora_rank: int # head_dim_ckv + qk_rope_head_dim: int # head_dim_kpe + qk_nope_head_dim: int + v_head_dim: int + num_seq: int + page_size: int + q_dtype: torch.dtype + kv_dtype: torch.dtype + sm_scale: Optional[float] = None + + def __hash__(self): + """Convert all fields to a string representation and concatenate them.""" + return hash("_".join([str(getattr(self, f.name)) for f in fields(self)])) + + +@dataclass +class MLAPrefillPlanParams: + """Parameters that affect the FlashInfer MLA prefill execution plan.""" + + num_heads: int + num_kv_heads: int # For MLA with expanded KV, same as num_heads + head_dim_qk: int # qk_nope_head_dim + qk_rope_head_dim + head_dim_vo: int # v_head_dim (value/output head dimension) + num_seq: int + q_dtype: torch.dtype + kv_dtype: torch.dtype + sm_scale: Optional[float] = None + + def __hash__(self): + """Convert all fields to a string representation and concatenate them.""" + return hash("_".join([str(getattr(self, f.name)) for f in fields(self)])) + + +class _FlashInferMLAPlanner: + """A class interface to handle FlashInfer MLA-related planning/wrapping operations. + + For MLA attention: + - Regular prefill uses BatchPrefillWithRaggedKVCacheWrapper with expanded K, V tensors + - Chunked prefill uses BatchMLAPagedAttentionWrapper with matrix absorption (same as decode) + - Decode uses BatchMLAPagedAttentionWrapper with paged compressed KV cache + """ + + workspace_buffer: Optional[torch.Tensor] + prefill_wrapper: Optional[flashinfer.prefill.BatchPrefillWithRaggedKVCacheWrapper] + decode_wrapper: Optional["flashinfer.mla.BatchMLAPagedAttentionWrapper"] + # Separate wrapper for chunked/incremental prefill (uses same kernel as decode but different planning) + chunked_prefill_wrapper: Optional["flashinfer.mla.BatchMLAPagedAttentionWrapper"] + cached_cuda_graph_decode_wrappers: Dict[ + MLADecodePlanParams, "flashinfer.mla.BatchMLAPagedAttentionWrapper" + ] + plan_params_prefill: Optional[MLAPrefillPlanParams] + plan_params_decode: Optional[MLADecodePlanParams] + plan_params_chunked_prefill: Optional[MLADecodePlanParams] + kv_layout: Literal["NHD", "HND"] = "NHD" + + def __init__(self): + self.workspace_buffer = None + self.prefill_wrapper = None + self.decode_wrapper = None + self.chunked_prefill_wrapper = None + self.cached_cuda_graph_decode_wrappers = {} + self.plan_params_prefill = None + self.plan_params_decode = None + self.plan_params_chunked_prefill = None + + def _init_decode_wrapper( + self, + use_cuda_graph: bool = False, + qo_indptr: Optional[torch.Tensor] = None, + kv_indptr: Optional[torch.Tensor] = None, + kv_indices: Optional[torch.Tensor] = None, + kv_len_arr: Optional[torch.Tensor] = None, + ): + assert self.workspace_buffer is not None + if use_cuda_graph: + return flashinfer.mla.BatchMLAPagedAttentionWrapper( + self.workspace_buffer, + use_cuda_graph=True, + qo_indptr=qo_indptr, + kv_indptr=kv_indptr, + kv_indices=kv_indices, + kv_len_arr=kv_len_arr, + ) + else: + return flashinfer.mla.BatchMLAPagedAttentionWrapper( + self.workspace_buffer, + use_cuda_graph=False, + ) + + def reset(self, device: torch.device) -> None: + self.plan_params_prefill = None + self.plan_params_decode = None + self.plan_params_chunked_prefill = None + + if isinstance(self.workspace_buffer, torch.Tensor): + return + + self.__init__() # reset all state + + # NOTE: avoid OOM for many cudagraphs + self.workspace_buffer = torch.empty(320 * 1024 * 1024, device=device, dtype=torch.uint8) + + # Prefill uses BatchPrefillWithRaggedKVCacheWrapper with expanded K, V + self.prefill_wrapper = flashinfer.prefill.BatchPrefillWithRaggedKVCacheWrapper( + self.workspace_buffer, + self.kv_layout, + ) + # Decode uses BatchMLAPagedAttentionWrapper with paged compressed KV cache + self.decode_wrapper = self._init_decode_wrapper() + # Chunked prefill uses same kernel as decode but with variable-length queries + self.chunked_prefill_wrapper = self._init_decode_wrapper() + + def plan_prefill( + self, + qo_indptr_host: torch.Tensor, + kv_indptr_host: torch.Tensor, + plan_params: MLAPrefillPlanParams, + custom_mask: Optional[torch.Tensor] = None, + ) -> flashinfer.prefill.BatchPrefillWithRaggedKVCacheWrapper: + """Plan prefill using BatchPrefillWithRaggedKVCacheWrapper. + + For MLA prefill, we expand compressed_kv to get full K, V tensors + and use standard ragged KV cache attention. + + Args: + qo_indptr_host: Cumulative query/output lengths on host. + kv_indptr_host: Cumulative key/value lengths on host. + plan_params: Parameters for planning (hashable, no tensors). + custom_mask: Optional custom attention mask tensor on CUDA. + """ + # Always re-plan when custom_mask is provided since the mask changes per batch + needs_replan = (plan_params != self.plan_params_prefill) or (custom_mask is not None) + + if needs_replan: + # Use causal=True as fallback when no custom_mask is provided + use_causal = custom_mask is None + + # When using custom_mask, FlashInfer expects indptr tensors on CUDA + if custom_mask is not None: + qo_indptr = qo_indptr_host.to(custom_mask.device) + kv_indptr = kv_indptr_host.to(custom_mask.device) + else: + qo_indptr = qo_indptr_host + kv_indptr = kv_indptr_host + + self.prefill_wrapper.plan( + qo_indptr, + kv_indptr, + plan_params.num_heads, + plan_params.num_kv_heads, + plan_params.head_dim_qk, + head_dim_vo=plan_params.head_dim_vo, + use_fp16_qk_reduction=False, + causal=use_causal, + q_data_type=plan_params.q_dtype, + kv_data_type=plan_params.kv_dtype, + sm_scale=plan_params.sm_scale, + custom_mask=custom_mask, + ) + self.plan_params_prefill = plan_params + + return self.prefill_wrapper + + def _plan_mla_wrapper( + self, + wrapper: "flashinfer.mla.BatchMLAPagedAttentionWrapper", + qo_indptr: torch.Tensor, + kv_page_indptr: torch.Tensor, + kv_page_indices: torch.Tensor, + kv_last_page_len: torch.Tensor, + plan_params: MLADecodePlanParams, + ): + """Helper to plan a BatchMLAPagedAttentionWrapper.""" + # Compute actual KV lengths from paging metadata: + # kv_len = (num_pages - 1) * page_size + last_page_len + num_pages_per_seq = kv_page_indptr[1:] - kv_page_indptr[:-1] + kv_len_arr = (num_pages_per_seq - 1) * plan_params.page_size + kv_last_page_len + wrapper.plan( + qo_indptr, + kv_page_indptr, + kv_page_indices, + kv_len_arr, + plan_params.num_heads, + plan_params.kv_lora_rank, # head_dim_ckv + plan_params.qk_rope_head_dim, # head_dim_kpe + plan_params.page_size, + causal=True, + q_data_type=plan_params.q_dtype, + kv_data_type=plan_params.kv_dtype, + sm_scale=plan_params.sm_scale, + ) + + def plan_decode( + self, + kv_page_indptr: torch.Tensor, + kv_page_indices: torch.Tensor, + kv_last_page_len: torch.Tensor, + plan_params: MLADecodePlanParams, + ) -> "flashinfer.mla.BatchMLAPagedAttentionWrapper": + """Plan decode using BatchMLAPagedAttentionWrapper. + + For MLA decode, we use the paged compressed KV cache with + FlashInfer's optimized MLA kernels. Each sequence generates 1 token. + + Args: + kv_page_indptr: Cumulative page counts [batch_size + 1]. + kv_page_indices: Page indices for the KV cache. + kv_last_page_len: Length of the last page per sequence. + plan_params: Parameters for planning. + """ + # Decode qo_indptr: [0, 1, 2, ..., batch_size] (1 token per sequence) + batch_size = kv_page_indptr.shape[0] - 1 + qo_indptr = torch.arange(batch_size + 1, device=kv_page_indptr.device, dtype=torch.int32) + + # Compute kv_len_arr for CUDA graph wrapper initialization + num_pages_per_seq = kv_page_indptr[1:] - kv_page_indptr[:-1] + kv_len_arr = (num_pages_per_seq - 1) * plan_params.page_size + kv_last_page_len + + # we want to plan during warm-up of cuda graph capture to ensure we have the plan cached + if ( + cuda_graph_state.in_warm_up() + and plan_params not in self.cached_cuda_graph_decode_wrappers + ): + # During CUDA graph capture, the metadata tensors provided by auto-deploy are stable. + # Pass the buffer tensors to the wrapper for use_cuda_graph=True + wrapper = self._init_decode_wrapper( + use_cuda_graph=True, + qo_indptr=qo_indptr, + kv_indptr=kv_page_indptr, + kv_indices=kv_page_indices, + kv_len_arr=kv_len_arr, + ) + self.cached_cuda_graph_decode_wrappers[plan_params] = wrapper + self._plan_mla_wrapper( + wrapper, qo_indptr, kv_page_indptr, kv_page_indices, kv_last_page_len, plan_params + ) + + # check if we are in cuda graph capture and just return the pre-cached decode wrapper + if torch.cuda.is_current_stream_capturing() or cuda_graph_state.in_warm_up(): + wrapper = self.cached_cuda_graph_decode_wrappers[plan_params] + return wrapper + + # Re-plan if plan_params changed + if plan_params != self.plan_params_decode: + self._plan_mla_wrapper( + self.decode_wrapper, + qo_indptr, + kv_page_indptr, + kv_page_indices, + kv_last_page_len, + plan_params, + ) + self.plan_params_decode = plan_params + + return self.decode_wrapper + + def plan_chunked_prefill( + self, + qo_indptr: torch.Tensor, + kv_page_indptr: torch.Tensor, + kv_page_indices: torch.Tensor, + kv_last_page_len: torch.Tensor, + plan_params: MLADecodePlanParams, + ) -> "flashinfer.mla.BatchMLAPagedAttentionWrapper": + """Plan chunked/incremental prefill using BatchMLAPagedAttentionWrapper. + + For chunked prefill (input_pos > 0), we use the same kernel as decode but with + variable-length queries. Each sequence can have multiple tokens. + + Args: + qo_indptr: Cumulative query lengths [batch_size + 1]. + kv_page_indptr: Cumulative page counts [batch_size + 1]. + kv_page_indices: Page indices for the KV cache. + kv_last_page_len: Length of the last page per sequence. + plan_params: Parameters for planning. + """ + # Re-plan if plan_params changed + if plan_params != self.plan_params_chunked_prefill: + self._plan_mla_wrapper( + self.chunked_prefill_wrapper, + qo_indptr, + kv_page_indptr, + kv_page_indices, + kv_last_page_len, + plan_params, + ) + self.plan_params_chunked_prefill = plan_params + + return self.chunked_prefill_wrapper + + def plan_generate_only( + self, + num_seq: int, + cu_num_pages: torch.Tensor, + cache_loc: torch.Tensor, + last_page_len: torch.Tensor, + ): + """Plan decode-only batches for cached CUDA graph wrappers. + + This is called from the host-side preparation function to plan + the decode wrappers for decode-only batches before the actual + attention op is invoked. + + Args: + num_seq: Number of sequences in the decode batch. + cu_num_pages: Cumulative page counts, already sliced to [: num_seq + 1]. + cache_loc: Page indices for the KV cache. + last_page_len: Length of the last page per sequence, already sliced to [:num_seq]. + """ + for plan_params in self.cached_cuda_graph_decode_wrappers: + if plan_params.num_seq == num_seq: + wrapper = self.cached_cuda_graph_decode_wrappers[plan_params] + + # For a pure decode batch, qo_indptr is just [0, 1, 2, ..., batch_size] + qo_indptr = torch.arange(num_seq + 1, device=cu_num_pages.device, dtype=torch.int32) + + # Compute actual KV lengths from paging metadata: + # kv_len = (num_pages - 1) * page_size + last_page_len + num_pages_per_seq = cu_num_pages[1:] - cu_num_pages[:-1] + kv_len_arr = (num_pages_per_seq - 1) * plan_params.page_size + last_page_len + + wrapper.plan( + qo_indptr, + cu_num_pages, # kv_page_indptr + cache_loc, # kv_page_indices + kv_len_arr, + plan_params.num_heads, + plan_params.kv_lora_rank, # head_dim_ckv + plan_params.qk_rope_head_dim, # head_dim_kpe + plan_params.page_size, + causal=True, + q_data_type=plan_params.q_dtype, + kv_data_type=plan_params.kv_dtype, + sm_scale=plan_params.sm_scale, + ) + + +_GlobalFlashInferMLAPlanner = _FlashInferMLAPlanner() + + +def _compute_ragged_causal_mask( + cu_seqlen_host: torch.Tensor, + num_seq: int, + device: torch.device, +) -> torch.Tensor: + """Compute a causal mask for FlashInfer ragged prefill (fully GPU-vectorized). + + For ragged attention with multiple sequences, this function creates a + causal mask where each sequence has its own lower triangular (causal) mask. + The masks are flattened and concatenated into a single 1D tensor. + + This implementation is fully vectorized on GPU with no Python loops. + + Args: + cu_seqlen_host: Cumulative sequence lengths on host [num_seq + 1]. + seq_len[i] = cu_seqlen_host[i+1] - cu_seqlen_host[i] + num_seq: Number of sequences in the batch. + device: Device to create the mask on. + + Returns: + A flattened boolean mask tensor where True = attend, False = mask out. + Shape: (sum(seq_len[i] * seq_len[i] for i in range(num_seq)),) + """ + # Move sequence lengths to GPU for vectorized computation + seq_lens = (cu_seqlen_host[1 : num_seq + 1] - cu_seqlen_host[:num_seq]).to(device) + + # Compute squared lengths (mask size per sequence) and total + sq_lens = seq_lens * seq_lens + total_mask_size = sq_lens.sum().item() + + # Cumulative squared lengths (offsets into flattened mask) + cu_sq_lens = torch.zeros(num_seq + 1, device=device, dtype=seq_lens.dtype) + cu_sq_lens[1:] = torch.cumsum(sq_lens, dim=0) + + # Create position indices for the entire mask + positions = torch.arange(total_mask_size, device=device) + + # Find which sequence each position belongs to using searchsorted + # right=True ensures positions at boundaries map to the next sequence + seq_ids = torch.searchsorted(cu_sq_lens[1:], positions, right=True) + + # Compute local position within each sequence's mask + local_pos = positions - cu_sq_lens[seq_ids] + + # Get sequence length for each position + seq_len_per_pos = seq_lens[seq_ids] + + # Compute row and column within each sequence's 2D mask + row = local_pos // seq_len_per_pos + col = local_pos % seq_len_per_pos + + # Causal mask: attend if row >= col (lower triangular) + return row >= col + + +@torch.library.custom_op("auto_deploy::flashinfer_mla_prepare_metadata", mutates_args=()) +def prepare_flashinfer_mla_metadata( + position_ids: torch.Tensor, + batch_info_host: torch.Tensor, + cu_seqlen: torch.Tensor, + seq_len_with_cache: torch.Tensor, +) -> List[torch.Tensor]: + """Prepare metadata for FlashInfer MLA attention. + + This prepares batch_indices and positions for cache appends, similar to + the standard FlashInfer attention preparation. + """ + # retrieve host-side metadata + num_prefill, num_prefill_tokens, num_decode = batch_info_host.tolist() + num_seq = num_prefill + num_decode + num_tokens = num_prefill_tokens + num_decode + + _GlobalFlashInferMLAPlanner.reset(position_ids.device) + + qo_indptr = cu_seqlen[: num_seq + 1] + + # Compute batch_indices and positions for cache appends + batch_indices, positions = flashinfer.get_batch_indices_positions( + qo_indptr, seq_len_with_cache[:num_seq], num_tokens + ) + + return batch_indices, positions + + +@prepare_flashinfer_mla_metadata.register_fake +def prepare_flashinfer_mla_metadata_fake( + position_ids: torch.Tensor, + batch_info_host: torch.Tensor, + cu_seqlen: torch.Tensor, + seq_len_with_cache: torch.Tensor, +): + num_tokens = position_ids.shape[0] * position_ids.shape[1] + return ( + torch.empty(num_tokens, dtype=torch.int32, device=position_ids.device), # batch_indices + torch.empty(num_tokens, dtype=torch.int32, device=position_ids.device), # positions + ) + + +def prepare_flashinfer_mla_metadata_host( + batch_info_host: torch.Tensor, + cu_num_pages_host: torch.Tensor, + cache_loc_host: torch.Tensor, + last_page_len_host: torch.Tensor, +) -> None: + """Host-side preparation for FlashInfer MLA attention. + + For decode-only batches, this function pre-plans the cached CUDA graph + wrappers to avoid planning during graph capture/replay. + """ + num_prefill, num_prefill_tokens, num_decode = batch_info_host.tolist() + + if num_prefill == 0: + _GlobalFlashInferMLAPlanner.plan_generate_only( + num_decode, + cu_num_pages_host[: num_decode + 1], + cache_loc_host, + last_page_len_host[:num_decode], + ) + + +@torch.library.custom_op("auto_deploy::flashinfer_mla_with_cache", mutates_args=()) +def flashinfer_mla_with_cache( + # 5 tensor args (matching torch_mla source op) + q_nope: torch.Tensor, # [B, S, N, qk_nope_head_dim] + q_pe: torch.Tensor, # [B, S, N, qk_rope_head_dim] + compressed_kv: torch.Tensor, # [B, S, kv_lora_rank] + kpe: torch.Tensor, # [B, S, 1, qk_rope_head_dim] + kv_b_proj_weight: torch.Tensor, # [N * (qk_nope_head_dim + v_head_dim), kv_lora_rank] + # Standard paged metadata + batch_info_host: torch.Tensor, + cu_seqlen_host: torch.Tensor, + cu_num_pages: torch.Tensor, + cu_num_pages_host: torch.Tensor, + cache_loc: torch.Tensor, + last_page_len: torch.Tensor, + last_page_len_host: torch.Tensor, + seq_len_with_cache_host: torch.Tensor, + # Extra FlashInfer metadata + flashinfer_batch_indices: torch.Tensor, + flashinfer_positions: torch.Tensor, + # Paged caches (two separate caches) + ckv_cache: torch.Tensor, # [num_pages, page_size, kv_lora_rank] + kpe_cache: torch.Tensor, # [num_pages, page_size, qk_rope_head_dim] + # Constants + scale: Optional[float], + kv_lora_rank: int, +) -> torch.Tensor: + """FlashInfer MLA attention with paged cache. + + Uses FlashInfer's optimized kernels: + - Prefill: BatchPrefillWithRaggedKVCacheWrapper with expanded K, V tensors + - Decode: BatchMLAPagedAttentionWrapper with paged compressed KV cache + + FlashInfer MLA Cache Layout (two separate caches): + ckv_cache: [num_pages, page_size, kv_lora_rank] + kpe_cache: [num_pages, page_size, qk_rope_head_dim] + + Args: + q_nope: Query non-positional component [B, S, N, qk_nope_head_dim] + q_pe: Query positional component [B, S, N, qk_rope_head_dim] + compressed_kv: Compressed KV latent [B, S, kv_lora_rank] + kpe: Key positional encoding [B, S, 1, qk_rope_head_dim] + kv_b_proj_weight: Projection weight [N * (qk_nope_head_dim + v_head_dim), kv_lora_rank] + (metadata args): Standard paged attention metadata + ckv_cache: Paged cache for compressed KV + kpe_cache: Paged cache for key positional encoding + scale: Softmax scale factor + kv_lora_rank: Rank of compressed KV + + Returns: + Attention output [B, S, N, v_head_dim] + """ + # Get dimensions + b, s = q_nope.shape[:2] + num_heads = q_nope.shape[2] + qk_nope_head_dim = q_nope.shape[3] + qk_rope_head_dim = q_pe.shape[3] + qk_head_dim = qk_nope_head_dim + qk_rope_head_dim + + # Infer v_head_dim from kv_b_proj_weight + out_features = kv_b_proj_weight.shape[0] + kv_head_dim = out_features // num_heads + v_head_dim = kv_head_dim - qk_nope_head_dim + + # Get batch info + num_prefill, num_prefill_tokens, num_decode = batch_info_host.tolist() + num_seq = num_prefill + num_decode + num_total_tokens = num_prefill_tokens + num_decode + + # Set scale + if scale is None: + scale = 1.0 / math.sqrt(qk_head_dim) + + page_size = ckv_cache.shape[1] + + # Flatten inputs to [total_tokens, ...] format + bs = b * s + q_nope_flat = q_nope.contiguous().view(bs, num_heads, qk_nope_head_dim) + q_pe_flat = q_pe.contiguous().view(bs, num_heads, qk_rope_head_dim) + compressed_kv_flat = compressed_kv.contiguous().view(bs, kv_lora_rank) + kpe_flat = kpe.contiguous().view(bs, qk_rope_head_dim) + + # Convert cache dtype if needed + if ckv_cache.dtype == torch.float8_e4m3fn: + compressed_kv_flat = compressed_kv_flat.to(torch.float8_e4m3fn) + kpe_flat = kpe_flat.to(torch.float8_e4m3fn) + + # Append to paged cache using FlashInfer's append function + # Note: caches are guaranteed contiguous by CachedSequenceInterface._create_kv_cache_manager + flashinfer.page.append_paged_mla_kv_cache( + compressed_kv_flat, + kpe_flat, + flashinfer_batch_indices, + flashinfer_positions, + ckv_cache, + kpe_cache, + cache_loc, + cu_num_pages[: num_seq + 1], + last_page_len[:num_seq], + ) + + # Pre-allocate output + if num_prefill > 0 and num_decode > 0: + y = torch.empty(bs, num_heads, v_head_dim, dtype=q_nope.dtype, device=q_nope.device) + else: + y = None + + # ========================================================================= + # PREFILL phase: Use BatchPrefillWithRaggedKVCacheWrapper for regular prefill + # or BatchMLAPagedAttentionWrapper for chunked prefill + # ========================================================================= + if num_prefill > 0: + q_nope_prefill = q_nope_flat[:num_prefill_tokens] + q_pe_prefill = q_pe_flat[:num_prefill_tokens] + compressed_kv_prefill = compressed_kv_flat[:num_prefill_tokens] + kpe_prefill = kpe_flat[:num_prefill_tokens] + + # Check if any prefill sequence has cached tokens (chunked prefill) + # seq_len_with_cache > current_seq_len means there are cached tokens + q_lens = cu_seqlen_host[1 : num_prefill + 1] - cu_seqlen_host[:num_prefill] + kv_lens = seq_len_with_cache_host[:num_prefill] + is_chunked_prefill = (kv_lens > q_lens).any().item() + + if is_chunked_prefill: + # ================================================================= + # CHUNKED PREFILL: Use BatchMLAPagedAttentionWrapper with absorption + # Same approach as decode, but with variable-length Q sequences + # ================================================================= + + # Extract W_kn and W_v from kv_b_proj_weight + # kv_b_proj_weight: [N * (qk_nope_head_dim + v_head_dim), kv_lora_rank] + # Reshape to [N, qk_nope_head_dim + v_head_dim, kv_lora_rank] + kv_b_proj_reshaped = kv_b_proj_weight.view( + num_heads, qk_nope_head_dim + v_head_dim, kv_lora_rank + ) + # W_kn: [N, qk_nope_head_dim, kv_lora_rank] + w_kn = kv_b_proj_reshaped[:, :qk_nope_head_dim, :] + # W_v: [N, v_head_dim, kv_lora_rank] + w_v = kv_b_proj_reshaped[:, qk_nope_head_dim:, :] + + # Absorb W_kn into q_nope: + # q_nope_prefill: [num_prefill_tokens, N, qk_nope_head_dim] + # w_kn: [N, qk_nope_head_dim, kv_lora_rank] + # q_nope_absorbed: [num_prefill_tokens, N, kv_lora_rank] + q_nope_absorbed = torch.einsum("bnd,ndk->bnk", q_nope_prefill, w_kn).contiguous() + + # Build qo_indptr for variable-length prefill sequences + qo_indptr = cu_seqlen_host[: num_prefill + 1].to( + device=cu_num_pages.device, dtype=torch.int32 + ) + + pp_chunked = MLADecodePlanParams( + num_heads=num_heads, + kv_lora_rank=kv_lora_rank, + qk_rope_head_dim=qk_rope_head_dim, + qk_nope_head_dim=qk_nope_head_dim, + v_head_dim=v_head_dim, + num_seq=num_prefill, + page_size=page_size, + q_dtype=q_nope.dtype, + kv_dtype=ckv_cache.dtype, + sm_scale=scale, + ) + + wrapper_chunked = _GlobalFlashInferMLAPlanner.plan_chunked_prefill( + qo_indptr=qo_indptr, + kv_page_indptr=cu_num_pages[: num_prefill + 1], + kv_page_indices=cache_loc, + kv_last_page_len=last_page_len[:num_prefill], + plan_params=pp_chunked, + ) + + # Run paged MLA attention in compressed space + y_prefill_compressed = wrapper_chunked.run( + q_nope_absorbed, + q_pe_prefill, + ckv_cache, + kpe_cache, + ) + + # Project output back from latent space to v_head_dim + # y_prefill_compressed: [num_prefill_tokens, N, kv_lora_rank] + # w_v: [N, v_head_dim, kv_lora_rank] + # y_prefill: [num_prefill_tokens, N, v_head_dim] + y_prefill = torch.einsum("bnk,nvk->bnv", y_prefill_compressed, w_v) + + else: + # ================================================================= + # REGULAR PREFILL: Use BatchPrefillWithRaggedKVCacheWrapper + # Expand compressed_kv to K, V and use ragged attention + # ================================================================= + + # Expand compressed_kv using kv_b_proj_weight to get k_nope and v + # compressed_kv: [num_prefill_tokens, kv_lora_rank] + # kv_b_proj_weight: [N * (qk_nope_head_dim + v_head_dim), kv_lora_rank] + # kv_expanded: [num_prefill_tokens, N * (qk_nope_head_dim + v_head_dim)] + kv_expanded = torch.matmul(compressed_kv_prefill, kv_b_proj_weight.t()) + kv_expanded = kv_expanded.view( + num_prefill_tokens, num_heads, qk_nope_head_dim + v_head_dim + ) + + # Split into k_nope and v + k_nope_prefill = kv_expanded[:, :, :qk_nope_head_dim] # [tokens, N, qk_nope_head_dim] + v_prefill = kv_expanded[:, :, qk_nope_head_dim:].contiguous() # [tokens, N, v_head_dim] + + # Expand kpe to all heads: [tokens, qk_rope_head_dim] -> [tokens, N, qk_rope_head_dim] + kpe_expanded = kpe_prefill.unsqueeze(1).expand(-1, num_heads, -1).contiguous() + + # Concatenate to form full Q and K + # Q: [tokens, N, qk_head_dim] + q_prefill = torch.cat([q_nope_prefill, q_pe_prefill], dim=-1).contiguous() + # K: [tokens, N, qk_head_dim] + k_prefill = torch.cat([k_nope_prefill, kpe_expanded], dim=-1).contiguous() + + # Compute the causal mask for ragged prefill + custom_mask = _compute_ragged_causal_mask( + cu_seqlen_host=cu_seqlen_host, + num_seq=num_prefill, + device=q_nope.device, + ) + + pp_prefill = MLAPrefillPlanParams( + num_heads=num_heads, + num_kv_heads=num_heads, # For MLA with expanded KV, same as num_heads + head_dim_qk=qk_head_dim, + head_dim_vo=v_head_dim, + num_seq=num_prefill, + q_dtype=q_nope.dtype, + kv_dtype=k_prefill.dtype, + sm_scale=scale, + ) + + wrapper_prefill = _GlobalFlashInferMLAPlanner.plan_prefill( + qo_indptr_host=cu_seqlen_host[: num_prefill + 1], + kv_indptr_host=cu_seqlen_host[: num_prefill + 1], # Same as qo for self-attention + plan_params=pp_prefill, + custom_mask=custom_mask, + ) + + y_prefill = wrapper_prefill.run( + q_prefill, + k_prefill, + v_prefill, + ) + + if y is not None: + y[:num_prefill_tokens] = y_prefill + else: + y = y_prefill + + # ========================================================================= + # DECODE phase: Use BatchMLAPagedAttentionWrapper with paged compressed KV + # ========================================================================= + if num_decode > 0: + q_nope_decode = q_nope_flat[num_prefill_tokens:num_total_tokens].contiguous() + q_pe_decode = q_pe_flat[num_prefill_tokens:num_total_tokens].contiguous() + + # FlashInfer MLA operates in the compressed latent space. + # We need to: + # 1. Absorb W_kn (K-nope projection) into q_nope + # 2. Run attention in compressed space + # 3. Project output back using W_v + + # Extract W_kn and W_v from kv_b_proj_weight + # kv_b_proj_weight: [N * (qk_nope_head_dim + v_head_dim), kv_lora_rank] + # Reshape to [N, qk_nope_head_dim + v_head_dim, kv_lora_rank] + kv_b_proj_reshaped = kv_b_proj_weight.view( + num_heads, qk_nope_head_dim + v_head_dim, kv_lora_rank + ) + # W_kn: [N, qk_nope_head_dim, kv_lora_rank] + w_kn = kv_b_proj_reshaped[:, :qk_nope_head_dim, :] + # W_v: [N, v_head_dim, kv_lora_rank] + w_v = kv_b_proj_reshaped[:, qk_nope_head_dim:, :] + + # Absorb W_kn into q_nope: + # q_nope_decode: [num_decode, N, qk_nope_head_dim] + # w_kn: [N, qk_nope_head_dim, kv_lora_rank] + # q_nope_absorbed: [num_decode, N, kv_lora_rank] + q_nope_absorbed = torch.einsum("bnd,ndk->bnk", q_nope_decode, w_kn).contiguous() + + pp_decode = MLADecodePlanParams( + num_heads=num_heads, + kv_lora_rank=kv_lora_rank, + qk_rope_head_dim=qk_rope_head_dim, + qk_nope_head_dim=qk_nope_head_dim, + v_head_dim=v_head_dim, + num_seq=num_decode, + page_size=page_size, + q_dtype=q_nope.dtype, + kv_dtype=ckv_cache.dtype, + sm_scale=scale, + ) + + wrapper_decode = _GlobalFlashInferMLAPlanner.plan_decode( + kv_page_indptr=cu_num_pages[num_prefill : num_seq + 1], + kv_page_indices=cache_loc, + kv_last_page_len=last_page_len[num_prefill:num_seq], + plan_params=pp_decode, + ) + + # Run attention in compressed space + # y_decode_compressed: [num_decode, N, kv_lora_rank] + # Note: caches are guaranteed contiguous by CachedSequenceInterface._create_kv_cache_manager + y_decode_compressed = wrapper_decode.run( + q_nope_absorbed, + q_pe_decode, + ckv_cache, + kpe_cache, + ) + + # Project output back from latent space to v_head_dim + # y_decode_compressed: [num_decode, N, kv_lora_rank] + # w_v: [N, v_head_dim, kv_lora_rank] + # y_decode: [num_decode, N, v_head_dim] + y_decode = torch.einsum("bnk,nvk->bnv", y_decode_compressed, w_v) + + if y is not None: + y[num_prefill_tokens:num_total_tokens] = y_decode + else: + y = y_decode + + return y.view(b, s, num_heads, v_head_dim) + + +@flashinfer_mla_with_cache.register_fake +def flashinfer_mla_with_cache_fake( + q_nope: torch.Tensor, + q_pe: torch.Tensor, + compressed_kv: torch.Tensor, + kpe: torch.Tensor, + kv_b_proj_weight: torch.Tensor, + batch_info_host: torch.Tensor, + cu_seqlen_host: torch.Tensor, + cu_num_pages: torch.Tensor, + cu_num_pages_host: torch.Tensor, + cache_loc: torch.Tensor, + last_page_len: torch.Tensor, + last_page_len_host: torch.Tensor, + seq_len_with_cache_host: torch.Tensor, + flashinfer_batch_indices: torch.Tensor, + flashinfer_positions: torch.Tensor, + ckv_cache: torch.Tensor, + kpe_cache: torch.Tensor, + scale: Optional[float], + kv_lora_rank: int, +) -> torch.Tensor: + """Fake implementation for flashinfer_mla_with_cache.""" + num_heads = q_nope.shape[2] + qk_nope_head_dim = q_nope.shape[-1] + out_features = kv_b_proj_weight.shape[0] + kv_head_dim = out_features // num_heads + v_head_dim = kv_head_dim - qk_nope_head_dim + + return q_nope.new_empty( + q_nope.shape[0], q_nope.shape[1], q_nope.shape[2], v_head_dim + ).contiguous() + + +class MLAPagedResourceHandler(ResourceHandler): + """Handler for paged resources in MLA that require per-layer contiguous memory. + + While MLA uses paged caching, the underlying flashinfer MLA kernel uses a uint32_t to track the + strides for the cache. The KVCacheManager will allocate a contiguous tensor for the cache + across all layers with dim 0 representing the layer index. Hence, the per-layer cache has very + large strides to jump between pages which causes overflow in the MLA kernel that uses uint32_t + for strides. + + We use a separate handler for this purpose to avoid registering the cache with the + KVCacheManager and instead rely on local allocation. + """ + + @property + def is_paged(self) -> bool: + """Whether the resource is paged.""" + return True + + def __init__(self, *token_shape: int, dtype: torch.dtype) -> None: + """Initialize the ContiguousPagedResourceHandler. + + Args: + token_shape: The shape of the resource per token. + dtype: The dtype of the resource. + """ + self.token_shape = token_shape + self.dtype = dtype + + def _get_bytes_per_token(self) -> int: + """The size of the resource per token in bytes.""" + return prod(self.token_shape) * self.dtype.itemsize + + def allocate(self, sequence_info: SequenceInfo) -> torch.Tensor: + """Allocate contiguous paged resource. + + Args: + sequence_info: SequenceInfo with device and page information. + + Returns: + Contiguous tensor of shape [num_blocks, tokens_per_block, *token_shape]. + """ + return torch.empty( + sequence_info.num_blocks, + sequence_info.tokens_per_block, + *self.token_shape, + device=sequence_info.device, + dtype=self.dtype, + ) + + +@AttentionRegistry.register("flashinfer_mla") +class FlashInferMLAAttention(AttentionDescriptor): + """Attention descriptor for FlashInfer-based MLA with paged cache. + + This descriptor uses FlashInfer's optimized MLA kernels: + - Source op: torch_mla (same as torch_mla backend) + - Cached op: flashinfer_mla_with_cache with paged cache + + FlashInfer MLA Cache Layout (two separate caches): + ckv_cache: [num_pages, page_size, kv_lora_rank] + kpe_cache: [num_pages, page_size, qk_rope_head_dim] + - No num_heads dimension (MLA-specific optimization) + + Reference: https://docs.flashinfer.ai/api/mla.html + """ + + @classmethod + def _get_planner(cls) -> _FlashInferMLAPlanner: + return _GlobalFlashInferMLAPlanner + + @classmethod + def get_attention_layout(cls) -> AttentionLayout: + """Get the attention layout expected by the backend.""" + return "bsnd" + + @classmethod + def get_num_qkv_args(cls) -> int: + """Get the number of tensor arguments expected by the source op.""" + return 5 # q_nope, q_pe, compressed_kv, kpe, kv_b_proj_weight + + @classmethod + def get_source_attention_op(cls) -> OpOverloadPacket: + """Get the source attention op that we target for replacement.""" + return torch.ops.auto_deploy.torch_mla + + @classmethod + def get_cached_attention_op(cls) -> MHACallable: + """Get the cached attention op.""" + return torch.ops.auto_deploy.flashinfer_mla_with_cache.default + + @classmethod + def get_standard_metadata_args(cls) -> List[str]: + """Get the list of standard metadata arguments for paged attention.""" + return [ + "batch_info_host", + "cu_seqlen_host", + "cu_num_pages", + "cu_num_pages_host", + "cache_loc", + "last_page_len", + "last_page_len_host", + "seq_len_with_cache_host", + ] + + @classmethod + def get_prepare_extra_metadata_info( + cls, any_source_attn_node: Node + ) -> Tuple[Optional[PrepareMetadataCallable], int, List[Constant]]: + """Get the prepare_metadata op for FlashInfer MLA.""" + return (torch.ops.auto_deploy.flashinfer_mla_prepare_metadata.default, 2, []) + + @classmethod + def get_cache_initializers( + cls, source_attn_node: Node, cache_config: KvCacheConfig + ) -> ResourceHandlerDict: + """Get cache initializers using FlashInfer MLA paged cache layout. + + Creates two separate paged caches: + - ckv_cache: [num_pages, page_size, kv_lora_rank] + - kpe_cache: [num_pages, page_size, qk_rope_head_dim] + """ + # Extract dimensions from source node args + # torch_mla signature: q_nope, q_pe, compressed_kv, kpe, kv_b_proj_weight, ... + compressed_kv_fake: FakeTensor = source_attn_node.args[2].meta["val"] + kpe_fake: FakeTensor = source_attn_node.args[3].meta["val"] + + # Get dimensions + # compressed_kv: [B, S, kv_lora_rank] + # kpe: [B, S, 1, qk_rope_head_dim] + kv_lora_rank = compressed_kv_fake.shape[-1] + qk_rope_head_dim = kpe_fake.shape[-1] + + # flashinfer mla requires kv_lora_rank to be 512 and qk_rope_head_dim to be 64 + if kv_lora_rank != 512: + raise ValueError("kv_lora_rank must be 512 for flashinfer_mla") + if qk_rope_head_dim != 64: + raise ValueError("qk_rope_head_dim must be 64 for flashinfer_mla") + + cache_dtype = cls.resolve_cache_dtype(cache_config.dtype, compressed_kv_fake.dtype) + + # FlashInfer MLA uses two separate paged caches with no num_heads dimension + return { + "ckv_cache": MLAPagedResourceHandler( + kv_lora_rank, + dtype=cache_dtype, + ), + "kpe_cache": MLAPagedResourceHandler( + qk_rope_head_dim, + dtype=cache_dtype, + ), + } + + @classmethod + def get_host_prepare_metadata_function(cls) -> Optional[PrepareMetadataHostCallable]: + """Get function for host-side preparation.""" + return prepare_flashinfer_mla_metadata_host + + @classmethod + def get_constants(cls, source_attn_node: Node) -> List[Constant]: + """Get constants to pass to the cached attention op.""" + # Extract kv_lora_rank for cache operations + compressed_kv_fake = source_attn_node.args[2].meta["val"] + kv_lora_rank = compressed_kv_fake.shape[-1] + + # Get scale from kwargs + scale = source_attn_node.kwargs.get("scale", None) + + return [scale, kv_lora_rank] diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/mla.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/mla.py deleted file mode 100644 index f435fc58188e..000000000000 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/mla.py +++ /dev/null @@ -1,280 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Custom ops for MultiHead Latent attention.""" - -import math -from typing import List, Optional, Union - -import torch -from torch._ops import OpOverloadPacket -from torch.fx import Node - -from .....llmapi.llm_args import KvCacheConfig -from ..attention.triton_attention import _decode_attention, _prefill_attention -from ..attention_interface import ( - AttentionDescriptor, - AttentionLayout, - AttentionRegistry, - MHACallable, - ResourceHandlerDict, - UnpagedResourceHandler, -) - -Constant = Union[int, float, str, None] - - -def _precompute_inv_freq( - max_seq_len: int, head_dim: int, rope_theta: float, device: torch.device -) -> torch.Tensor: - inv_freq = 1.0 / ( - rope_theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim) - ) - t = torch.arange(max_seq_len, device=inv_freq.device, dtype=inv_freq.dtype) - - freqs = torch.outer(t, inv_freq.to(t.device)) - # Different from paper, but it uses a different permutation in order to obtain the same calculation - emb = torch.cat((freqs, freqs), dim=-1) - cos_sin_stacked = torch.stack([emb.cos().to(torch.bfloat16), emb.sin().to(torch.bfloat16)]) - return cos_sin_stacked - - -@torch.library.custom_op( - "auto_deploy::triton_attention_fused_flattened_mla_with_cache", mutates_args=() -) -def fused_flattened_mla_with_cache( - # Q, K, V - q_nope: torch.Tensor, - q_pe: torch.Tensor, - kv: torch.Tensor, - k_pe: torch.Tensor, - # STANDARD METADATA - batch_info_host: torch.Tensor, - seq_len: torch.Tensor, - input_pos: torch.Tensor, - cache_loc: torch.Tensor, - cu_seqlen: torch.Tensor, - # EXTRA METADATA - # - # CACHES - k_cache: torch.Tensor, - v_cache: torch.Tensor, - # CONSTANTS - softmax_scale: Optional[float] = None, - rope_theta: Optional[float] = None, -) -> torch.Tensor: - """Flattened & fused MLA with cache with triton kernels.""" - # b, s info - # NOTE: b, s are just the shapes of the input tensor q; not necessarily the number of sequences. - # Generally speaking, we expect one of two cases here: - # 1. b > 0, s==1: this indicates a generate-only batch of tokens. - # 2. b==1, s > 0: this indicates a mixed context+generate phase. The actual number of sequences - # and number of tokens per sequence are encoded in seq_len and seq_start. - - # check for sequence info and truncate metadata - num_prefill, num_prefill_tokens, num_decode = batch_info_host.tolist() - num_seq = num_prefill + num_decode - - seq_len = seq_len[:num_seq] - input_pos = input_pos[:num_seq] - cache_loc = cache_loc[:num_seq] - seq_start = cu_seqlen[:num_seq] - - # Get parameters - b, num_heads, s, qk_nope_head_dim = q_nope.shape - qk_rope_head_dim = q_pe.shape[-1] - v_head_dim = kv.shape[-1] - qk_nope_head_dim - - # Get k_nope and value_states - k_nope, value_states = torch.split(kv, [qk_nope_head_dim, v_head_dim], dim=-1) - - # Flatten inputs - if s == 1: - bs_view = (b, s) - else: - bs_view = (b * s,) - - # TODO(suyogg): do something about all these clones, transposes, and contiguous-es - q_nope = q_nope.clone().transpose(1, 2).view(*bs_view, num_heads, qk_nope_head_dim).contiguous() - q_pe = q_pe.clone().transpose(1, 2).view(*bs_view, num_heads, qk_rope_head_dim).contiguous() - k_nope = k_nope.clone().transpose(1, 2).view(*bs_view, num_heads, qk_nope_head_dim).contiguous() - k_pe = k_pe.clone().transpose(1, 2).view(*bs_view, -1, qk_rope_head_dim).contiguous() - value_states = value_states.transpose(1, 2).view(*bs_view, -1, v_head_dim).contiguous() - # Apply RoPE - if rope_theta is not None: - max_seq_len = (input_pos + seq_len).max().item() - cos_sin_stacked = _precompute_inv_freq( - max_seq_len, qk_rope_head_dim, rope_theta, q_pe.device - ) - - # Extract cos and sin from freqs_cis - cos_base = cos_sin_stacked[0, ...] - sin_base = cos_sin_stacked[1, ...] - - # TODO: Use triton kernels for RoPE - # TODO: Add yarn support - for i in range(seq_len.shape[0]): - start = seq_start[i] - length = seq_len[i] - - # build position_ids - if s == 1: - idx = (input_pos[i] + length - 1).item() - pos_ids = torch.tensor(idx, device=cos_base.device) - else: - pos_ids = torch.arange(input_pos[i], input_pos[i] + length, device=cos_base.device) - - cos = cos_base[pos_ids] # [..., 1, head_dim] - sin = sin_base[pos_ids] - q_slice = q_pe[start : start + length] - k_slice = k_pe[start : start + length] - - q_rot, k_rot = torch.ops.auto_deploy.torch_rope_with_qk_interleaving( - q_slice, - k_slice, - cos, - sin, - -2, - ) - - q_pe[start : start + length] = q_rot - k_pe[start : start + length] = k_rot - - # Create query_states, key_states - query_states = torch.cat((q_nope, q_pe), dim=-1) # [b*s,n,d] - key_states = torch.cat((k_nope, k_pe.expand(*bs_view, num_heads, -1)), dim=-1) # [b*s,n,d] - - # Compute scale if not provided - qk_head_dim = qk_nope_head_dim + qk_rope_head_dim - scale = softmax_scale if softmax_scale is not None else 1.0 / math.sqrt(qk_head_dim) - - # Compute attention - y = torch.empty_like(value_states) - if s == 1: - # generate-only phase (decode) - _decode_attention( - query_states.contiguous(), - key_states.contiguous(), - value_states.contiguous(), - k_cache, - v_cache, - cache_loc, - input_pos, - scale, - y, - ) - - else: - # mixed context + generate phase (prefill) - _prefill_attention( - query_states.contiguous(), - key_states.contiguous(), - value_states.contiguous(), - k_cache, - v_cache, - input_pos, - cache_loc, - seq_len, - seq_start, - scale, - y, - ) - - y = ( - y.view(b, s, -1, v_head_dim).transpose(1, 2).contiguous() - ) # BNSD format as expected by the callsite. - return y - - -@fused_flattened_mla_with_cache.register_fake -def fused_flattened_mla_with_cache_fake( - # Q, K, V - q_nope: torch.Tensor, - q_pe: torch.Tensor, - kv: torch.Tensor, - k_pe: torch.Tensor, - # STANDARD METADATA - batch_info_host: torch.Tensor, - seq_len: torch.Tensor, - input_pos: torch.Tensor, - cache_loc: torch.Tensor, - cu_seqlen: torch.Tensor, - # EXTRA METADATA - # - # CACHES - k_cache: torch.Tensor, - v_cache: torch.Tensor, - # CONSTANTS - softmax_scale: Optional[float] = None, - rope_theta: Optional[float] = None, -): - v_head_dim = kv.shape[-1] - q_nope.shape[-1] - return torch.empty_like(kv[..., -v_head_dim:]) - - -@AttentionRegistry.register("MultiHeadLatentAttention") -class MultiHeadLatentAttention(AttentionDescriptor): - @classmethod - def get_attention_layout(cls) -> AttentionLayout: - """Get the attention layout expected by the backend.""" - return "bnsd" - - @classmethod - def get_num_qkv_args(cls) -> int: - """Get the number of qkv arguments expected by the source op.""" - return 4 - - @classmethod - def get_source_attention_op(cls) -> OpOverloadPacket: - return torch.ops.auto_deploy.torch_attention_deepseek_fused_mla - - @classmethod - def get_cached_attention_op(cls) -> MHACallable: - return torch.ops.auto_deploy.triton_attention_fused_flattened_mla_with_cache.default - - @classmethod - def get_standard_metadata_args(cls) -> List[str]: - return ["batch_info_host", "seq_len", "input_pos", "cache_loc", "cu_seqlen"] - - @classmethod - def get_cache_initializers( - cls, source_attn_node: Node, cache_config: KvCacheConfig - ) -> ResourceHandlerDict: - q_nope_fake = source_attn_node.args[0].meta["val"] - q_pe_fake = source_attn_node.args[1].meta["val"] - kv_fake = source_attn_node.args[2].meta["val"] - - num_kv_heads = kv_fake.shape[1] - head_dim = q_nope_fake.shape[-1] - rope_dim = q_pe_fake.shape[-1] - - return { - "k_cache": UnpagedResourceHandler( - num_kv_heads, - head_dim + rope_dim, - dtype=cls.resolve_cache_dtype(cache_config.dtype, kv_fake.dtype), - ), - "v_cache": UnpagedResourceHandler( - num_kv_heads, - head_dim, - dtype=cls.resolve_cache_dtype(cache_config.dtype, kv_fake.dtype), - ), - } - - @classmethod - def get_constants(cls, source_attn_node: Node) -> List[Constant]: - softmax_scale = None - rope_theta = 10000.0 # TODO: remove once MLA is unfused - return [softmax_scale, rope_theta] diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/torch_backend_mla.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/torch_backend_mla.py new file mode 100644 index 000000000000..28cda4cb0ef5 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/torch_backend_mla.py @@ -0,0 +1,518 @@ +"""Custom ops for MultiHead Latent Attention (MLA) with FlashInfer-compatible cache. + +This module provides: +- torch_cached_mla_with_cache: cached backend op +- TorchBackendMLAAttention: attention descriptor + +FlashInfer MLA Cache Layout: + mla_cache: [max_batch, max_seq, kv_lora_rank + qk_rope_head_dim] + - No num_heads dimension (MLA-specific optimization) + - compressed_kv_cached = mla_cache[:, :, :kv_lora_rank] (zero-copy slice) + - kpe_cached = mla_cache[:, :, kv_lora_rank:] (zero-copy slice) + +The implementation uses: +- Prefill: Expand compressed_kv -> full K, V, compute normal attention +- Generate: Weight absorption for efficiency (Q @ W^T instead of expanding cached KV) + +Reference: https://docs.flashinfer.ai/tutorials/kv_layout.html#mla-page-layout +""" + +import math +from typing import List, Optional + +import torch +from torch._ops import OpOverloadPacket +from torch.fx import Node + +from .....llmapi.llm_args import KvCacheConfig +from ..attention_interface import ( + AttentionDescriptor, + AttentionLayout, + AttentionRegistry, + Constant, + MHACallable, + ResourceHandlerDict, + UnpagedResourceHandler, +) + + +def _update_mla_cache( + compressed_kv: torch.Tensor, # [total_tokens, kv_lora_rank] + kpe: torch.Tensor, # [total_tokens, qk_rope_head_dim] + mla_cache: torch.Tensor, # [max_batch, max_seq, kv_lora_rank + qk_rope_head_dim] + seq_len: torch.Tensor, + input_pos: torch.Tensor, + slot_idx: torch.Tensor, + seq_start: torch.Tensor, + kv_lora_rank: int, +) -> None: + """Update FlashInfer MLA cache with compressed_kv and kpe values. + + FlashInfer MLA cache layout: [max_batch, max_seq, kv_lora_rank + qk_rope_head_dim] + - First kv_lora_rank dims: compressed KV latent (before kv_b_proj) + - Last qk_rope_head_dim dims: key positional encoding + """ + for idx in range(seq_len.shape[0]): + start = seq_start[idx].item() + length = seq_len[idx].item() + cache_idx = slot_idx[idx].item() + pos = input_pos[idx].item() + + # Update compressed_kv portion + mla_cache[cache_idx, pos : pos + length, :kv_lora_rank] = compressed_kv[ + start : start + length + ] + # Update kpe portion + mla_cache[cache_idx, pos : pos + length, kv_lora_rank:] = kpe[start : start + length] + + +def _torch_mla_generate_with_absorption( + q_nope: torch.Tensor, # [B, 1, N, qk_nope_head_dim] + q_pe: torch.Tensor, # [B, 1, N, qk_rope_head_dim] + compressed_kv: torch.Tensor, # [B, 1, kv_lora_rank] + kpe: torch.Tensor, # [B, 1, 1, qk_rope_head_dim] + kv_b_proj_weight: torch.Tensor, # [N * (qk_nope_head_dim + v_head_dim), kv_lora_rank] + mla_cache: torch.Tensor, # [max_batch, max_seq, kv_lora_rank + qk_rope_head_dim] + slot_idx: torch.Tensor, + input_pos: torch.Tensor, + scale: float, + kv_lora_rank: int, + num_heads: int, + qk_nope_head_dim: int, + v_head_dim: int, + out: torch.Tensor, +) -> None: + """Generate-only MLA attention with weight absorption. + + Weight absorption: Instead of expanding all cached KV, we absorb kv_b_proj into Q. + Q_absorbed = Q_nope @ W_k^T where W_k is the k_nope portion of kv_b_proj_weight + + This avoids expanding potentially thousands of cached tokens. + """ + b = q_nope.shape[0] + + # Extract k_nope and v portions from kv_b_proj_weight + # kv_b_proj_weight: [N * (qk_nope_head_dim + v_head_dim), kv_lora_rank] + # Reshape to [N, qk_nope_head_dim + v_head_dim, kv_lora_rank] + weight_reshaped = kv_b_proj_weight.view(num_heads, qk_nope_head_dim + v_head_dim, kv_lora_rank) + w_k_nope = weight_reshaped[:, :qk_nope_head_dim, :] # [N, qk_nope_head_dim, kv_lora_rank] + w_v = weight_reshaped[:, qk_nope_head_dim:, :] # [N, v_head_dim, kv_lora_rank] + + # Update cache with new tokens + compressed_kv_flat = compressed_kv.squeeze(1) # [B, kv_lora_rank] + kpe_flat = kpe.squeeze(1).squeeze(1) # [B, qk_rope_head_dim] + + for i in range(b): + cache_idx = slot_idx[i].item() + pos = input_pos[i].item() + mla_cache[cache_idx, pos, :kv_lora_rank] = compressed_kv_flat[i] + mla_cache[cache_idx, pos, kv_lora_rank:] = kpe_flat[i] + + # Compute attention for each sequence using weight absorption + for i in range(b): + cache_idx = slot_idx[i].item() + pos = input_pos[i].item() + + # Get query for this sequence: [N, qk_nope_head_dim], [N, qk_rope_head_dim] + q_nope_i = q_nope[i, 0] # [N, qk_nope_head_dim] + q_pe_i = q_pe[i, 0] # [N, qk_rope_head_dim] + + # Retrieve cached data up to current position + cached_data = mla_cache[cache_idx, : pos + 1] # [seq_len, kv_lora_rank + qk_rope_head_dim] + compressed_kv_cached = cached_data[:, :kv_lora_rank] # [seq_len, kv_lora_rank] + kpe_cached = cached_data[:, kv_lora_rank:] # [seq_len, qk_rope_head_dim] + + # ===================================================================== + # Weight absorption for Q_nope part + # ===================================================================== + # q_absorbed = q_nope @ w_k_nope^T (absorb k_nope projection into query) + # q_nope_i: [N, qk_nope_head_dim] + # w_k_nope: [N, qk_nope_head_dim, kv_lora_rank] + # q_absorbed: [N, kv_lora_rank] + q_absorbed = torch.einsum("nd,ndk->nk", q_nope_i, w_k_nope) + + # Attention scores from absorbed Q and compressed KV + # Compute in fp32 to match FlashInfer's use_fp16_qk_reduction=False + # q_absorbed: [N, kv_lora_rank], compressed_kv_cached: [seq_len, kv_lora_rank] + # scores_nope: [N, seq_len] + scores_nope = torch.matmul(q_absorbed.float(), compressed_kv_cached.float().t()) + + # ===================================================================== + # Q_pe part - standard attention with kpe + # ===================================================================== + # q_pe_i: [N, qk_rope_head_dim], kpe_cached: [seq_len, qk_rope_head_dim] + # scores_pe: [N, seq_len] + scores_pe = torch.matmul(q_pe_i.float(), kpe_cached.float().t()) + + # Combined attention scores (already in fp32) + attn_scores = (scores_nope + scores_pe) * scale # [N, seq_len] + + # Softmax (already in fp32, convert back to input dtype) + attn_weights = torch.softmax(attn_scores, dim=-1).to(q_nope.dtype) # [N, seq_len] + + # ===================================================================== + # Compute output with absorbed value projection + # ===================================================================== + # v_out = attn_weights @ compressed_kv @ w_v^T + # First: weighted_kv = attn_weights @ compressed_kv_cached -> [N, kv_lora_rank] + weighted_kv = torch.matmul(attn_weights, compressed_kv_cached) # [N, kv_lora_rank] + + # Then: attn_out = weighted_kv @ w_v^T -> [N, v_head_dim] + # w_v: [N, v_head_dim, kv_lora_rank] + # weighted_kv: [N, kv_lora_rank] + attn_out = torch.einsum("nk,nvk->nv", weighted_kv, w_v) # [N, v_head_dim] + + out[i] = attn_out + + +def _torch_mla_context_with_expansion( + q_nope: torch.Tensor, # [total_tokens, N, qk_nope_head_dim] + q_pe: torch.Tensor, # [total_tokens, N, qk_rope_head_dim] + compressed_kv: torch.Tensor, # [total_tokens, kv_lora_rank] + kpe: torch.Tensor, # [total_tokens, 1, qk_rope_head_dim] + kv_b_proj_weight: torch.Tensor, # [N * (qk_nope_head_dim + v_head_dim), kv_lora_rank] + mla_cache: torch.Tensor, # [max_batch, max_seq, kv_lora_rank + qk_rope_head_dim] + input_pos: torch.Tensor, + slot_idx: torch.Tensor, + seq_len: torch.Tensor, + seq_start: torch.Tensor, + scale: float, + kv_lora_rank: int, + num_heads: int, + qk_nope_head_dim: int, + v_head_dim: int, + out: torch.Tensor, +) -> None: + """Context MLA attention with kv_b_proj expansion. + + For prefill, we expand compressed_kv using kv_b_proj_weight and compute + standard attention. This is more efficient than absorption for prefill + since we only expand the current tokens, not the full cache. + """ + + # Flatten kpe: [total_tokens, 1, qk_rope_head_dim] -> [total_tokens, qk_rope_head_dim] + kpe_flat = kpe.squeeze(1) + + # Update cache first with compressed representation + _update_mla_cache( + compressed_kv, + kpe_flat, + mla_cache, + seq_len, + input_pos, + slot_idx, + seq_start, + kv_lora_rank, + ) + + # Compute attention for each sequence + attn_outputs = [] + for idx in range(seq_len.shape[0]): + seq_len_i = seq_len[idx].item() + input_pos_i = input_pos[idx].item() + slot_idx_i = slot_idx[idx].item() + seq_start_i = seq_start[idx].item() + + if seq_len_i == 0: + continue + + # Get query for this sequence + q_nope_seq = q_nope[ + seq_start_i : seq_start_i + seq_len_i + ] # [seq_len_i, N, qk_nope_head_dim] + q_pe_seq = q_pe[seq_start_i : seq_start_i + seq_len_i] # [seq_len_i, N, qk_rope_head_dim] + + # Get cached data for attention (includes just-added tokens) + kv_seq_len = input_pos_i + seq_len_i + cached_data = mla_cache[ + slot_idx_i, :kv_seq_len + ] # [kv_seq_len, kv_lora_rank + qk_rope_head_dim] + compressed_kv_cached = cached_data[:, :kv_lora_rank] # [kv_seq_len, kv_lora_rank] + kpe_cached = cached_data[:, kv_lora_rank:] # [kv_seq_len, qk_rope_head_dim] + + # ===================================================================== + # Expand compressed_kv using kv_b_proj_weight for this sequence + # ===================================================================== + # compressed_kv_cached: [kv_seq_len, kv_lora_rank] + # kv_b_proj_weight: [N * (qk_nope_head_dim + v_head_dim), kv_lora_rank] + # kv_expanded: [kv_seq_len, N * (qk_nope_head_dim + v_head_dim)] + kv_expanded = torch.matmul(compressed_kv_cached, kv_b_proj_weight.t()) + + # Reshape to [kv_seq_len, N, qk_nope_head_dim + v_head_dim] + kv_expanded = kv_expanded.view(kv_seq_len, num_heads, qk_nope_head_dim + v_head_dim) + + # Split into k_nope and v + k_nope_expanded = kv_expanded[:, :, :qk_nope_head_dim] # [kv_seq_len, N, qk_nope_head_dim] + v_expanded = kv_expanded[:, :, qk_nope_head_dim:] # [kv_seq_len, N, v_head_dim] + + # Expand kpe to all heads + kpe_expanded = kpe_cached.unsqueeze(1).expand( + -1, num_heads, -1 + ) # [kv_seq_len, N, qk_rope_head_dim] + + # Construct full query and key + query_full = torch.cat([q_nope_seq, q_pe_seq], dim=-1) # [seq_len_i, N, qk_head_dim] + key_full = torch.cat( + [k_nope_expanded, kpe_expanded], dim=-1 + ) # [kv_seq_len, N, qk_head_dim] + + # Transpose for attention: [1, N, seq_len, head_dim] + query_t = query_full.transpose(0, 1).unsqueeze(0) # [1, N, seq_len_i, qk_head_dim] + key_t = key_full.transpose(0, 1).unsqueeze(0) # [1, N, kv_seq_len, qk_head_dim] + + # Compute attention scores in fp32 to match FlashInfer's use_fp16_qk_reduction=False + # FlashInfer uses fp32 accumulation for QK^T, so we do the same for numerical consistency + attn_scores = ( + torch.matmul(query_t.float(), key_t.float().transpose(-2, -1)) * scale + ) # [1, N, seq_len_i, kv_seq_len] in fp32 + + # Apply causal mask + causal_mask = torch.triu( + torch.ones(seq_len_i, kv_seq_len, device=q_nope.device, dtype=torch.bool), + diagonal=kv_seq_len - seq_len_i + 1, + ) + attn_scores.masked_fill_(causal_mask.unsqueeze(0).unsqueeze(0), float("-inf")) + + # Softmax (already in fp32, convert back to input dtype) + attn_weights = torch.softmax(attn_scores, dim=-1).to(q_nope.dtype) + + # Value: [1, N, kv_seq_len, v_head_dim] + v_t = v_expanded.transpose(0, 1).unsqueeze(0) + + # Compute output + attn_out = torch.matmul(attn_weights, v_t) # [1, N, seq_len_i, v_head_dim] + attn_out = attn_out[0].transpose(0, 1) # [seq_len_i, N, v_head_dim] + + attn_outputs.append(attn_out) + + # Concatenate all outputs + if len(attn_outputs) == 0: + out.zero_() + elif len(attn_outputs) == 1: + out.copy_(attn_outputs[0]) + else: + out.copy_(torch.cat(attn_outputs, dim=0)) + + +@torch.library.custom_op("auto_deploy::torch_cached_mla_with_cache", mutates_args=()) +def torch_backend_mla_with_cache( + # 5 tensor args (get_num_qkv_args = 5) + q_nope: torch.Tensor, # [B, S, N, qk_nope_head_dim] + q_pe: torch.Tensor, # [B, S, N, qk_rope_head_dim] + compressed_kv: torch.Tensor, # [B, S, kv_lora_rank] + kpe: torch.Tensor, # [B, S, 1, qk_rope_head_dim] + kv_b_proj_weight: torch.Tensor, # [N * (qk_nope_head_dim + v_head_dim), kv_lora_rank] + # Standard metadata + batch_info_host: torch.Tensor, + seq_len: torch.Tensor, + input_pos: torch.Tensor, + slot_idx: torch.Tensor, + cu_seqlen: torch.Tensor, + # Cache (FlashInfer layout) + mla_cache: torch.Tensor, # [max_batch, max_seq, kv_lora_rank + qk_rope_head_dim] + # Constants + scale: Optional[float] = None, + kv_lora_rank: int = 512, +) -> torch.Tensor: + """Torch backend MLA with FlashInfer-compatible compressed cache. + + FlashInfer MLA Cache Layout: + mla_cache: [max_batch, max_seq, kv_lora_rank + qk_rope_head_dim] + - compressed_kv = mla_cache[:, :, :kv_lora_rank] (zero-copy slice) + - kpe = mla_cache[:, :, kv_lora_rank:] (zero-copy slice) + + Prefill (context): Expand compressed_kv, compute normal attention + Generate (decode): Use weight absorption for efficiency + """ + # Get dimensions + b, s = q_nope.shape[:2] + num_heads = q_nope.shape[2] + qk_nope_head_dim = q_nope.shape[3] + qk_rope_head_dim = q_pe.shape[3] + qk_head_dim = qk_nope_head_dim + qk_rope_head_dim + + # Infer v_head_dim from kv_b_proj_weight + out_features = kv_b_proj_weight.shape[0] + kv_head_dim = out_features // num_heads + v_head_dim = kv_head_dim - qk_nope_head_dim + + # Get cleaned up metadata + num_prefill, num_prefill_tokens, num_decode = batch_info_host.tolist() + num_seq = num_prefill + num_decode + seq_len = seq_len[:num_seq] + input_pos = input_pos[:num_seq] + slot_idx = slot_idx[:num_seq] + seq_start = cu_seqlen[:num_seq] + + # Set scale + if scale is None: + scale = 1.0 / math.sqrt(qk_head_dim) + + # Define output shape: [B, S, N, v_head_dim] + output_shape = (b, s, num_heads, v_head_dim) + + if s == 1: + # ===================================================================== + # Generate phase: Use weight absorption + # ===================================================================== + y = q_nope.new_empty(b, num_heads, v_head_dim).contiguous() + + _torch_mla_generate_with_absorption( + q_nope, + q_pe, + compressed_kv, + kpe, + kv_b_proj_weight, + mla_cache, + slot_idx, + input_pos, + scale, + kv_lora_rank, + num_heads, + qk_nope_head_dim, + v_head_dim, + y, + ) + + return y.unsqueeze(1) # [B, 1, N, v_head_dim] + else: + # ===================================================================== + # Context phase: Expand and compute normal attention + # ===================================================================== + bs_view = (b * s,) + + q_nope_flat = q_nope.contiguous().view(*bs_view, num_heads, qk_nope_head_dim) + q_pe_flat = q_pe.contiguous().view(*bs_view, num_heads, qk_rope_head_dim) + compressed_kv_flat = compressed_kv.contiguous().view(*bs_view, kv_lora_rank) + kpe_flat = kpe.contiguous().view(*bs_view, 1, qk_rope_head_dim) + + y = q_nope.new_empty(*bs_view, num_heads, v_head_dim).contiguous() + + _torch_mla_context_with_expansion( + q_nope_flat, + q_pe_flat, + compressed_kv_flat, + kpe_flat, + kv_b_proj_weight, + mla_cache, + input_pos, + slot_idx, + seq_len, + seq_start, + scale, + kv_lora_rank, + num_heads, + qk_nope_head_dim, + v_head_dim, + y, + ) + + return y.view(*output_shape) + + +@torch_backend_mla_with_cache.register_fake +def torch_backend_mla_with_cache_fake( + q_nope: torch.Tensor, + q_pe: torch.Tensor, + compressed_kv: torch.Tensor, + kpe: torch.Tensor, + kv_b_proj_weight: torch.Tensor, + batch_info_host: torch.Tensor, + seq_len: torch.Tensor, + input_pos: torch.Tensor, + slot_idx: torch.Tensor, + cu_seqlen: torch.Tensor, + mla_cache: torch.Tensor, + scale: Optional[float] = None, + kv_lora_rank: int = 512, +) -> torch.Tensor: + """Fake implementation for torch_backend_mla_with_cache.""" + num_heads = q_nope.shape[2] + qk_nope_head_dim = q_nope.shape[-1] + out_features = kv_b_proj_weight.shape[0] + kv_head_dim = out_features // num_heads + v_head_dim = kv_head_dim - qk_nope_head_dim + + return q_nope.new_empty( + q_nope.shape[0], q_nope.shape[1], q_nope.shape[2], v_head_dim + ).contiguous() + + +@AttentionRegistry.register("torch_mla") +class TorchBackendMLAAttention(AttentionDescriptor): + """Attention descriptor for Multi-head Latent Attention (MLA). + + This descriptor uses FlashInfer-compatible compressed cache: + - torch_mla: source op that expands compressed_kv for attention + - torch_cached_mla_with_cache: cached op with absorption for generate + + FlashInfer MLA Cache Layout: + mla_cache: [max_batch, max_seq, head_dim_ckv + head_dim_kpe] + - No num_heads dimension (MLA-specific optimization) + - ckv_cached = mla_cache[:, :, :head_dim_ckv] (zero-copy slice) + - kpe_cached = mla_cache[:, :, head_dim_ckv:] (zero-copy slice) + + Reference: https://docs.flashinfer.ai/tutorials/kv_layout.html#mla-page-layout + """ + + @classmethod + def get_attention_layout(cls) -> AttentionLayout: + """Get the attention layout expected by the backend.""" + return "bsnd" + + @classmethod + def get_num_qkv_args(cls) -> int: + """Get the number of tensor arguments expected by the source op.""" + return 5 # q_nope, q_pe, compressed_kv, kpe, kv_b_proj_weight + + @classmethod + def get_source_attention_op(cls) -> OpOverloadPacket: + """Get the source attention op that we target for replacement.""" + return torch.ops.auto_deploy.torch_mla + + @classmethod + def get_cached_attention_op(cls) -> MHACallable: + """Get the cached attention op.""" + return torch.ops.auto_deploy.torch_cached_mla_with_cache.default + + @classmethod + def get_standard_metadata_args(cls) -> List[str]: + """Get the list of standard metadata arguments.""" + return ["batch_info_host", "seq_len", "input_pos", "slot_idx", "cu_seqlen"] + + @classmethod + def get_cache_initializers( + cls, source_attn_node: Node, cache_config: KvCacheConfig + ) -> ResourceHandlerDict: + """Get cache initializers using FlashInfer MLA cache layout.""" + # Extract dimensions from source node args + # torch_mla signature: q_nope, q_pe, compressed_kv, kpe, kv_b_proj_weight, ... + compressed_kv_fake = source_attn_node.args[2].meta["val"] + kpe_fake = source_attn_node.args[3].meta["val"] + + # Get dimensions + # compressed_kv: [B, S, kv_lora_rank] + # kpe: [B, S, 1, qk_rope_head_dim] + kv_lora_rank = compressed_kv_fake.shape[-1] + qk_rope_head_dim = kpe_fake.shape[-1] + + # FlashInfer MLA cache: [max_batch, max_seq, kv_lora_rank + qk_rope_head_dim] + # No num_heads dimension - this is the key MLA optimization + return { + "mla_cache": UnpagedResourceHandler( + kv_lora_rank + qk_rope_head_dim, + dtype=cls.resolve_cache_dtype(cache_config.dtype, compressed_kv_fake.dtype), + ), + } + + @classmethod + def get_constants(cls, source_attn_node: Node) -> List[Constant]: + """Get constants to pass to the cached attention op.""" + # Extract kv_lora_rank for cache slicing + compressed_kv_fake = source_attn_node.args[2].meta["val"] + kv_lora_rank = compressed_kv_fake.shape[-1] + + # Get scale from kwargs + scale = source_attn_node.kwargs.get("scale", None) + + return [scale, kv_lora_rank] diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/torch_mla.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/torch_mla.py new file mode 100644 index 000000000000..5be00a764eb6 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/torch_mla.py @@ -0,0 +1,157 @@ +"""Torch reference implementation for Multi-head Latent Attention (MLA). + +This module provides the source op for MLA that: +- Accepts compressed_kv (before kv_b_proj) for FlashInfer-compatible caching +- Expands compressed_kv using kv_b_proj_weight for attention computation +- Computes standard attention with the expanded K, V +""" + +import math +from typing import Optional + +import torch + + +@torch.library.custom_op("auto_deploy::torch_mla", mutates_args=()) +def torch_mla( + q_nope: torch.Tensor, # [B, S, N, qk_nope_head_dim] + q_pe: torch.Tensor, # [B, S, N, qk_rope_head_dim] (RoPE applied) + compressed_kv: torch.Tensor, # [B, S, kv_lora_rank] - BEFORE kv_b_proj + kpe: torch.Tensor, # [B, S, 1, qk_rope_head_dim] (RoPE applied) + kv_b_proj_weight: torch.Tensor, # [num_heads * (qk_nope_head_dim + v_head_dim), kv_lora_rank] + is_causal: bool = True, + scale: Optional[float] = None, + layout: str = "bsnd", +) -> torch.Tensor: + """Multi-head Latent Attention (MLA) with FlashInfer-compatible compressed KV. + + This op expands compressed_kv using kv_b_proj_weight and computes attention. + For prefill, this is the standard formulation. For the cached version, + weight absorption is used for efficiency. + + Args: + q_nope: Query non-positional component [B, S, N, qk_nope_head_dim] (bsnd) + q_pe: Query positional component with RoPE applied [B, S, N, qk_rope_head_dim] (bsnd) + compressed_kv: Compressed KV latent [B, S, kv_lora_rank] (before kv_b_proj) + kpe: Key positional encoding with RoPE applied [B, S, 1, qk_rope_head_dim] (bsnd) + kv_b_proj_weight: Projection weights [num_heads * (qk_nope_head_dim + v_head_dim), kv_lora_rank] + is_causal: Whether to apply causal masking (default: True) + scale: Softmax scale factor (default: 1/sqrt(qk_head_dim)) + layout: Input/output layout, either "bsnd" or "bnsd" (default: "bsnd") + + Returns: + Attention output with shape [B, S, N, v_head_dim] (bsnd) + """ + if layout not in ("bnsd", "bsnd"): + raise ValueError(f"layout must be 'bnsd' or 'bsnd', got {layout!r}") + + # Get dimensions + if layout == "bsnd": + bs, s_q, num_heads, qk_nope_head_dim = q_nope.shape + qk_rope_head_dim = q_pe.shape[-1] + else: + bs, num_heads, s_q, qk_nope_head_dim = q_nope.shape + qk_rope_head_dim = q_pe.shape[-1] + + s_k = compressed_kv.shape[1] + + # Infer dimensions from kv_b_proj_weight + # kv_b_proj_weight: [num_heads * (qk_nope_head_dim + v_head_dim), kv_lora_rank] + out_features = kv_b_proj_weight.shape[0] + kv_head_dim = out_features // num_heads # qk_nope_head_dim + v_head_dim + v_head_dim = kv_head_dim - qk_nope_head_dim + + # Set scale + qk_head_dim = qk_nope_head_dim + qk_rope_head_dim + if scale is None: + scale = 1.0 / math.sqrt(qk_head_dim) + + # ========================================================================= + # Expand compressed_kv using kv_b_proj_weight (this is the prefill path) + # ========================================================================= + # compressed_kv: [B, S, kv_lora_rank] + # kv_b_proj_weight: [num_heads * kv_head_dim, kv_lora_rank] + # kv = compressed_kv @ kv_b_proj_weight.T -> [B, S, num_heads * kv_head_dim] + kv = torch.matmul(compressed_kv, kv_b_proj_weight.t()) + + # Reshape to [B, S, N, kv_head_dim] + kv = kv.view(bs, s_k, num_heads, kv_head_dim) + + # Split into k_nope and value_states + k_nope, value_states = torch.split(kv, [qk_nope_head_dim, v_head_dim], dim=-1) + + # k_nope and value_states are always [B, S, N, D] from the kv reshape above. + # We need them in [B, N, S, D] for attention computation. + k_nope = k_nope.transpose(1, 2).contiguous() + value_states = value_states.transpose(1, 2).contiguous() + + # Convert inputs to computation layout [B, N, S, D] if they come in bsnd format + if layout == "bsnd": + # [B, S, N, D] -> [B, N, S, D] + q_nope = q_nope.transpose(1, 2).contiguous() + q_pe = q_pe.transpose(1, 2).contiguous() + kpe = kpe.transpose(1, 2).contiguous() + + # kpe is [B, 1, S, qk_rope_head_dim], expand to num_heads + kpe_expanded = kpe.expand(bs, num_heads, s_k, qk_rope_head_dim) + + # Construct full query and key states + # query_states: [B, N, S, qk_head_dim] + query_states = torch.cat([q_nope, q_pe], dim=-1) + # key_states: [B, N, S, qk_head_dim] + key_states = torch.cat([k_nope, kpe_expanded], dim=-1) + + # Compute attention scores: Q @ K^T + attn_scores = ( + torch.matmul(query_states, key_states.transpose(-2, -1)) * scale + ) # [B, N, s_q, s_k] + + # Apply causal mask if specified + if is_causal and s_q == s_k: + causal_mask = torch.triu( + torch.ones(s_q, s_k, device=q_nope.device, dtype=torch.bool), + diagonal=1, + ) + attn_scores.masked_fill_(causal_mask.unsqueeze(0).unsqueeze(0), float("-inf")) + + # Compute attention weights and output + attn_weights = torch.softmax(attn_scores, dim=-1, dtype=torch.float32).to(q_nope.dtype) + attn_out = torch.matmul(attn_weights, value_states) # [B, N, s_q, v_head_dim] + + # Convert back to requested layout + if layout == "bsnd": + return attn_out.transpose(1, 2).contiguous() # [B, S, N, v_head_dim] + else: + return attn_out.contiguous() # [B, N, S, v_head_dim] + + +@torch_mla.register_fake +def torch_mla_fake( + q_nope: torch.Tensor, + q_pe: torch.Tensor, + compressed_kv: torch.Tensor, + kpe: torch.Tensor, + kv_b_proj_weight: torch.Tensor, + is_causal: bool = True, + scale: Optional[float] = None, + layout: str = "bsnd", +) -> torch.Tensor: + """Fake implementation for torch_mla.""" + # Infer v_head_dim from kv_b_proj_weight + qk_nope_head_dim = q_nope.shape[-1] + num_heads = q_nope.shape[2] if layout == "bsnd" else q_nope.shape[1] + out_features = kv_b_proj_weight.shape[0] + kv_head_dim = out_features // num_heads + v_head_dim = kv_head_dim - qk_nope_head_dim + + # Output shape depends on layout + if layout == "bsnd": + # Input: [B, S, N, D], Output: [B, S, N, v_head_dim] + return q_nope.new_empty( + q_nope.shape[0], q_nope.shape[1], q_nope.shape[2], v_head_dim + ).contiguous() + else: + # Input: [B, N, S, D], Output: [B, N, S, v_head_dim] + return q_nope.new_empty( + q_nope.shape[0], q_nope.shape[1], q_nope.shape[2], v_head_dim + ).contiguous() diff --git a/tensorrt_llm/_torch/auto_deploy/models/custom/__init__.py b/tensorrt_llm/_torch/auto_deploy/models/custom/__init__.py index 4ad9e96cb874..b9fcb1e0f01a 100644 --- a/tensorrt_llm/_torch/auto_deploy/models/custom/__init__.py +++ b/tensorrt_llm/_torch/auto_deploy/models/custom/__init__.py @@ -1,9 +1,11 @@ -from .modeling_eagle import Eagle3DrafterForCausalLM +from .modeling_deepseek import DeepSeekV3ForCausalLM +from .modeling_glm4_moe_lite import Glm4MoeLiteForCausalLM from .modeling_nemotron_flash import NemotronFlashForCausalLM, NemotronFlashPreTrainedTokenizerFast from .modeling_nemotron_h import NemotronHForCausalLM __all__ = ( - "Eagle3DrafterForCausalLM", + "DeepSeekV3ForCausalLM", + "Glm4MoeLiteForCausalLM", "NemotronFlashForCausalLM", "NemotronFlashPreTrainedTokenizerFast", "NemotronHForCausalLM", diff --git a/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_deepseek.py b/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_deepseek.py new file mode 100644 index 000000000000..add0de21fc46 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_deepseek.py @@ -0,0 +1,655 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + +"""Slimmed down PyTorch DeepSeekV3 model implementation for auto_deploy export. + +Source: +https://huggingface.co/deepseek-ai/DeepSeek-R1/blob/main/modeling_deepseek.py + +This implementation differs from the original in the following ways: +* Simplified for prefill-only inference (no KV caching) +* Uses auto_deploy custom ops for export compatibility +* Removed flash attention variants (uses torch_mla custom op) +* Removed gradient checkpointing and training code paths +* Removed attention dropout (inference only) + +This allows us to have a clean export-ready implementation with auto_deploy custom ops. +""" + +import math +from dataclasses import dataclass +from typing import Optional, Tuple + +import torch +import torch.nn.functional as F +from torch import nn +from transformers.activations import ACT2FN +from transformers.generation import GenerationMixin +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import ModelOutput + +from tensorrt_llm._torch.auto_deploy.models.hf import AutoModelForCausalLMFactory +from tensorrt_llm._torch.utils import ActivationType + + +class DeepSeekV3RMSNorm(nn.Module): + """RMS Normalization for DeepSeekV3.""" + + def __init__(self, hidden_size: int, eps: float = 1e-6): + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return torch.ops.auto_deploy.triton_rms_norm( + hidden_states, self.weight, self.variance_epsilon + ).to(hidden_states.dtype) + + +class DeepSeekV3RotaryEmbedding(nn.Module): + """Rotary Position Embedding for DeepSeekV3. + + Simplified version that precomputes and caches cos/sin values. + Returns full cached values (not sliced by seq_len) to enable export. + """ + + def __init__(self, dim: int, max_position_embeddings: int = 2048, base: float = 10000.0): + super().__init__() + self.dim = dim + self.max_position_embeddings = max_position_embeddings + self.base = base + + inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float() / self.dim)) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + # Build cos/sin cache + self._set_cos_sin_cache(max_position_embeddings) + + def _set_cos_sin_cache(self, seq_len: int): + self.max_seq_len_cached = seq_len + t = torch.arange(seq_len, dtype=self.inv_freq.dtype) + freqs = torch.outer(t, self.inv_freq) + emb = torch.cat((freqs, freqs), dim=-1) + self.register_buffer("cos_cached", emb.cos(), persistent=False) + self.register_buffer("sin_cached", emb.sin(), persistent=False) + + def forward( + self, x: torch.Tensor, seq_len: Optional[int] = None + ) -> Tuple[torch.Tensor, torch.Tensor]: + # Return full cached cos/sin (not sliced) for export compatibility + return ( + self.cos_cached.to(dtype=x.dtype, device=x.device), + self.sin_cached.to(dtype=x.dtype, device=x.device), + ) + + +class DeepSeekV3YarnRotaryEmbedding(DeepSeekV3RotaryEmbedding): + """YaRN-extended rotary embedding for DeepSeekV3.""" + + def __init__( + self, + dim: int, + max_position_embeddings: int = 2048, + base: float = 10000.0, + scaling_factor: float = 1.0, + original_max_position_embeddings: int = 4096, + beta_fast: int = 32, + beta_slow: int = 1, + mscale: float = 1.0, + mscale_all_dim: float = 0.0, + ): + self.scaling_factor = scaling_factor + self.original_max_position_embeddings = original_max_position_embeddings + self.beta_fast = beta_fast + self.beta_slow = beta_slow + self.mscale = mscale + self.mscale_all_dim = mscale_all_dim + super().__init__(dim, max_position_embeddings, base) + + def _set_cos_sin_cache(self, seq_len: int): + self.max_seq_len_cached = seq_len + dim = self.dim + + freq_extra = 1.0 / (self.base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + freq_inter = 1.0 / ( + self.scaling_factor * self.base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim) + ) + + low, high = self._yarn_find_correction_range( + self.beta_fast, + self.beta_slow, + dim, + self.base, + self.original_max_position_embeddings, + ) + inv_freq_mask = 1.0 - self._yarn_linear_ramp_mask(low, high, dim // 2) + inv_freq = freq_inter * (1 - inv_freq_mask) + freq_extra * inv_freq_mask + self.register_buffer("inv_freq", inv_freq, persistent=False) + + t = torch.arange(seq_len, dtype=torch.float32) + freqs = torch.outer(t, inv_freq) + + _mscale = float( + self._yarn_get_mscale(self.scaling_factor, self.mscale) + / self._yarn_get_mscale(self.scaling_factor, self.mscale_all_dim) + ) + + emb = torch.cat((freqs, freqs), dim=-1) + self.register_buffer("cos_cached", (emb.cos() * _mscale), persistent=False) + self.register_buffer("sin_cached", (emb.sin() * _mscale), persistent=False) + + @staticmethod + def _yarn_find_correction_dim( + num_rotations: float, dim: int, base: float = 10000, max_position_embeddings: int = 2048 + ) -> float: + return (dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi))) / ( + 2 * math.log(base) + ) + + def _yarn_find_correction_range( + self, low_rot: int, high_rot: int, dim: int, base: float, max_position_embeddings: int + ) -> Tuple[int, int]: + low = math.floor( + self._yarn_find_correction_dim(low_rot, dim, base, max_position_embeddings) + ) + high = math.ceil( + self._yarn_find_correction_dim(high_rot, dim, base, max_position_embeddings) + ) + return max(low, 0), min(high, dim - 1) + + @staticmethod + def _yarn_get_mscale(scale: float = 1.0, mscale: float = 1.0) -> float: + if scale <= 1: + return 1.0 + return 0.1 * mscale * math.log(scale) + 1.0 + + @staticmethod + def _yarn_linear_ramp_mask(min_val: float, max_val: float, dim: int) -> torch.Tensor: + if min_val == max_val: + max_val += 0.001 + linear_func = (torch.arange(dim, dtype=torch.float32) - min_val) / (max_val - min_val) + return torch.clamp(linear_func, 0, 1) + + +class DeepSeekV3MLP(nn.Module): + """MLP layer for DeepSeekV3 (SwiGLU activation).""" + + def __init__( + self, config, hidden_size: Optional[int] = None, intermediate_size: Optional[int] = None + ): + super().__init__() + self.config = config + self.hidden_size = hidden_size or config.hidden_size + self.intermediate_size = intermediate_size or config.intermediate_size + + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + + +class DeepSeekV3MoEGate(nn.Module): + """MoE Gating for DeepSeekV3 with noaux_tc top-k selection.""" + + def __init__(self, config): + super().__init__() + self.config = config + self.top_k = config.num_experts_per_tok + self.n_routed_experts = config.n_routed_experts + self.routed_scaling_factor = config.routed_scaling_factor + self.n_group = config.n_group + self.topk_group = config.topk_group + + self.weight = nn.Parameter( + torch.empty((self.n_routed_experts, config.hidden_size), dtype=torch.float32) + ) + self.register_buffer( + "e_score_correction_bias", + torch.zeros(self.n_routed_experts, dtype=torch.float32), + ) + self.reset_parameters() + + def reset_parameters(self) -> None: + """Initialize gate weights using kaiming uniform (matches original DeepSeek implementation).""" + torch.nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + + def forward(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Forward pass returning (selected_experts, routing_weights).""" + bsz, seq_len, hidden_dim = hidden_states.shape + hidden_states_flat = hidden_states.view(-1, hidden_dim) + + # Compute router logits + if self.weight.dtype == torch.float32: + router_logits = F.linear(hidden_states_flat.float(), self.weight) + else: + router_logits = torch.ops.trtllm.dsv3_router_gemm_op( + hidden_states_flat, self.weight.t(), bias=None, out_dtype=torch.float32 + ) + + # Use fused noaux_tc_op kernel for top-k selection + topk_weights, topk_indices = torch.ops.trtllm.noaux_tc_op( + router_logits, + self.e_score_correction_bias, + self.n_group, + self.topk_group, + self.top_k, + self.routed_scaling_factor, + ) + + return topk_indices, topk_weights + + +class DeepSeekV3MoE(nn.Module): + """Mixture of Experts layer for DeepSeekV3.""" + + def __init__(self, config): + super().__init__() + self.config = config + self.num_experts_per_tok = config.num_experts_per_tok + + # Routed experts + self.experts = nn.ModuleList( + [ + DeepSeekV3MLP(config, intermediate_size=config.moe_intermediate_size) + for _ in range(config.n_routed_experts) + ] + ) + + # Gate + self.gate = DeepSeekV3MoEGate(config) + + # Shared experts (if configured) + if config.n_shared_experts is not None: + intermediate_size = config.moe_intermediate_size * config.n_shared_experts + self.shared_experts = DeepSeekV3MLP(config, intermediate_size=intermediate_size) + else: + self.shared_experts = None + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + identity = hidden_states + orig_shape = hidden_states.shape + + selected_experts, routing_weights = self.gate(hidden_states) + + # Use torch_moe custom op for routed experts + final_hidden_states = torch.ops.auto_deploy.torch_moe( + hidden_states.view(-1, hidden_states.shape[-1]), + selected_experts, + routing_weights, + w1_weight=[expert.gate_proj.weight for expert in self.experts], + w2_weight=[expert.down_proj.weight for expert in self.experts], + w3_weight=[expert.up_proj.weight for expert in self.experts], + is_gated_mlp=True, + act_fn=int(ActivationType.Silu), + ) + + final_hidden_states = final_hidden_states.view(*orig_shape) + + # Add shared experts output if present + if self.shared_experts is not None: + final_hidden_states = final_hidden_states + self.shared_experts(identity) + + return final_hidden_states.to(hidden_states.dtype) + + +class DeepSeekV3Attention(nn.Module): + """Multi-head Latent Attention (MLA) for DeepSeekV3. + + Uses compressed KV representation with latent projections. + """ + + def __init__(self, config, layer_idx: Optional[int] = None): + super().__init__() + self.config = config + self.layer_idx = layer_idx + + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.q_lora_rank = config.q_lora_rank + self.kv_lora_rank = config.kv_lora_rank + self.qk_nope_head_dim = config.qk_nope_head_dim + self.qk_rope_head_dim = config.qk_rope_head_dim + self.v_head_dim = config.v_head_dim + self.q_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim + + self.max_position_embeddings = config.max_position_embeddings + self.rope_theta = config.rope_theta + + # Q projection (with optional LoRA) + if self.q_lora_rank is None: + self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.q_head_dim, bias=False) + else: + self.q_a_proj = nn.Linear( + self.hidden_size, config.q_lora_rank, bias=config.attention_bias + ) + self.q_a_layernorm = DeepSeekV3RMSNorm(config.q_lora_rank) + self.q_b_proj = nn.Linear( + config.q_lora_rank, self.num_heads * self.q_head_dim, bias=False + ) + + # KV projection with MQA + self.kv_a_proj_with_mqa = nn.Linear( + self.hidden_size, + self.kv_lora_rank + self.qk_rope_head_dim, + bias=config.attention_bias, + ) + self.kv_a_layernorm = DeepSeekV3RMSNorm(self.kv_lora_rank) + self.kv_b_proj = nn.Linear( + self.kv_lora_rank, + self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), + bias=False, + ) + + # Output projection + self.o_proj = nn.Linear( + self.num_heads * self.v_head_dim, self.hidden_size, bias=config.attention_bias + ) + + # Initialize rotary embedding + self._init_rope() + + # Softmax scale + self.softmax_scale = self.q_head_dim ** (-0.5) + if config.rope_scaling is not None: + mscale_all_dim = config.rope_scaling.get("mscale_all_dim", 0) + scaling_factor = config.rope_scaling["factor"] + if mscale_all_dim: + mscale = DeepSeekV3YarnRotaryEmbedding._yarn_get_mscale( + scaling_factor, mscale_all_dim + ) + self.softmax_scale = self.softmax_scale * mscale * mscale + + def _init_rope(self): + if self.config.rope_scaling is None: + self.rotary_emb = DeepSeekV3RotaryEmbedding( + self.qk_rope_head_dim, + max_position_embeddings=self.max_position_embeddings, + base=self.rope_theta, + ) + else: + scaling_type = self.config.rope_scaling["type"] + scaling_factor = self.config.rope_scaling["factor"] + + if scaling_type == "yarn": + kwargs = { + key: self.config.rope_scaling[key] + for key in [ + "original_max_position_embeddings", + "beta_fast", + "beta_slow", + "mscale", + "mscale_all_dim", + ] + if key in self.config.rope_scaling + } + self.rotary_emb = DeepSeekV3YarnRotaryEmbedding( + self.qk_rope_head_dim, + max_position_embeddings=self.max_position_embeddings, + scaling_factor=scaling_factor, + base=self.rope_theta, + **kwargs, + ) + else: + # Default to base rotary embedding for unsupported types + self.rotary_emb = DeepSeekV3RotaryEmbedding( + self.qk_rope_head_dim, + max_position_embeddings=self.max_position_embeddings, + base=self.rope_theta, + ) + + def forward( + self, + hidden_states: torch.Tensor, + position_ids: torch.Tensor, + ) -> torch.Tensor: + bsz, q_len, _ = hidden_states.size() + + # Q projection + if self.q_lora_rank is None: + q = self.q_proj(hidden_states) + else: + q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))) + + # Shape: [B, S, N, q_head_dim] (BSND layout) + q = q.view(bsz, q_len, self.num_heads, self.q_head_dim) + q_nope, q_pe = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + + # KV projection - keep compressed form + kv_a_output = self.kv_a_proj_with_mqa(hidden_states) + compressed_kv, k_pe = torch.split( + kv_a_output, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 + ) + + # Apply layernorm to compressed_kv + compressed_kv = self.kv_a_layernorm(compressed_kv) + + # k_pe: [B, S, 1, qk_rope_head_dim] (BSND layout, shared across heads) + k_pe = k_pe.view(bsz, q_len, 1, self.qk_rope_head_dim) + + kv_seq_len = q_len + + # Get cos/sin for RoPE + cos, sin = self.rotary_emb(hidden_states, seq_len=kv_seq_len) + cos = cos[position_ids] # [B, S, head_dim] + sin = sin[position_ids] # [B, S, head_dim] + + # Apply RoPE using custom op + q_pe_rotated, kpe = torch.ops.auto_deploy.torch_rope_with_qk_interleaving( + q_pe, + k_pe, + cos, + sin, + 2, # unsqueeze_dim=2 for BSND layout + ) + + # Call MLA with compressed KV + attn_output = torch.ops.auto_deploy.torch_mla( + q_nope, # [B, S, N, qk_nope_head_dim] + q_pe_rotated, # [B, S, N, qk_rope_head_dim] + compressed_kv, # [B, S, kv_lora_rank] + kpe, # [B, S, 1, qk_rope_head_dim] + self.kv_b_proj.weight, # [N*(qk_nope+v), kv_lora_rank] + True, # is_causal + self.softmax_scale, + "bsnd", # layout + ) + + # Output: [B, S, N, v_head_dim] -> [B, S, N * v_head_dim] + attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.v_head_dim) + attn_output = self.o_proj(attn_output) + + return attn_output + + +class DeepSeekV3DecoderLayer(nn.Module): + """Transformer decoder layer for DeepSeekV3.""" + + def __init__(self, config, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + self.layer_idx = layer_idx + + # Attention + self.self_attn = DeepSeekV3Attention(config, layer_idx=layer_idx) + + # MLP or MoE + # MoE layers are used after first_k_dense_replace and at moe_layer_freq intervals + use_moe = ( + config.n_routed_experts is not None + and layer_idx >= config.first_k_dense_replace + and layer_idx % config.moe_layer_freq == 0 + ) + if use_moe: + self.mlp = DeepSeekV3MoE(config) + else: + self.mlp = DeepSeekV3MLP(config) + + # Layer norms + self.input_layernorm = DeepSeekV3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = DeepSeekV3RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + def forward( + self, + hidden_states: torch.Tensor, + position_ids: torch.Tensor, + ) -> torch.Tensor: + # Self attention + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self.self_attn(hidden_states, position_ids) + hidden_states = residual + hidden_states + + # MLP/MoE + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + return hidden_states + + +@dataclass +class DeepSeekV3Output(ModelOutput): + """Output for DeepSeekV3Model.""" + + last_hidden_state: Optional[torch.FloatTensor] = None + + +@dataclass +class DeepSeekV3CausalLMOutput(ModelOutput): + """Output for DeepSeekV3ForCausalLM.""" + + logits: Optional[torch.FloatTensor] = None + + +class DeepSeekV3PreTrainedModel(PreTrainedModel): + """Base class for DeepSeekV3 models.""" + + base_model_prefix = "model" + _no_split_modules = ["DeepSeekV3DecoderLayer"] + supports_gradient_checkpointing = False + + def _init_weights(self, module): + std = self.config.initializer_range + if isinstance(module, nn.Linear): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + + +class DeepSeekV3Model(DeepSeekV3PreTrainedModel): + """DeepSeekV3 transformer decoder model.""" + + def __init__(self, config): + super().__init__(config) + self.config = config + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList( + [ + DeepSeekV3DecoderLayer(config, layer_idx=idx) + for idx in range(config.num_hidden_layers) + ] + ) + self.norm = DeepSeekV3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + self.post_init() + + def get_input_embeddings(self): + return self.embed_tokens + + def set_input_embeddings(self, value): + self.embed_tokens = value + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + **kwargs, + ) -> DeepSeekV3Output: + if input_ids is not None and inputs_embeds is not None: + raise ValueError("Cannot specify both input_ids and inputs_embeds") + elif input_ids is None and inputs_embeds is None: + raise ValueError("Must specify either input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + batch_size, seq_length = inputs_embeds.shape[:2] + + if position_ids is None: + device = input_ids.device if input_ids is not None else inputs_embeds.device + position_ids = torch.arange(seq_length, dtype=torch.long, device=device) + position_ids = position_ids.unsqueeze(0).expand(batch_size, -1) + + hidden_states = inputs_embeds + + for decoder_layer in self.layers: + hidden_states = decoder_layer(hidden_states, position_ids) + + hidden_states = self.norm(hidden_states) + + return DeepSeekV3Output(last_hidden_state=hidden_states) + + +class DeepSeekV3ForCausalLM(DeepSeekV3PreTrainedModel, GenerationMixin): + """DeepSeekV3 model with language modeling head.""" + + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config): + super().__init__(config) + self.model = DeepSeekV3Model(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def get_decoder(self): + return self.model + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + **kwargs, + ) -> DeepSeekV3CausalLMOutput: + outputs = self.model( + input_ids=input_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + **kwargs, + ) + + hidden_states = outputs.last_hidden_state + logits = self.lm_head(hidden_states).float() + + return DeepSeekV3CausalLMOutput(logits=logits) + + +# Register with AutoModelForCausalLMFactory +AutoModelForCausalLMFactory.register_custom_model_cls("DeepseekV3Config", DeepSeekV3ForCausalLM) diff --git a/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_glm4_moe_lite.py b/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_glm4_moe_lite.py new file mode 100644 index 000000000000..17aec049406a --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_glm4_moe_lite.py @@ -0,0 +1,830 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + +"""Slimmed down PyTorch GLM4 MoE Lite model implementation for auto_deploy export. + +Source: +https://huggingface.co/zai-org/GLM-4.7-Flash + +This implementation differs from the original HuggingFace version in the following ways: +* Bundled config class to work with transformers v4.57 (model requires v5.0) +* Simplified for prefill-only inference (no KV caching) +* Uses auto_deploy custom ops for export compatibility +* Removed flash attention variants (uses torch_mla custom op) +* Removed gradient checkpointing and training code paths +* Removed attention dropout (inference only) + +The GLM4 MoE Lite model uses Multi-head Latent Attention (MLA), similar to DeepSeek V3. +""" + +import math +from dataclasses import dataclass +from typing import Optional, Tuple + +import torch +import torch.nn.functional as F +from torch import nn +from transformers import AutoConfig, PretrainedConfig +from transformers.activations import ACT2FN +from transformers.generation import GenerationMixin +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import ModelOutput + +from tensorrt_llm._torch.auto_deploy.models.hf import AutoModelForCausalLMFactory +from tensorrt_llm._torch.utils import ActivationType + + +class Glm4MoeLiteConfig(PretrainedConfig): + """Configuration class for GLM4 MoE Lite model. + + This config class is bundled with the custom model implementation to enable + loading on transformers v4.57 (the model requires v5.0 where the config is + natively registered). + """ + + model_type = "glm4_moe_lite" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + vocab_size: int = 154880, + hidden_size: int = 2048, + intermediate_size: int = 10240, + num_hidden_layers: int = 47, + num_attention_heads: int = 20, + num_key_value_heads: int = 20, + hidden_act: str = "silu", + max_position_embeddings: int = 202752, + rms_norm_eps: float = 1e-5, + # MLA parameters + q_lora_rank: int = 768, + kv_lora_rank: int = 512, + qk_nope_head_dim: int = 192, + qk_rope_head_dim: int = 64, + v_head_dim: int = 256, + # MoE parameters + n_routed_experts: int = 64, + n_shared_experts: int = 1, + num_experts_per_tok: int = 4, + moe_intermediate_size: int = 1536, + n_group: int = 1, + topk_group: int = 1, + routed_scaling_factor: float = 1.8, + norm_topk_prob: bool = True, + first_k_dense_replace: int = 1, + # RoPE parameters + rope_theta: float = 1000000.0, + rope_scaling: Optional[dict] = None, + # Other parameters + attention_bias: bool = False, + attention_dropout: float = 0.0, + tie_word_embeddings: bool = False, + pad_token_id: int = 154820, + initializer_range: float = 0.02, + **kwargs, + ): + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.hidden_act = hidden_act + self.max_position_embeddings = max_position_embeddings + self.rms_norm_eps = rms_norm_eps + # MLA + self.q_lora_rank = q_lora_rank + self.kv_lora_rank = kv_lora_rank + self.qk_nope_head_dim = qk_nope_head_dim + self.qk_rope_head_dim = qk_rope_head_dim + self.v_head_dim = v_head_dim + self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim + # MoE + self.n_routed_experts = n_routed_experts + self.n_shared_experts = n_shared_experts + self.num_experts_per_tok = num_experts_per_tok + self.moe_intermediate_size = moe_intermediate_size + self.n_group = n_group + self.topk_group = topk_group + self.routed_scaling_factor = routed_scaling_factor + self.norm_topk_prob = norm_topk_prob + self.first_k_dense_replace = first_k_dense_replace + # RoPE + self.rope_theta = rope_theta + self.rope_scaling = rope_scaling + # Other + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.initializer_range = initializer_range + + super().__init__( + pad_token_id=pad_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + + +# Register config with transformers' AutoConfig so it can be loaded from HF hub +# Use exist_ok=True to handle cases where transformers already has this model type registered +# (e.g., transformers v5.0+). In those cases, AutoConfig will use the built-in config, +# but AutoModelForCausalLMFactory will still use our custom model implementation. +try: + AutoConfig.register("glm4_moe_lite", Glm4MoeLiteConfig, exist_ok=True) +except TypeError: + # Older transformers versions don't support exist_ok parameter + try: + AutoConfig.register("glm4_moe_lite", Glm4MoeLiteConfig) + except ValueError: + # Already registered by transformers, that's fine + pass + + +class Glm4MoeLiteRMSNorm(nn.Module): + """RMS Normalization for GLM4 MoE Lite. + + Uses standard torch operations so AD fusion passes can replace with + the appropriate backend (flashinfer/triton) based on config. + """ + + def __init__(self, hidden_size: int, eps: float = 1e-6): + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return self.weight * hidden_states.to(input_dtype) + + +class Glm4MoeLiteRotaryEmbedding(nn.Module): + """Rotary Position Embedding for GLM4 MoE Lite. + + Simplified version that precomputes and caches cos/sin values. + Returns full cached values (not sliced by seq_len) to enable export. + + Uses _ad_ prefix for buffer names to work with AutoDeploy's lift_to_meta. + """ + + def __init__( + self, + dim: int, + max_position_embeddings: int = 2048, + base: float = 10000.0, + attention_scaling: float = 1.0, + ): + super().__init__() + self.dim = dim + self.max_position_embeddings = max_position_embeddings + self.base = base + self.attention_scaling = attention_scaling + + inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float() / self.dim)) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + # Build cos/sin cache with AD-specific naming + self._set_cos_sin_cache(max_position_embeddings) + + def _set_cos_sin_cache(self, seq_len: int): + self.max_seq_len_cached = seq_len + t = torch.arange(seq_len, dtype=self.inv_freq.dtype) + freqs = torch.outer(t, self.inv_freq) + emb = torch.cat((freqs, freqs), dim=-1) + # Use _ad_ prefix for AutoDeploy compatibility with lift_to_meta + self.register_buffer("_ad_cos_cached", emb.cos() * self.attention_scaling, persistent=False) + self.register_buffer("_ad_sin_cached", emb.sin() * self.attention_scaling, persistent=False) + + def forward( + self, x: torch.Tensor, seq_len: Optional[int] = None + ) -> Tuple[torch.Tensor, torch.Tensor]: + # Return full cached cos/sin (not sliced) for export compatibility + return ( + self._ad_cos_cached.to(dtype=x.dtype, device=x.device), + self._ad_sin_cached.to(dtype=x.dtype, device=x.device), + ) + + +class Glm4MoeLiteYarnRotaryEmbedding(Glm4MoeLiteRotaryEmbedding): + """YaRN-extended rotary embedding for GLM4 MoE Lite.""" + + def __init__( + self, + dim: int, + max_position_embeddings: int = 2048, + base: float = 10000.0, + scaling_factor: float = 1.0, + original_max_position_embeddings: int = 4096, + beta_fast: int = 32, + beta_slow: int = 1, + mscale: float = 1.0, + mscale_all_dim: float = 0.0, + attention_scaling: float = 1.0, + ): + self.scaling_factor = scaling_factor + self.original_max_position_embeddings = original_max_position_embeddings + self.beta_fast = beta_fast + self.beta_slow = beta_slow + self.mscale = mscale + self.mscale_all_dim = mscale_all_dim + super().__init__(dim, max_position_embeddings, base, attention_scaling) + + def _set_cos_sin_cache(self, seq_len: int): + self.max_seq_len_cached = seq_len + dim = self.dim + + freq_extra = 1.0 / (self.base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + freq_inter = 1.0 / ( + self.scaling_factor * self.base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim) + ) + + low, high = self._yarn_find_correction_range( + self.beta_fast, + self.beta_slow, + dim, + self.base, + self.original_max_position_embeddings, + ) + inv_freq_mask = 1.0 - self._yarn_linear_ramp_mask(low, high, dim // 2) + inv_freq = freq_inter * (1 - inv_freq_mask) + freq_extra * inv_freq_mask + self.register_buffer("inv_freq", inv_freq, persistent=False) + + t = torch.arange(seq_len, dtype=torch.float32) + freqs = torch.outer(t, inv_freq) + + _mscale = float( + self._yarn_get_mscale(self.scaling_factor, self.mscale) + / self._yarn_get_mscale(self.scaling_factor, self.mscale_all_dim) + ) + + emb = torch.cat((freqs, freqs), dim=-1) + # Use _ad_ prefix for AutoDeploy compatibility with lift_to_meta + # Note: attention_scaling is already incorporated in _mscale for YaRN + self.register_buffer("_ad_cos_cached", (emb.cos() * _mscale), persistent=False) + self.register_buffer("_ad_sin_cached", (emb.sin() * _mscale), persistent=False) + + @staticmethod + def _yarn_find_correction_dim( + num_rotations: float, dim: int, base: float = 10000, max_position_embeddings: int = 2048 + ) -> float: + return (dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi))) / ( + 2 * math.log(base) + ) + + def _yarn_find_correction_range( + self, low_rot: int, high_rot: int, dim: int, base: float, max_position_embeddings: int + ) -> Tuple[int, int]: + low = math.floor( + self._yarn_find_correction_dim(low_rot, dim, base, max_position_embeddings) + ) + high = math.ceil( + self._yarn_find_correction_dim(high_rot, dim, base, max_position_embeddings) + ) + return max(low, 0), min(high, dim - 1) + + @staticmethod + def _yarn_get_mscale(scale: float = 1.0, mscale: float = 1.0) -> float: + if scale <= 1: + return 1.0 + return 0.1 * mscale * math.log(scale) + 1.0 + + @staticmethod + def _yarn_linear_ramp_mask(min_val: float, max_val: float, dim: int) -> torch.Tensor: + if min_val == max_val: + max_val += 0.001 + linear_func = (torch.arange(dim, dtype=torch.float32) - min_val) / (max_val - min_val) + return torch.clamp(linear_func, 0, 1) + + +class Glm4MoeLiteMLP(nn.Module): + """MLP layer for GLM4 MoE Lite (SwiGLU activation).""" + + def __init__( + self, config, hidden_size: Optional[int] = None, intermediate_size: Optional[int] = None + ): + super().__init__() + self.config = config + self.hidden_size = hidden_size or config.hidden_size + self.intermediate_size = intermediate_size or config.intermediate_size + + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + + +class Glm4MoeLiteMoEGate(nn.Module): + """MoE Gating for GLM4 MoE Lite with top-k selection. + + Uses fused TensorRT-LLM custom ops for efficient routing: + - dsv3_router_gemm_op: Fused router GEMM for non-float32 weights + - noaux_tc_op: Fused sigmoid + bias + group top-k + normalize + scale + """ + + def __init__(self, config): + super().__init__() + self.config = config + self.top_k = config.num_experts_per_tok + self.n_routed_experts = config.n_routed_experts + self.routed_scaling_factor = config.routed_scaling_factor + self.n_group = config.n_group + self.topk_group = config.topk_group + self.norm_topk_prob = getattr(config, "norm_topk_prob", True) + + # noaux_tc_op always normalizes, so norm_topk_prob must be True + if not self.norm_topk_prob: + raise ValueError( + "Glm4MoeLiteMoEGate requires norm_topk_prob=True when using fused ops. " + "The noaux_tc_op kernel always normalizes routing weights." + ) + + self.weight = nn.Parameter( + torch.empty((self.n_routed_experts, config.hidden_size), dtype=torch.float32) + ) + self.register_buffer( + "e_score_correction_bias", + torch.zeros(self.n_routed_experts, dtype=torch.float32), + ) + self.reset_parameters() + + def reset_parameters(self) -> None: + """Initialize gate weights using kaiming uniform.""" + torch.nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + + def forward(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Forward pass returning (selected_experts, routing_weights). + + Uses fused TensorRT-LLM ops for efficient routing: + 1. dsv3_router_gemm_op: Router GEMM (when weights are not float32) + 2. noaux_tc_op: Fused sigmoid + bias + group top-k + normalize + scale + """ + bsz, seq_len, hidden_dim = hidden_states.shape + hidden_states_flat = hidden_states.view(-1, hidden_dim) + + # Router GEMM - use fused op when weights are not float32 + if self.weight.dtype == torch.float32: + router_logits = F.linear(hidden_states_flat.float(), self.weight) + else: + router_logits = torch.ops.trtllm.dsv3_router_gemm_op( + hidden_states_flat, self.weight.t(), bias=None, out_dtype=torch.float32 + ) + + # Fused routing: sigmoid + bias + group top-k + normalize + scale + # noaux_tc_op internally applies: + # 1. Sigmoid to router_logits + # 2. Adds e_score_correction_bias + # 3. Group-wise top-2 scoring and top group selection + # 4. Top-k expert selection from selected groups + # 5. Gathers weights from sigmoid scores + # 6. Normalizes and scales by routed_scaling_factor + topk_weights, topk_indices = torch.ops.trtllm.noaux_tc_op( + router_logits, + self.e_score_correction_bias, + self.n_group, + self.topk_group, + self.top_k, + self.routed_scaling_factor, + ) + + return topk_indices, topk_weights + + +class Glm4MoeLiteMoE(nn.Module): + """Mixture of Experts layer for GLM4 MoE Lite.""" + + def __init__(self, config): + super().__init__() + self.config = config + self.num_experts_per_tok = config.num_experts_per_tok + + # Routed experts - use ModuleList with individual expert modules + # This creates state_dict keys like: experts.0.gate_proj.weight + # which matches the checkpoint structure + self.experts = nn.ModuleList( + [ + Glm4MoeLiteMLP(config, intermediate_size=config.moe_intermediate_size) + for _ in range(config.n_routed_experts) + ] + ) + + # Gate + self.gate = Glm4MoeLiteMoEGate(config) + + # Shared experts + if config.n_shared_experts is not None and config.n_shared_experts > 0: + intermediate_size = config.moe_intermediate_size * config.n_shared_experts + self.shared_experts = Glm4MoeLiteMLP(config, intermediate_size=intermediate_size) + else: + self.shared_experts = None + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + identity = hidden_states + orig_shape = hidden_states.shape + + selected_experts, routing_weights = self.gate(hidden_states) + + # Use torch_moe custom op for routed experts + final_hidden_states = torch.ops.auto_deploy.torch_moe( + hidden_states.view(-1, hidden_states.shape[-1]), + selected_experts, + routing_weights, + w1_weight=[expert.gate_proj.weight for expert in self.experts], + w2_weight=[expert.down_proj.weight for expert in self.experts], + w3_weight=[expert.up_proj.weight for expert in self.experts], + is_gated_mlp=True, + act_fn=int(ActivationType.Silu), + ) + + final_hidden_states = final_hidden_states.view(*orig_shape) + + # Add shared experts output if present + if self.shared_experts is not None: + final_hidden_states = final_hidden_states + self.shared_experts(identity) + + return final_hidden_states.to(hidden_states.dtype) + + +class Glm4MoeLiteAttention(nn.Module): + """Multi-head Latent Attention (MLA) for GLM4 MoE Lite. + + Uses compressed KV representation with latent projections. + Receives position embeddings from the model level (shared rotary embedding). + """ + + def __init__(self, config, layer_idx: Optional[int] = None): + super().__init__() + self.config = config + self.layer_idx = layer_idx + + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.q_lora_rank = config.q_lora_rank + self.kv_lora_rank = config.kv_lora_rank + self.qk_nope_head_dim = config.qk_nope_head_dim + self.qk_rope_head_dim = config.qk_rope_head_dim + self.v_head_dim = config.v_head_dim + self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim + + # Q projection (with optional LoRA) + if self.q_lora_rank is None: + self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.qk_head_dim, bias=False) + else: + self.q_a_proj = nn.Linear( + self.hidden_size, config.q_lora_rank, bias=config.attention_bias + ) + self.q_a_layernorm = Glm4MoeLiteRMSNorm(config.q_lora_rank) + self.q_b_proj = nn.Linear( + config.q_lora_rank, self.num_heads * self.qk_head_dim, bias=False + ) + + # KV projection with MQA + self.kv_a_proj_with_mqa = nn.Linear( + self.hidden_size, + self.kv_lora_rank + self.qk_rope_head_dim, + bias=config.attention_bias, + ) + self.kv_a_layernorm = Glm4MoeLiteRMSNorm(self.kv_lora_rank) + self.kv_b_proj = nn.Linear( + self.kv_lora_rank, + self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), + bias=False, + ) + + # Output projection + self.o_proj = nn.Linear( + self.num_heads * self.v_head_dim, self.hidden_size, bias=config.attention_bias + ) + + # Softmax scale + self.softmax_scale = self.qk_head_dim ** (-0.5) + # Apply mscale adjustment if using YaRN scaling with factor + if ( + config.rope_scaling is not None + and isinstance(config.rope_scaling, dict) + and "factor" in config.rope_scaling + ): + mscale_all_dim = config.rope_scaling.get("mscale_all_dim", 0) + scaling_factor = config.rope_scaling["factor"] + if mscale_all_dim: + mscale = Glm4MoeLiteYarnRotaryEmbedding._yarn_get_mscale( + scaling_factor, mscale_all_dim + ) + self.softmax_scale = self.softmax_scale * mscale * mscale + + def forward( + self, + hidden_states: torch.Tensor, + position_ids: torch.Tensor, + position_embeddings: Tuple[torch.Tensor, torch.Tensor], + ) -> torch.Tensor: + bsz, q_len, _ = hidden_states.size() + + # Q projection + if self.q_lora_rank is None: + q = self.q_proj(hidden_states) + else: + q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))) + + # Shape: [B, S, N, qk_head_dim] (BSND layout) + q = q.view(bsz, q_len, self.num_heads, self.qk_head_dim) + q_nope, q_pe = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + + # KV projection - keep compressed form + kv_a_output = self.kv_a_proj_with_mqa(hidden_states) + compressed_kv, k_pe = torch.split( + kv_a_output, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 + ) + + # Apply layernorm to compressed_kv + compressed_kv = self.kv_a_layernorm(compressed_kv) + + # k_pe: [B, S, 1, qk_rope_head_dim] (BSND layout, shared across heads) + k_pe = k_pe.view(bsz, q_len, 1, self.qk_rope_head_dim) + + # Get cos/sin from position_embeddings (full cached from shared rotary embedding) + cos, sin = position_embeddings # Full table: [max_seq_len, head_dim] + cos = cos[position_ids] # [B, S, head_dim] + sin = sin[position_ids] # [B, S, head_dim] + + # Apply RoPE using custom op + q_pe_rotated, kpe = torch.ops.auto_deploy.torch_rope_with_qk_interleaving( + q_pe, + k_pe, + cos, + sin, + 2, # unsqueeze_dim=2 for BSND layout + ) + + # Call MLA with compressed KV + attn_output = torch.ops.auto_deploy.torch_mla( + q_nope, # [B, S, N, qk_nope_head_dim] + q_pe_rotated, # [B, S, N, qk_rope_head_dim] + compressed_kv, # [B, S, kv_lora_rank] + kpe, # [B, S, 1, qk_rope_head_dim] + self.kv_b_proj.weight, # [N*(qk_nope+v), kv_lora_rank] + True, # is_causal + self.softmax_scale, + "bsnd", # layout + ) + + # Output: [B, S, N, v_head_dim] -> [B, S, N * v_head_dim] + attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.v_head_dim) + attn_output = self.o_proj(attn_output) + + return attn_output + + +class Glm4MoeLiteDecoderLayer(nn.Module): + """Transformer decoder layer for GLM4 MoE Lite.""" + + def __init__(self, config, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + self.layer_idx = layer_idx + + # Attention + self.self_attn = Glm4MoeLiteAttention(config, layer_idx=layer_idx) + + # MLP or MoE + # Layer 0 to first_k_dense_replace-1 use dense MLP, rest use MoE + use_moe = config.n_routed_experts is not None and layer_idx >= config.first_k_dense_replace + if use_moe: + self.mlp = Glm4MoeLiteMoE(config) + else: + self.mlp = Glm4MoeLiteMLP(config) + + # Layer norms + self.input_layernorm = Glm4MoeLiteRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = Glm4MoeLiteRMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + def forward( + self, + hidden_states: torch.Tensor, + position_ids: torch.Tensor, + position_embeddings: Tuple[torch.Tensor, torch.Tensor], + ) -> torch.Tensor: + # Self attention + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self.self_attn(hidden_states, position_ids, position_embeddings) + hidden_states = residual + hidden_states + + # MLP/MoE + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + return hidden_states + + +@dataclass +class Glm4MoeLiteOutput(ModelOutput): + """Output for Glm4MoeLiteModel.""" + + last_hidden_state: Optional[torch.FloatTensor] = None + + +@dataclass +class Glm4MoeLiteCausalLMOutput(ModelOutput): + """Output for Glm4MoeLiteForCausalLM.""" + + logits: Optional[torch.FloatTensor] = None + + +class Glm4MoeLitePreTrainedModel(PreTrainedModel): + """Base class for GLM4 MoE Lite models.""" + + config_class = Glm4MoeLiteConfig + base_model_prefix = "model" + _no_split_modules = ["Glm4MoeLiteDecoderLayer"] + supports_gradient_checkpointing = False + + def _init_weights(self, module): + std = self.config.initializer_range + if isinstance(module, nn.Linear): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + + +class Glm4MoeLiteModel(Glm4MoeLitePreTrainedModel): + """GLM4 MoE Lite transformer decoder model.""" + + def __init__(self, config): + super().__init__(config) + self.config = config + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList( + [ + Glm4MoeLiteDecoderLayer(config, layer_idx=idx) + for idx in range(config.num_hidden_layers) + ] + ) + self.norm = Glm4MoeLiteRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + # Shared rotary embedding at model level (not per-layer) + # This creates a single set of cos/sin buffers for all layers + self.rotary_emb = self._init_rope(config) + + self.post_init() + + def _init_rope(self, config): + """Initialize shared rotary embedding for all layers.""" + qk_rope_head_dim = config.qk_rope_head_dim + + # Compute attention_scaling for RoPE (same logic as in attention) + attention_scaling = 1.0 + if ( + config.rope_scaling is not None + and isinstance(config.rope_scaling, dict) + and "factor" in config.rope_scaling + ): + mscale_all_dim = config.rope_scaling.get("mscale_all_dim", 0) + scaling_factor = config.rope_scaling["factor"] + if mscale_all_dim: + mscale = Glm4MoeLiteYarnRotaryEmbedding._yarn_get_mscale( + scaling_factor, mscale_all_dim + ) + attention_scaling = mscale + + # Check if rope_scaling is None, empty, or missing required "factor" key + use_yarn = ( + config.rope_scaling is not None + and isinstance(config.rope_scaling, dict) + and "factor" in config.rope_scaling + ) + + if not use_yarn: + return Glm4MoeLiteRotaryEmbedding( + qk_rope_head_dim, + max_position_embeddings=config.max_position_embeddings, + base=config.rope_theta, + attention_scaling=attention_scaling, + ) + else: + scaling_factor = config.rope_scaling["factor"] + kwargs = { + key: config.rope_scaling[key] + for key in [ + "original_max_position_embeddings", + "beta_fast", + "beta_slow", + "mscale", + "mscale_all_dim", + ] + if key in config.rope_scaling + } + return Glm4MoeLiteYarnRotaryEmbedding( + qk_rope_head_dim, + max_position_embeddings=config.max_position_embeddings, + scaling_factor=scaling_factor, + base=config.rope_theta, + attention_scaling=attention_scaling, + **kwargs, + ) + + def get_input_embeddings(self): + return self.embed_tokens + + def set_input_embeddings(self, value): + self.embed_tokens = value + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + **kwargs, + ) -> Glm4MoeLiteOutput: + if input_ids is not None and inputs_embeds is not None: + raise ValueError("Cannot specify both input_ids and inputs_embeds") + elif input_ids is None and inputs_embeds is None: + raise ValueError("Must specify either input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + batch_size, seq_length = inputs_embeds.shape[:2] + + if position_ids is None: + device = input_ids.device if input_ids is not None else inputs_embeds.device + position_ids = torch.arange(seq_length, dtype=torch.long, device=device) + position_ids = position_ids.unsqueeze(0).expand(batch_size, -1) + + # Compute position embeddings once from shared rotary embedding + # This returns full cached cos/sin tables + position_embeddings = self.rotary_emb(inputs_embeds) + + hidden_states = inputs_embeds + + for decoder_layer in self.layers: + hidden_states = decoder_layer(hidden_states, position_ids, position_embeddings) + + hidden_states = self.norm(hidden_states) + + return Glm4MoeLiteOutput(last_hidden_state=hidden_states) + + +class Glm4MoeLiteForCausalLM(Glm4MoeLitePreTrainedModel, GenerationMixin): + """GLM4 MoE Lite model with language modeling head.""" + + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config, **kwargs): + super().__init__(config) + self.model = Glm4MoeLiteModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def get_decoder(self): + return self.model + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + **kwargs, + ) -> Glm4MoeLiteCausalLMOutput: + outputs = self.model( + input_ids=input_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + **kwargs, + ) + + hidden_states = outputs.last_hidden_state + logits = self.lm_head(hidden_states).float() + + return Glm4MoeLiteCausalLMOutput(logits=logits) + + +# Register with AutoModelForCausalLMFactory +AutoModelForCausalLMFactory.register_custom_model_cls("Glm4MoeLiteConfig", Glm4MoeLiteForCausalLM) diff --git a/tensorrt_llm/_torch/auto_deploy/models/patches/deepseek.py b/tensorrt_llm/_torch/auto_deploy/models/patches/deepseek.py deleted file mode 100644 index f30bc0c6fac5..000000000000 --- a/tensorrt_llm/_torch/auto_deploy/models/patches/deepseek.py +++ /dev/null @@ -1,185 +0,0 @@ -import types -import warnings -from typing import Dict, Optional - -import torch -import torch.utils.checkpoint -from transformers import AutoModelForCausalLM -from transformers.cache_utils import Cache - - -def deepseek_v3_attention( - self, - hidden_states: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.IntTensor] = None, - past_key_value: Optional[Cache] = None, - output_attentions: bool = False, - use_cache: bool = False, - **kwargs, -): - """DeepSeekV3Attention forward function rewritten to wrap MultiheadLatentAttention as a custom op.""" - if "padding_mask" in kwargs: - warnings.warn( - "Passing `padding_mask` is deprecated and will be removed in v4.37. " - "Please make sure use `attention_mask` instead.`" - ) - bsz, q_len, _ = hidden_states.size() - - # If else paths are determined by config.json - if self.q_lora_rank is None: - q = self.q_proj(hidden_states) - else: - q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))) - q = q.view(bsz, q_len, self.num_heads, self.q_head_dim).transpose(1, 2) - q_nope, q_pe = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) - - compressed_kv = self.kv_a_proj_with_mqa(hidden_states) - compressed_kv, k_pe = torch.split( - compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 - ) - k_pe = k_pe.view(bsz, q_len, 1, self.qk_rope_head_dim).transpose(1, 2) - kv = ( - self.kv_b_proj(self.kv_a_layernorm(compressed_kv)) - .view(bsz, q_len, self.num_heads, self.qk_nope_head_dim + self.v_head_dim) - .transpose(1, 2) - ) - _, value_states = torch.split(kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1) - kv_seq_len = value_states.shape[-2] - if past_key_value is not None: - raise ValueError("past_key_value is not supported") - cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len) - - # Use custom op to capture mla. This does not handle KV cache - # as passing transformers Cache into a custom op is throwing an error. - # Would not be an issue, cause we intend to replace mla op with our implementation further along the pipeline - attn_output = torch.ops.auto_deploy.torch_attention_deepseek_fused_mla( - q_nope, - q_pe, - kv, - k_pe, - cos, - sin, - position_ids, - attention_mask, - self.softmax_scale, - ) - - attn_output = attn_output.transpose(1, 2).contiguous() - attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.v_head_dim) - attn_output = self.o_proj(attn_output) - - return attn_output, None, past_key_value - - -# This patched module matches exactly with HF generate -@torch.inference_mode() -def deepseek_v3_moe_exact(self, hidden_states): - """DeepSeekV3MoE forward function rewritten to enable torch export. - - This custom implementation matches exactly with the deepseek implementation. There are - some errors in the output tensors when the index_add based implementation is used, leading - to some mismatch in the outputs for some prompts. This ensures exact match between HF output - without custom patch and with custom patch. - """ - identity = hidden_states - batch_size, sequence_length, hidden_dim = hidden_states.shape - - selected_experts, routing_weights, *_ = self.gate(hidden_states) - - hidden_states = hidden_states.view(-1, hidden_dim) - idxs = torch.argsort(selected_experts.view(-1), stable=True) - - expert_mask = torch.nn.functional.one_hot( - selected_experts, num_classes=self.experts_per_rank - ).permute(2, 1, 0) - outputs = [] - for expert_idx in range(len(self.experts)): - expert_layer = self.experts[expert_idx] - _, top_x = torch.where(expert_mask[expert_idx]) - # Sort the top_xs and idx - sorted, _ = torch.sort(top_x) - tokens_for_this_expert = hidden_states[None, sorted].reshape(-1, hidden_dim) - expert_out = expert_layer(tokens_for_this_expert) - outputs.append(expert_out) - - outs = torch.cat(outputs, dim=0) - # Wrap torch.zeros() in a custom op to fix meta device issue during inference. - new_x = torch.zeros( - (*selected_experts.view(-1).shape, hidden_dim), - device=selected_experts.device, - dtype=outs.dtype, - ) - new_x[idxs] = outs - final_hidden_states = ( - new_x.view(*selected_experts.shape, -1) - .type(routing_weights.dtype) - .mul_(routing_weights.unsqueeze(-1)) - .sum(dim=1) - .type(new_x.dtype) - ) - final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim) - - if self.config.n_shared_experts is not None: - final_hidden_states = final_hidden_states + self.shared_experts(identity) - - return final_hidden_states.to(hidden_states.dtype) - - -@torch.inference_mode() -def deepseek_v3_moe(self, hidden_states): - """DeepSeekV3MoE forward function rewritten in Mixtral style to enable torch export.""" - - selected_experts, routing_weights, *_ = self.gate(hidden_states) - final_hidden_states = torch.ops.auto_deploy.torch_moe( - hidden_states, - selected_experts, - routing_weights, - w1_weight=[expert.gate_proj.weight for expert in self.experts], - w2_weight=[expert.down_proj.weight for expert in self.experts], - w3_weight=[expert.up_proj.weight for expert in self.experts], - ) - - if self.config.n_shared_experts is not None: - final_hidden_states = final_hidden_states + self.shared_experts(hidden_states) - - return final_hidden_states.to(hidden_states.dtype) - - -def deepseek_v3_rope(self, x, seq_len=None): - """DeepSeekV3 Rotary Embedding forward function rewritten to enable torch export. - We return the full cached cos and sin values, instead of slicing them based on seq_len as this - would cause an issue during the generate phase (when seq_len=1 from input_ids). We also move the cos - sin buffers to appropriate device to enable export. - """ - - return ( - self.cos_cached.to(dtype=x.dtype).to(device=x.device), - self.sin_cached.to(dtype=x.dtype).to(device=x.device), - ) - - -_from_config_original = AutoModelForCausalLM.from_config - -CUSTOM_MODULE_PATCHES: Dict[str, callable] = { - "DeepseekV3MoE": deepseek_v3_moe, - "DeepseekV2MoE": deepseek_v3_moe, - "DeepseekV3RotaryEmbedding": deepseek_v3_rope, - "DeepseekV3YarnRotaryEmbedding": deepseek_v3_rope, - "DeepseekV2RotaryEmbedding": deepseek_v3_rope, - "DeepseekV2YarnRotaryEmbedding": deepseek_v3_rope, -} - - -def get_model_from_config_patched(config, **kwargs): - model = _from_config_original(config, **kwargs) - # Patch modules - for _, module in model.named_modules(): - if type(module).__name__ in CUSTOM_MODULE_PATCHES.keys(): - module.forward = types.MethodType(CUSTOM_MODULE_PATCHES[type(module).__name__], module) - - return model - - -# TODO: figure out how this can be incorporated into the export patch system -AutoModelForCausalLM.from_config = get_model_from_config_patched diff --git a/tensorrt_llm/_torch/auto_deploy/models/quant_config_reader.py b/tensorrt_llm/_torch/auto_deploy/models/quant_config_reader.py index 5ec9d69627ba..898b46c76910 100644 --- a/tensorrt_llm/_torch/auto_deploy/models/quant_config_reader.py +++ b/tensorrt_llm/_torch/auto_deploy/models/quant_config_reader.py @@ -83,7 +83,7 @@ def has(cls, reader_cls: str) -> bool: @QuantConfigReaderRegistry.register("modelopt") class ModelOPTQuantConfigReader(QuantConfigReader): - _ALWAYS_EXCLUDE = ("lm_head", "model.embed_tokens", "*.mixer.gate*") + _ALWAYS_EXCLUDE = ("lm_head", "model.embed_tokens", "*.mixer.gate*", "*.mlp.gate") def read_config(self, config: Dict) -> Dict: producer = config.get("producer", {}).get("name") diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/fused_moe.py b/tensorrt_llm/_torch/auto_deploy/transform/library/fused_moe.py index 94dc64892ca9..bde3c4844329 100644 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/fused_moe.py +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/fused_moe.py @@ -1659,6 +1659,7 @@ def _extract_op_args(node): "w3_weight", "w1_input_scale", "w2_input_scale", + "w3_input_scale", "w1_weight_scale", "w2_weight_scale", "w3_weight_scale", @@ -1680,20 +1681,52 @@ def _prepare_args_cutlass_format_nvfp4(): if is_gated_mlp: # For gated MLP, concatenate w1 and w3 as [w3, w1] fc1_expert_weights = torch.cat([w3_stacked, w1_stacked], dim=1).contiguous() - # Expect w3 input scale and alpha to be the same as w1 - fc1_act_scale = w1_input_scale_stacked - fc1_alpha_stacked = w1_alpha_stacked fc1_weight_blockscale_fp8_stacked = torch.cat( [w3_weight_blockscale_fp8_stacked, w1_weight_blockscale_fp8_stacked], dim=1 ).contiguous() + + # Check if all w1 and w3 input scales are identical across experts + all_scales_equal = ( + torch.all(w1_input_scale_stacked == w1_input_scale_stacked[0]) + and torch.all(w3_input_scale_stacked == w3_input_scale_stacked[0]) + and torch.all(w1_input_scale_stacked == w3_input_scale_stacked) + ) + + if all_scales_equal: + # All scales are identical, no need for min() or alpha recomputation + fc1_act_scale = w1_input_scale_stacked[0] + fc1_alpha_stacked = w1_alpha_stacked + else: + # Scales differ across experts - use global min scale and recompute alpha + # Use min() because scales are in kernel format (2688/amax): smaller scale = larger amax. + fc1_act_scale = torch.minimum( + w1_input_scale_stacked.min(), w3_input_scale_stacked.min() + ) + # Recompute alpha using global input scale instead of per-expert input scale. + # Formula: new_alpha = old_alpha * per_expert_input_scale / global_input_scale + # This ensures alpha is consistent with the global fc1_act_scale used by the kernel. + fc1_alpha_stacked = w1_alpha_stacked * w1_input_scale_stacked / fc1_act_scale else: fc1_expert_weights = w1_stacked - fc1_act_scale = w1_input_scale_stacked - fc1_alpha_stacked = w1_alpha_stacked fc1_weight_blockscale_fp8_stacked = w1_weight_blockscale_fp8_stacked + # Check if all w1 input scales are identical across experts + all_scales_equal = torch.all(w1_input_scale_stacked == w1_input_scale_stacked[0]) + + if all_scales_equal: + # All scales are identical, no need for min() or alpha recomputation + fc1_act_scale = w1_input_scale_stacked[0] + fc1_alpha_stacked = w1_alpha_stacked + else: + # Scales differ across experts - use global min scale and recompute alpha + fc1_act_scale = w1_input_scale_stacked.min() + fc1_alpha_stacked = w1_alpha_stacked * w1_input_scale_stacked / fc1_act_scale + fc2_expert_weights = w2_stacked + # Keep fc2_act_scale per-expert (no global scale aggregation for fc2) fc2_act_scale = w2_input_scale_stacked + # No alpha recomputation needed since fc2 uses per-expert input scales + fc2_alpha_stacked = w2_alpha_stacked fc2_weight_blockscale_fp8_stacked = w2_weight_blockscale_fp8_stacked new_key_fc1_expert_weights = f"nvfp4_moe_w3_w1_stacked_{fused_key_counter}" @@ -1764,7 +1797,7 @@ def _prepare_args_cutlass_format_nvfp4(): _register_parameter(gm, new_key_fc1_act_scale, fc1_act_scale) _register_parameter(gm, new_key_fc2_act_scale, fc2_act_scale) _register_parameter(gm, new_key_fc1_alpha, fc1_alpha_stacked) - _register_parameter(gm, new_key_fc2_alpha, w2_alpha_stacked) + _register_parameter(gm, new_key_fc2_alpha, fc2_alpha_stacked) with graph.inserting_before(node): args = ( @@ -1804,6 +1837,7 @@ def _prepare_args_cutlass_format_nvfp4(): w3_list, w1_input_scale, w2_input_scale, + w3_input_scale, w1_weight_scale, w2_weight_scale, w3_weight_scale, @@ -1822,6 +1856,12 @@ def _prepare_args_cutlass_format_nvfp4(): # Scales are buffers, not parameters w1_input_scale_stacked = _stack(w1_input_scale, dim=0) w2_input_scale_stacked = _stack(w2_input_scale, dim=0) + w3_input_scale_stacked = _stack( + w3_input_scale, + dim=0, + device=w1_input_scale_stacked.device, + dtype=w1_input_scale_stacked.dtype, + ) # Use .view() not .to() to reinterpret bytes as float8, not value conversion w1_weight_blockscale_fp8_stacked = _stack(w1_weight_scale, dim=0).view(torch.float8_e4m3fn) diff --git a/tests/integration/defs/accuracy/test_llm_api_autodeploy.py b/tests/integration/defs/accuracy/test_llm_api_autodeploy.py index 3ee8eb2d62c5..ee9934cda3ca 100644 --- a/tests/integration/defs/accuracy/test_llm_api_autodeploy.py +++ b/tests/integration/defs/accuracy/test_llm_api_autodeploy.py @@ -411,3 +411,70 @@ def test_fp4(self, world_size): task.evaluate(llm, sampling_params=sampling_params) task = GSM8K(self.MODEL_NAME) task.evaluate(llm) + + +class TestGLM4Flash(LlmapiAccuracyTestHarness): + """Accuracy regression tests for GLM-4.7-Flash. + + TODO: enable in CI, see https://github.com/NVIDIA/TensorRT-LLM/issues/11117 + + In the meantime, you should run this test locally: + + ``` + cd tests/integration/defs + TRTLLM_ACCURACY_NO_REFERENCE=1 pytest -svv "accuracy/test_llm_api_autodeploy.py::TestGLM4Flash::test_auto_dtype[True]" + ``` + """ + + MODEL_NAME = "zai-org/GLM-4.7-Flash" + MODEL_PATH = MODEL_NAME # Model is in HF_CACHE + # Set minimum possible seq len + small buffer, for test speed & memory usage + MAX_SEQ_LEN = max(MMLU.MAX_INPUT_LEN + MMLU.MAX_OUTPUT_LEN, + GSM8K.MAX_INPUT_LEN + GSM8K.MAX_OUTPUT_LEN) + MAX_NUM_TOKENS = MAX_SEQ_LEN + + def get_default_kwargs(self, enable_chunked_prefill=False): + config = { + "skip_tokenizer_init": False, + "trust_remote_code": True, + "compile_backend": "torch-cudagraph", + "max_batch_size": 128, + "max_seq_len": self.MAX_SEQ_LEN, + "max_num_tokens": self.MAX_NUM_TOKENS, + "skip_loading_weights": False, + "disable_overlap_scheduler": False, + "cuda_graph_batch_sizes": [1, 2, 4, 8, 16, 32, 64, 128], + "kv_cache_config": { + "enable_block_reuse": False, + "free_gpu_memory_fraction": 0.88 + }, + "model_kwargs": { + "torch_dtype": "bfloat16" + }, + } + if enable_chunked_prefill: + config["enable_chunked_prefill"] = True + config[ + "max_num_tokens"] = 512 # NOTE: must be > max(tokens_per_block, max_batch_size) + return config + + def get_default_sampling_params(self): + eos_id = -1 + beam_width = 1 + return SamplingParams(end_id=eos_id, + pad_id=eos_id, + n=beam_width, + use_beam_search=beam_width > 1) + + @pytest.mark.skip_less_device_memory(32000) + @pytest.mark.parametrize("enable_chunked_prefill", [True, False]) + def test_auto_dtype(self, enable_chunked_prefill): + kwargs = self.get_default_kwargs(enable_chunked_prefill) + sampling_params = self.get_default_sampling_params() + with AutoDeployLLM(model=self.MODEL_PATH, + tokenizer=self.MODEL_PATH, + **kwargs) as llm: + task = MMLU(self.MODEL_NAME) + task.evaluate(llm, sampling_params=sampling_params) + task = GSM8K(self.MODEL_NAME) + task.evaluate(llm) diff --git a/tests/unittest/_torch/auto_deploy/_utils_test/_model_test_utils.py b/tests/unittest/_torch/auto_deploy/_utils_test/_model_test_utils.py index d585e4e08855..e1922ae751bd 100644 --- a/tests/unittest/_torch/auto_deploy/_utils_test/_model_test_utils.py +++ b/tests/unittest/_torch/auto_deploy/_utils_test/_model_test_utils.py @@ -400,7 +400,8 @@ def apply_rotary_pos_emb_ds(q, k, cos, sin, position_ids, unsqueeze_dim=1): "num_hidden_layers": 2, "hidden_size": 32, "intermediate_size": 64, - "kv_lora_rank": 128, + "kv_lora_rank": 512, # NOTE: must be 512 (default) for flashinfer_mla to work + "qk_rope_head_dim": 64, # NOTE: must be 64 (default) for flashinfer_mla to work "moe_intermediate_size": 128, "n_group": 2, "topk_group": 2, diff --git a/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_flashinfer_mla_op.py b/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_flashinfer_mla_op.py new file mode 100644 index 000000000000..c9020ffe49cd --- /dev/null +++ b/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_flashinfer_mla_op.py @@ -0,0 +1,2173 @@ +"""Test FlashInfer MLA backend operations. + +Tests the flashinfer_mla_with_cache cached op and compares it with the +torch_backend_mla_with_cache reference implementation. + +Key features tested: +- 5 tensor arguments: q_nope, q_pe, compressed_kv, kpe, kv_b_proj_weight +- Paged caches: ckv_cache [num_pages, page_size, kv_lora_rank] and kpe_cache [num_pages, page_size, qk_rope_head_dim] +- Prefill: Expand compressed_kv, compute normal attention via BatchPrefillWithRaggedKVCacheWrapper +- Decode: BatchMLAPagedAttentionWrapper with paged compressed KV cache + +Reference: https://docs.flashinfer.ai/api/mla.html +""" + +import flashinfer +import pytest +import torch + +import tensorrt_llm._torch.auto_deploy # noqa: F401 +from tensorrt_llm._torch.auto_deploy.custom_ops.mla.flashinfer_mla import ( + _GlobalFlashInferMLAPlanner, +) +from tensorrt_llm._torch.auto_deploy.utils.cuda_graph import CudaGraphWarmUpPhase + + +def _create_mla_inputs( + batch_size: int, + seq_len: int, + num_heads: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + kv_lora_rank: int, + v_head_dim: int, + dtype: torch.dtype, + device: str, +): + """Create MLA input tensors. + + Args: + batch_size: Batch size + seq_len: Sequence length + num_heads: Number of attention heads + qk_nope_head_dim: Dimension of query/key non-positional part + qk_rope_head_dim: Dimension of query/key positional (RoPE) part + kv_lora_rank: Rank of compressed KV (LoRA rank) + v_head_dim: Dimension of value head + dtype: Data type + device: Device + + Returns: + Dictionary with q_nope, q_pe, compressed_kv, kpe, kv_b_proj_weight + """ + kv_head_dim = qk_nope_head_dim + v_head_dim + + # Scale factor for Xavier-like initialization to keep values bounded + # This helps reduce numerical differences by keeping output magnitudes smaller + q_scale = 1.0 / (qk_nope_head_dim**0.5) + kv_scale = 1.0 / (kv_lora_rank**0.5) + + # q_nope: [B, S, N, qk_nope_head_dim] + q_nope = ( + torch.randn(batch_size, seq_len, num_heads, qk_nope_head_dim, dtype=dtype, device=device) + * q_scale + ) + + # q_pe: [B, S, N, qk_rope_head_dim] + q_pe = ( + torch.randn(batch_size, seq_len, num_heads, qk_rope_head_dim, dtype=dtype, device=device) + * q_scale + ) + + # compressed_kv: [B, S, kv_lora_rank] + compressed_kv = ( + torch.randn(batch_size, seq_len, kv_lora_rank, dtype=dtype, device=device) * kv_scale + ) + + # kpe: [B, S, 1, qk_rope_head_dim] + kpe = ( + torch.randn(batch_size, seq_len, 1, qk_rope_head_dim, dtype=dtype, device=device) * q_scale + ) + + # kv_b_proj_weight: [N * (qk_nope_head_dim + v_head_dim), kv_lora_rank] + # Xavier initialization for the projection weight + weight_scale = 1.0 / (kv_lora_rank**0.5) + kv_b_proj_weight = ( + torch.randn(num_heads * kv_head_dim, kv_lora_rank, dtype=dtype, device=device) + * weight_scale + ) + + return { + "q_nope": q_nope, + "q_pe": q_pe, + "compressed_kv": compressed_kv, + "kpe": kpe, + "kv_b_proj_weight": kv_b_proj_weight, + } + + +def _create_unpaged_cache_and_metadata( + batch_size: int, + max_seq_len: int, + kv_lora_rank: int, + qk_rope_head_dim: int, + dtype: torch.dtype, + device: str, + seq_lengths: list, + input_positions: list, +): + """Create unpaged (torch backend) cache and metadata. + + Args: + batch_size: Batch size + max_seq_len: Maximum sequence length + kv_lora_rank: Rank of compressed KV + qk_rope_head_dim: Dimension of RoPE + dtype: Data type + device: Device + seq_lengths: List of sequence lengths per batch + input_positions: List of input positions (cache offsets) per batch + + Returns: + Dictionary with cache and metadata tensors + """ + # FlashInfer MLA cache (unpaged): [max_batch, max_seq, kv_lora_rank + qk_rope_head_dim] + mla_cache = torch.zeros( + batch_size, max_seq_len, kv_lora_rank + qk_rope_head_dim, dtype=dtype, device=device + ) + + # Metadata + seq_len_tensor = torch.tensor(seq_lengths, dtype=torch.int32, device=device) + input_pos = torch.tensor(input_positions, dtype=torch.int32, device=device) + slot_idx = torch.arange(batch_size, dtype=torch.int32, device=device) + + # Compute cu_seqlen (cumulative sequence lengths) + cu_seqlen = torch.zeros(batch_size + 1, dtype=torch.int32, device=device) + cu_seqlen[1:] = torch.cumsum(seq_len_tensor, dim=0) + + # Determine if this is context (prefill) or generate (decode) + total_tokens = sum(seq_lengths) + is_decode = all(s == 1 for s in seq_lengths) + + if is_decode: + # Decode phase + batch_info_host = torch.tensor([0, 0, batch_size], dtype=torch.int32, device=device) + else: + # Context/prefill phase + batch_info_host = torch.tensor( + [batch_size, total_tokens, 0], dtype=torch.int32, device=device + ) + + return { + "mla_cache": mla_cache, + "batch_info_host": batch_info_host, + "seq_len": seq_len_tensor, + "input_pos": input_pos, + "slot_idx": slot_idx, + "cu_seqlen": cu_seqlen[:-1], # Exclude last element for seq_start + } + + +def _create_paged_cache_and_metadata( + batch_size: int, + max_num_pages: int, + page_size: int, + kv_lora_rank: int, + qk_rope_head_dim: int, + dtype: torch.dtype, + device: str, + seq_lengths: list, + input_positions: list, +): + """Create paged (flashinfer backend) cache and metadata. + + Args: + batch_size: Batch size + max_num_pages: Maximum number of pages + page_size: Size of each page + kv_lora_rank: Rank of compressed KV + qk_rope_head_dim: Dimension of RoPE + dtype: Data type + device: Device + seq_lengths: List of sequence lengths per batch + input_positions: List of input positions (cache offsets) per batch + + Returns: + Dictionary with paged cache and metadata tensors + """ + # Paged MLA caches (two separate caches) + ckv_cache = torch.zeros(max_num_pages, page_size, kv_lora_rank, dtype=dtype, device=device) + kpe_cache = torch.zeros(max_num_pages, page_size, qk_rope_head_dim, dtype=dtype, device=device) + + # Compute total KV lengths (input_pos + seq_len for each sequence) + kv_lengths = [pos + seq_len for pos, seq_len in zip(input_positions, seq_lengths)] + + # Compute number of pages per sequence + pages_per_seq = [(kv_len - 1) // page_size + 1 if kv_len > 0 else 1 for kv_len in kv_lengths] + + # Assign pages (simple sequential assignment) + page_assignments = [] + next_page = 0 + for num_pages in pages_per_seq: + page_assignments.append(list(range(next_page, next_page + num_pages))) + next_page += num_pages + + # Create FlashInfer paging metadata + seq_len_tensor = torch.tensor(seq_lengths, dtype=torch.int32, device=device) + + # qo_indptr: cumulative query/output lengths + qo_indptr = torch.zeros(batch_size + 1, dtype=torch.int32, device=device) + qo_indptr[1:] = torch.cumsum(seq_len_tensor, dim=0) + + # cu_num_pages: cumulative number of pages per sequence + num_pages_per_seq = torch.tensor(pages_per_seq, dtype=torch.int32, device=device) + cu_num_pages = torch.zeros(batch_size + 1, dtype=torch.int32, device=device) + cu_num_pages[1:] = torch.cumsum(num_pages_per_seq, dim=0) + + # cache_loc (paged_kv_indices): flattened list of page indices + cache_loc = torch.tensor( + [p for pages in page_assignments for p in pages], dtype=torch.int32, device=device + ) + + # last_page_len: number of valid tokens in the last page of each sequence + last_page_len = torch.tensor( + [((kv_len - 1) % page_size) + 1 if kv_len > 0 else 0 for kv_len in kv_lengths], + dtype=torch.int32, + device=device, + ) + + # seq_len_with_cache: total KV lengths + seq_len_with_cache = torch.tensor(kv_lengths, dtype=torch.int32, device=device) + + # Host copies + qo_indptr_host = qo_indptr.cpu() + cu_num_pages_host = cu_num_pages.cpu() + last_page_len_host = last_page_len.cpu() + seq_len_with_cache_host = seq_len_with_cache.cpu() + + # Determine if this is context (prefill) or generate (decode) + total_tokens = sum(seq_lengths) + is_decode = all(s == 1 for s in seq_lengths) + + if is_decode: + # Decode phase + batch_info_host = torch.tensor([0, 0, batch_size], dtype=torch.int32, device=device) + else: + # Context/prefill phase + batch_info_host = torch.tensor( + [batch_size, total_tokens, 0], dtype=torch.int32, device=device + ) + + return { + "ckv_cache": ckv_cache, + "kpe_cache": kpe_cache, + "batch_info_host": batch_info_host, + "cu_seqlen_host": qo_indptr_host, + "cu_num_pages": cu_num_pages, + "cu_num_pages_host": cu_num_pages_host, + "cache_loc": cache_loc, + "last_page_len": last_page_len, + "last_page_len_host": last_page_len_host, + "seq_len_with_cache_host": seq_len_with_cache_host, + "page_size": page_size, + } + + +def _copy_unpaged_to_paged_cache( + unpaged_cache: torch.Tensor, + ckv_cache: torch.Tensor, + kpe_cache: torch.Tensor, + batch_size: int, + tokens_per_seq: list, + page_size: int, + cu_num_pages: torch.Tensor, + cache_loc: torch.Tensor, + kv_lora_rank: int, +): + """Copy unpaged cache data to paged cache format. + + This is used to initialize paged cache with the same data as unpaged cache + for comparison tests. + + Args: + unpaged_cache: Source cache [batch, max_seq, dim] + ckv_cache: Destination paged ckv cache [num_pages, page_size, kv_lora_rank] + kpe_cache: Destination paged kpe cache [num_pages, page_size, qk_rope_head_dim] + batch_size: Number of sequences + tokens_per_seq: Number of tokens to copy per sequence + page_size: Number of tokens per page + cu_num_pages: Cumulative page counts from flashinfer metadata [batch_size + 1] + cache_loc: Page indices from flashinfer metadata + kv_lora_rank: Rank of compressed KV (split dimension) + """ + for batch_idx in range(batch_size): + num_tokens = tokens_per_seq[batch_idx] + if num_tokens == 0: + continue + + # Get page assignments for this sequence from flashinfer metadata + page_start_idx = cu_num_pages[batch_idx].item() + page_end_idx = cu_num_pages[batch_idx + 1].item() + + token_offset = 0 + for i in range(page_start_idx, page_end_idx): + page_num = cache_loc[i].item() + tokens_to_copy = min(page_size, num_tokens - token_offset) + if tokens_to_copy <= 0: + break + + # Split unpaged cache into ckv and kpe portions + unpaged_data = unpaged_cache[batch_idx, token_offset : token_offset + tokens_to_copy] + ckv_cache[page_num, :tokens_to_copy] = unpaged_data[:, :kv_lora_rank] + kpe_cache[page_num, :tokens_to_copy] = unpaged_data[:, kv_lora_rank:] + token_offset += tokens_to_copy + + +@pytest.mark.parametrize("seq_length", [32, 128]) +@pytest.mark.parametrize("num_heads", [1]) +@pytest.mark.parametrize("batch_size", [8, 64]) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@pytest.mark.parametrize("device", ["cuda"]) +def test_flashinfer_mla_op_context(seq_length, num_heads, batch_size, dtype, device): + """Test FlashInfer MLA context (prefill) phase. + + Compares flashinfer_mla_with_cache against torch_backend_mla_with_cache + for context (prefill) operations where seq_length > 1. + """ + # MLA dimensions (similar to DeepSeek) + qk_nope_head_dim = 128 + qk_rope_head_dim = 64 # Must be 64 for FlashInfer MLA. + kv_lora_rank = 512 # Must be 512 for FlashInfer MLA. + v_head_dim = 128 + page_size = seq_length # Use seq_length as page size for simpler comparison + + max_seq_len = 256 + max_num_pages = batch_size * (max_seq_len // page_size + 1) + + # Create input tensors + inputs = _create_mla_inputs( + batch_size, + seq_length, + num_heads, + qk_nope_head_dim, + qk_rope_head_dim, + kv_lora_rank, + v_head_dim, + dtype, + device, + ) + + # Flatten inputs for context phase + total_tokens = batch_size * seq_length + q_nope_flat = inputs["q_nope"].view(1, total_tokens, num_heads, qk_nope_head_dim) + q_pe_flat = inputs["q_pe"].view(1, total_tokens, num_heads, qk_rope_head_dim) + compressed_kv_flat = inputs["compressed_kv"].view(1, total_tokens, kv_lora_rank) + kpe_flat = inputs["kpe"].view(1, total_tokens, 1, qk_rope_head_dim) + + # Sequence lengths and positions + seq_lengths = [seq_length] * batch_size + input_positions = [0] * batch_size # Context starts at position 0 + + # ========================================================================= + # Run torch backend (reference) + # ========================================================================= + torch_meta = _create_unpaged_cache_and_metadata( + batch_size, + max_seq_len, + kv_lora_rank, + qk_rope_head_dim, + dtype, + device, + seq_lengths, + input_positions, + ) + + torch_output = torch.ops.auto_deploy.torch_cached_mla_with_cache( + q_nope_flat, + q_pe_flat, + compressed_kv_flat, + kpe_flat, + inputs["kv_b_proj_weight"], + torch_meta["batch_info_host"], + torch_meta["seq_len"], + torch_meta["input_pos"], + torch_meta["slot_idx"], + torch_meta["cu_seqlen"], + torch_meta["mla_cache"], + None, # scale + kv_lora_rank, + ) + + # ========================================================================= + # Run FlashInfer backend + # ========================================================================= + # Reset planner + _GlobalFlashInferMLAPlanner.reset(torch.device(device)) + + # Create paged metadata + flashinfer_meta = _create_paged_cache_and_metadata( + batch_size, + max_num_pages, + page_size, + kv_lora_rank, + qk_rope_head_dim, + dtype, + device, + seq_lengths, + input_positions, + ) + + # Compute FlashInfer batch indices and positions + qo_indptr = torch.zeros(batch_size + 1, dtype=torch.int32, device=device) + qo_indptr[1:] = torch.cumsum(torch.tensor(seq_lengths, device=device), dim=0).int() + + batch_indices, positions = flashinfer.get_batch_indices_positions( + qo_indptr, + flashinfer.get_seq_lens( + flashinfer_meta["cu_num_pages"], + flashinfer_meta["last_page_len"], + page_size=page_size, + ), + total_tokens, + ) + + flashinfer_output = torch.ops.auto_deploy.flashinfer_mla_with_cache( + inputs["q_nope"], + inputs["q_pe"], + inputs["compressed_kv"], + inputs["kpe"], + inputs["kv_b_proj_weight"], + flashinfer_meta["batch_info_host"], + flashinfer_meta["cu_seqlen_host"], + flashinfer_meta["cu_num_pages"], + flashinfer_meta["cu_num_pages_host"], + flashinfer_meta["cache_loc"], + flashinfer_meta["last_page_len"], + flashinfer_meta["last_page_len_host"], + flashinfer_meta["seq_len_with_cache_host"], + batch_indices, + positions, + flashinfer_meta["ckv_cache"], + flashinfer_meta["kpe_cache"], + None, # scale + kv_lora_rank, + ) + + # ========================================================================= + # Compare outputs + # ========================================================================= + # Reshape for comparison + torch_output_reshaped = torch_output.view(batch_size, seq_length, num_heads, v_head_dim) + flashinfer_output_reshaped = flashinfer_output.view( + batch_size, seq_length, num_heads, v_head_dim + ) + + # FlashInfer uses fused kernels with different computation order/precision than the + # torch reference. With bfloat16 and scaled inputs, tighter tolerances are achievable. + assert torch.allclose( + flashinfer_output_reshaped.cpu().to(torch.float32), + torch_output_reshaped.cpu().to(torch.float32), + atol=0.05, + rtol=0.02, + ), ( + f"FlashInfer MLA context output doesn't match torch backend. " + f"Max diff: {(flashinfer_output_reshaped - torch_output_reshaped).abs().max():.6f}" + ) + + +@pytest.mark.parametrize("prefill_seq_length", [64, 128]) +@pytest.mark.parametrize("num_heads", [1]) +@pytest.mark.parametrize("batch_size", [4, 64]) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@pytest.mark.parametrize("device", ["cuda"]) +def test_flashinfer_mla_op_decode(prefill_seq_length, num_heads, batch_size, dtype, device): + """Test FlashInfer MLA decode (generate) phase. + + Compares flashinfer_mla_with_cache against torch_backend_mla_with_cache + for decode operations where seq_length = 1. + """ + # MLA dimensions + qk_nope_head_dim = 128 + qk_rope_head_dim = 64 # Must be 64 for FlashInfer MLA. + kv_lora_rank = 512 # Must be 512 for FlashInfer MLA. + v_head_dim = 192 + page_size = 64 + + seq_length = 1 # Decode phase + + max_seq_len = 256 + max_num_pages = batch_size * (max_seq_len // page_size + 1) + + # Create input tensors + inputs = _create_mla_inputs( + batch_size, + seq_length, + num_heads, + qk_nope_head_dim, + qk_rope_head_dim, + kv_lora_rank, + v_head_dim, + dtype, + device, + ) + + # Sequence lengths and positions + seq_lengths = [seq_length] * batch_size + input_positions = [prefill_seq_length] * batch_size + + # ========================================================================= + # Setup caches with pre-filled data + # ========================================================================= + # Create unpaged cache with prefilled data + torch_meta = _create_unpaged_cache_and_metadata( + batch_size, + max_seq_len, + kv_lora_rank, + qk_rope_head_dim, + dtype, + device, + seq_lengths, + input_positions, + ) + + # Pre-fill cache with random data if prefill_seq_length > 0 + if prefill_seq_length > 0: + torch_meta["mla_cache"][:, :prefill_seq_length, :] = torch.randn( + batch_size, + prefill_seq_length, + kv_lora_rank + qk_rope_head_dim, + dtype=dtype, + device=device, + ) + + # Create paged cache + flashinfer_meta = _create_paged_cache_and_metadata( + batch_size, + max_num_pages, + page_size, + kv_lora_rank, + qk_rope_head_dim, + dtype, + device, + seq_lengths, + input_positions, + ) + + # Copy unpaged cache to paged format + if prefill_seq_length > 0: + _copy_unpaged_to_paged_cache( + torch_meta["mla_cache"], + flashinfer_meta["ckv_cache"], + flashinfer_meta["kpe_cache"], + batch_size, + [prefill_seq_length] * batch_size, # Number of tokens to copy + page_size, + flashinfer_meta["cu_num_pages"], + flashinfer_meta["cache_loc"], + kv_lora_rank, + ) + + # ========================================================================= + # Run torch backend (reference) + # ========================================================================= + torch_output = torch.ops.auto_deploy.torch_cached_mla_with_cache( + inputs["q_nope"], + inputs["q_pe"], + inputs["compressed_kv"], + inputs["kpe"], + inputs["kv_b_proj_weight"], + torch_meta["batch_info_host"], + torch_meta["seq_len"], + torch_meta["input_pos"], + torch_meta["slot_idx"], + torch_meta["cu_seqlen"], + torch_meta["mla_cache"], + None, # scale + kv_lora_rank, + ) + + # ========================================================================= + # Run FlashInfer backend + # ========================================================================= + _GlobalFlashInferMLAPlanner.reset(torch.device(device)) + + # Compute FlashInfer batch indices and positions + qo_indptr = torch.zeros(batch_size + 1, dtype=torch.int32, device=device) + qo_indptr[1:] = torch.cumsum(torch.tensor(seq_lengths, device=device), dim=0).int() + + batch_indices, positions = flashinfer.get_batch_indices_positions( + qo_indptr, + flashinfer.get_seq_lens( + flashinfer_meta["cu_num_pages"], + flashinfer_meta["last_page_len"], + page_size=page_size, + ), + batch_size * seq_length, + ) + + flashinfer_output = torch.ops.auto_deploy.flashinfer_mla_with_cache( + inputs["q_nope"], + inputs["q_pe"], + inputs["compressed_kv"], + inputs["kpe"], + inputs["kv_b_proj_weight"], + flashinfer_meta["batch_info_host"], + flashinfer_meta["cu_seqlen_host"], + flashinfer_meta["cu_num_pages"], + flashinfer_meta["cu_num_pages_host"], + flashinfer_meta["cache_loc"], + flashinfer_meta["last_page_len"], + flashinfer_meta["last_page_len_host"], + flashinfer_meta["seq_len_with_cache_host"], + batch_indices, + positions, + flashinfer_meta["ckv_cache"], + flashinfer_meta["kpe_cache"], + None, # scale + kv_lora_rank, + ) + + # ========================================================================= + # Compare outputs + # ========================================================================= + # FlashInfer uses fused kernels with different computation order/precision than the + # torch reference. With bfloat16 and scaled inputs, tighter tolerances are achievable. + assert torch.allclose( + flashinfer_output.cpu().to(torch.float32), + torch_output.cpu().to(torch.float32), + atol=0.05, + rtol=0.02, + ), ( + f"FlashInfer MLA decode output doesn't match torch backend. " + f"Max diff: {(flashinfer_output - torch_output).abs().max():.6f}" + ) + + +@pytest.mark.parametrize("prefill_seq_length", [16, 128, 1024]) +@pytest.mark.parametrize("num_heads", [1, 8]) +@pytest.mark.parametrize("batch_size", [4, 64, 256]) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@pytest.mark.parametrize("device", ["cuda"]) +def test_flashinfer_mla_context_and_generate( + prefill_seq_length, num_heads, batch_size, dtype, device +): + """Test FlashInfer MLA context (prefill) followed by generate (decode). + + This test verifies the full workflow: + 1. Context phase: Process initial sequence + 2. Generate phase: Generate additional tokens one at a time + """ + # MLA dimensions + qk_nope_head_dim = 128 + qk_rope_head_dim = 64 # Must be 64 for FlashInfer MLA. + kv_lora_rank = 512 # Must be 512 for FlashInfer MLA. + v_head_dim = 128 + # Use a fixed page_size of 64 for FlashInfer MLA. + page_size = 64 + + max_seq_len = 2048 + max_num_pages = batch_size * (max_seq_len // page_size + 1) + + # ========================================================================= + # Context phase + # ========================================================================= + inputs_context = _create_mla_inputs( + batch_size, + prefill_seq_length, + num_heads, + qk_nope_head_dim, + qk_rope_head_dim, + kv_lora_rank, + v_head_dim, + dtype, + device, + ) + + seq_lengths_context = [prefill_seq_length] * batch_size + input_positions_context = [0] * batch_size + + # Flatten context inputs + total_tokens = batch_size * prefill_seq_length + q_nope_flat = inputs_context["q_nope"].view(1, total_tokens, num_heads, qk_nope_head_dim) + q_pe_flat = inputs_context["q_pe"].view(1, total_tokens, num_heads, qk_rope_head_dim) + compressed_kv_flat = inputs_context["compressed_kv"].view(1, total_tokens, kv_lora_rank) + kpe_flat = inputs_context["kpe"].view(1, total_tokens, 1, qk_rope_head_dim) + + # Create caches + torch_meta = _create_unpaged_cache_and_metadata( + batch_size, + max_seq_len, + kv_lora_rank, + qk_rope_head_dim, + dtype, + device, + seq_lengths_context, + input_positions_context, + ) + + flashinfer_meta = _create_paged_cache_and_metadata( + batch_size, + max_num_pages, + page_size, + kv_lora_rank, + qk_rope_head_dim, + dtype, + device, + seq_lengths_context, + input_positions_context, + ) + + # Run torch backend context + torch_output_context = torch.ops.auto_deploy.torch_cached_mla_with_cache( + q_nope_flat, + q_pe_flat, + compressed_kv_flat, + kpe_flat, + inputs_context["kv_b_proj_weight"], + torch_meta["batch_info_host"], + torch_meta["seq_len"], + torch_meta["input_pos"], + torch_meta["slot_idx"], + torch_meta["cu_seqlen"], + torch_meta["mla_cache"], + None, + kv_lora_rank, + ) + + # Run FlashInfer context + _GlobalFlashInferMLAPlanner.reset(torch.device(device)) + + qo_indptr = torch.zeros(batch_size + 1, dtype=torch.int32, device=device) + qo_indptr[1:] = torch.cumsum(torch.tensor(seq_lengths_context, device=device), dim=0).int() + + batch_indices, positions = flashinfer.get_batch_indices_positions( + qo_indptr, + flashinfer.get_seq_lens( + flashinfer_meta["cu_num_pages"], + flashinfer_meta["last_page_len"], + page_size=page_size, + ), + total_tokens, + ) + + flashinfer_output_context = torch.ops.auto_deploy.flashinfer_mla_with_cache( + inputs_context["q_nope"], + inputs_context["q_pe"], + inputs_context["compressed_kv"], + inputs_context["kpe"], + inputs_context["kv_b_proj_weight"], + flashinfer_meta["batch_info_host"], + flashinfer_meta["cu_seqlen_host"], + flashinfer_meta["cu_num_pages"], + flashinfer_meta["cu_num_pages_host"], + flashinfer_meta["cache_loc"], + flashinfer_meta["last_page_len"], + flashinfer_meta["last_page_len_host"], + flashinfer_meta["seq_len_with_cache_host"], + batch_indices, + positions, + flashinfer_meta["ckv_cache"], + flashinfer_meta["kpe_cache"], + None, + kv_lora_rank, + ) + + # Verify context outputs match + torch_output_context_reshaped = torch_output_context.view( + batch_size, prefill_seq_length, num_heads, v_head_dim + ) + flashinfer_output_context_reshaped = flashinfer_output_context.view( + batch_size, prefill_seq_length, num_heads, v_head_dim + ) + + assert torch.allclose( + flashinfer_output_context_reshaped.cpu().to(torch.float32), + torch_output_context_reshaped.cpu().to(torch.float32), + atol=0.01, + rtol=0.01, + ), "Context phase outputs don't match" + + # ========================================================================= + # Generate phase (single token) + # ========================================================================= + inputs_gen = _create_mla_inputs( + batch_size, + 1, # seq_length = 1 for generate + num_heads, + qk_nope_head_dim, + qk_rope_head_dim, + kv_lora_rank, + v_head_dim, + dtype, + device, + ) + + seq_lengths_gen = [1] * batch_size + input_positions_gen = [prefill_seq_length] * batch_size + + # Update torch metadata for generate + torch_meta_gen = _create_unpaged_cache_and_metadata( + batch_size, + max_seq_len, + kv_lora_rank, + qk_rope_head_dim, + dtype, + device, + seq_lengths_gen, + input_positions_gen, + ) + # Use the same cache (already filled from context) + torch_meta_gen["mla_cache"] = torch_meta["mla_cache"] + + # Update flashinfer metadata for generate + flashinfer_meta_gen = _create_paged_cache_and_metadata( + batch_size, + max_num_pages, + page_size, + kv_lora_rank, + qk_rope_head_dim, + dtype, + device, + seq_lengths_gen, + input_positions_gen, + ) + # Use the same caches + flashinfer_meta_gen["ckv_cache"] = flashinfer_meta["ckv_cache"] + flashinfer_meta_gen["kpe_cache"] = flashinfer_meta["kpe_cache"] + + # Run torch backend generate + torch_output_gen = torch.ops.auto_deploy.torch_cached_mla_with_cache( + inputs_gen["q_nope"], + inputs_gen["q_pe"], + inputs_gen["compressed_kv"], + inputs_gen["kpe"], + inputs_context["kv_b_proj_weight"], # Use same weights + torch_meta_gen["batch_info_host"], + torch_meta_gen["seq_len"], + torch_meta_gen["input_pos"], + torch_meta_gen["slot_idx"], + torch_meta_gen["cu_seqlen"], + torch_meta_gen["mla_cache"], + None, + kv_lora_rank, + ) + + # Run FlashInfer generate + _GlobalFlashInferMLAPlanner.reset(torch.device(device)) + + qo_indptr_gen = torch.zeros(batch_size + 1, dtype=torch.int32, device=device) + qo_indptr_gen[1:] = torch.cumsum(torch.tensor(seq_lengths_gen, device=device), dim=0).int() + + batch_indices_gen, positions_gen = flashinfer.get_batch_indices_positions( + qo_indptr_gen, + flashinfer.get_seq_lens( + flashinfer_meta_gen["cu_num_pages"], + flashinfer_meta_gen["last_page_len"], + page_size=page_size, + ), + batch_size, + ) + + flashinfer_output_gen = torch.ops.auto_deploy.flashinfer_mla_with_cache( + inputs_gen["q_nope"], + inputs_gen["q_pe"], + inputs_gen["compressed_kv"], + inputs_gen["kpe"], + inputs_context["kv_b_proj_weight"], + flashinfer_meta_gen["batch_info_host"], + flashinfer_meta_gen["cu_seqlen_host"], + flashinfer_meta_gen["cu_num_pages"], + flashinfer_meta_gen["cu_num_pages_host"], + flashinfer_meta_gen["cache_loc"], + flashinfer_meta_gen["last_page_len"], + flashinfer_meta_gen["last_page_len_host"], + flashinfer_meta_gen["seq_len_with_cache_host"], + batch_indices_gen, + positions_gen, + flashinfer_meta_gen["ckv_cache"], + flashinfer_meta_gen["kpe_cache"], + None, + kv_lora_rank, + ) + + # Verify generate outputs match + assert torch.allclose( + flashinfer_output_gen.cpu().to(torch.float32), + torch_output_gen.cpu().to(torch.float32), + atol=0.05, + rtol=0.02, + ), ( + f"Generate phase outputs don't match. " + f"Max diff: {(flashinfer_output_gen - torch_output_gen).abs().max():.6f}" + ) + + +@pytest.mark.parametrize( + "seq_lengths", + [ + [8, 16], + [12, 24, 32], + [4, 8, 16, 32], + ], +) +@pytest.mark.parametrize("num_heads", [8]) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@pytest.mark.parametrize("device", ["cuda"]) +def test_flashinfer_mla_with_variable_seq_lengths(seq_lengths, num_heads, dtype, device): + """Test FlashInfer MLA with variable sequence lengths in a batch. + + This test verifies that the FlashInfer MLA backend handles batches + with different sequence lengths correctly. + """ + batch_size = len(seq_lengths) + + # MLA dimensions + qk_nope_head_dim = 128 + qk_rope_head_dim = 64 # Must be 64 for FlashInfer MLA. + kv_lora_rank = 512 # Must be 512 for FlashInfer MLA. + v_head_dim = 128 + page_size = 32 + + max_seq_len = 256 + max_num_pages = batch_size * (max_seq_len // page_size + 1) + + # Create individual inputs for each sequence + total_tokens = sum(seq_lengths) + + # Create batched inputs (we'll flatten later) + q_nope_list = [] + q_pe_list = [] + compressed_kv_list = [] + kpe_list = [] + + for seq_len in seq_lengths: + inputs = _create_mla_inputs( + 1, # Single sequence + seq_len, + num_heads, + qk_nope_head_dim, + qk_rope_head_dim, + kv_lora_rank, + v_head_dim, + dtype, + device, + ) + q_nope_list.append(inputs["q_nope"].squeeze(0)) + q_pe_list.append(inputs["q_pe"].squeeze(0)) + compressed_kv_list.append(inputs["compressed_kv"].squeeze(0)) + kpe_list.append(inputs["kpe"].squeeze(0)) + + # Concatenate into flattened format + q_nope_flat = torch.cat(q_nope_list, dim=0).unsqueeze(0) # [1, total_tokens, N, D] + q_pe_flat = torch.cat(q_pe_list, dim=0).unsqueeze(0) + compressed_kv_flat = torch.cat(compressed_kv_list, dim=0).unsqueeze(0) + kpe_flat = torch.cat(kpe_list, dim=0).unsqueeze(0) + + # Common kv_b_proj_weight + kv_head_dim = qk_nope_head_dim + v_head_dim + kv_b_proj_weight = torch.randn( + num_heads * kv_head_dim, kv_lora_rank, dtype=dtype, device=device + ) + + input_positions = [0] * batch_size + + # ========================================================================= + # Run FlashInfer backend + # ========================================================================= + _GlobalFlashInferMLAPlanner.reset(torch.device(device)) + + flashinfer_meta = _create_paged_cache_and_metadata( + batch_size, + max_num_pages, + page_size, + kv_lora_rank, + qk_rope_head_dim, + dtype, + device, + seq_lengths, + input_positions, + ) + + qo_indptr = torch.zeros(batch_size + 1, dtype=torch.int32, device=device) + qo_indptr[1:] = torch.cumsum(torch.tensor(seq_lengths, device=device), dim=0).int() + + batch_indices, positions = flashinfer.get_batch_indices_positions( + qo_indptr, + flashinfer.get_seq_lens( + flashinfer_meta["cu_num_pages"], + flashinfer_meta["last_page_len"], + page_size=page_size, + ), + total_tokens, + ) + + flashinfer_output = torch.ops.auto_deploy.flashinfer_mla_with_cache( + q_nope_flat, + q_pe_flat, + compressed_kv_flat, + kpe_flat, + kv_b_proj_weight, + flashinfer_meta["batch_info_host"], + flashinfer_meta["cu_seqlen_host"], + flashinfer_meta["cu_num_pages"], + flashinfer_meta["cu_num_pages_host"], + flashinfer_meta["cache_loc"], + flashinfer_meta["last_page_len"], + flashinfer_meta["last_page_len_host"], + flashinfer_meta["seq_len_with_cache_host"], + batch_indices, + positions, + flashinfer_meta["ckv_cache"], + flashinfer_meta["kpe_cache"], + None, + kv_lora_rank, + ) + + # Verify output shape + expected_shape = (1, total_tokens, num_heads, v_head_dim) + assert flashinfer_output.shape == expected_shape, ( + f"Expected shape {expected_shape}, got {flashinfer_output.shape}" + ) + + # Verify output is finite + assert torch.isfinite(flashinfer_output).all(), "Output contains NaN or Inf values" + + +@pytest.mark.parametrize( + "seq_lengths", + [ + [8, 16, 32], + [12, 24, 48, 64], + [16, 32, 64, 96, 128], + ], +) +@pytest.mark.parametrize("num_decode_steps", [3, 5]) +@pytest.mark.parametrize("num_heads", [8]) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@pytest.mark.parametrize("device", ["cuda"]) +def test_flashinfer_mla_variable_seq_multi_decode( + seq_lengths, num_decode_steps, num_heads, dtype, device +): + """Test FlashInfer MLA with variable sequence lengths and multiple decode steps. + + This test verifies the full workflow with variable sequence lengths: + 1. Context phase: Process initial sequences with different lengths + 2. Multiple decode steps: Generate multiple tokens, updating the cache each step + + Compares torch_backend_mla_with_cache against flashinfer_mla_with_cache. + """ + batch_size = len(seq_lengths) + + # MLA dimensions + qk_nope_head_dim = 128 + qk_rope_head_dim = 64 # Must be 64 for FlashInfer MLA. + kv_lora_rank = 512 # Must be 512 for FlashInfer MLA. + v_head_dim = 128 + page_size = 64 + + max_seq_len = max(seq_lengths) + num_decode_steps + 128 # Extra headroom + max_num_pages = batch_size * (max_seq_len // page_size + 2) + + # Create individual inputs for each sequence (context phase) + total_tokens = sum(seq_lengths) + + q_nope_list = [] + q_pe_list = [] + compressed_kv_list = [] + kpe_list = [] + + for seq_len in seq_lengths: + inputs = _create_mla_inputs( + 1, # Single sequence + seq_len, + num_heads, + qk_nope_head_dim, + qk_rope_head_dim, + kv_lora_rank, + v_head_dim, + dtype, + device, + ) + q_nope_list.append(inputs["q_nope"].squeeze(0)) + q_pe_list.append(inputs["q_pe"].squeeze(0)) + compressed_kv_list.append(inputs["compressed_kv"].squeeze(0)) + kpe_list.append(inputs["kpe"].squeeze(0)) + + # Concatenate into flattened format for context phase + q_nope_flat = torch.cat(q_nope_list, dim=0).unsqueeze(0) # [1, total_tokens, N, D] + q_pe_flat = torch.cat(q_pe_list, dim=0).unsqueeze(0) + compressed_kv_flat = torch.cat(compressed_kv_list, dim=0).unsqueeze(0) + kpe_flat = torch.cat(kpe_list, dim=0).unsqueeze(0) + + # Common kv_b_proj_weight + kv_head_dim = qk_nope_head_dim + v_head_dim + weight_scale = 1.0 / (kv_lora_rank**0.5) + kv_b_proj_weight = ( + torch.randn(num_heads * kv_head_dim, kv_lora_rank, dtype=dtype, device=device) + * weight_scale + ) + + input_positions_context = [0] * batch_size + + # ========================================================================= + # Context phase - Setup both backends + # ========================================================================= + + # Create torch unpaged cache + torch_meta = _create_unpaged_cache_and_metadata( + batch_size, + max_seq_len, + kv_lora_rank, + qk_rope_head_dim, + dtype, + device, + seq_lengths, + input_positions_context, + ) + + # Create flashinfer paged cache + flashinfer_meta = _create_paged_cache_and_metadata( + batch_size, + max_num_pages, + page_size, + kv_lora_rank, + qk_rope_head_dim, + dtype, + device, + seq_lengths, + input_positions_context, + ) + + # Run torch backend context + torch_output_context = torch.ops.auto_deploy.torch_cached_mla_with_cache( + q_nope_flat, + q_pe_flat, + compressed_kv_flat, + kpe_flat, + kv_b_proj_weight, + torch_meta["batch_info_host"], + torch_meta["seq_len"], + torch_meta["input_pos"], + torch_meta["slot_idx"], + torch_meta["cu_seqlen"], + torch_meta["mla_cache"], + None, + kv_lora_rank, + ) + + # Run FlashInfer context + _GlobalFlashInferMLAPlanner.reset(torch.device(device)) + + qo_indptr = torch.zeros(batch_size + 1, dtype=torch.int32, device=device) + qo_indptr[1:] = torch.cumsum(torch.tensor(seq_lengths, device=device), dim=0).int() + + batch_indices, positions = flashinfer.get_batch_indices_positions( + qo_indptr, + flashinfer.get_seq_lens( + flashinfer_meta["cu_num_pages"], + flashinfer_meta["last_page_len"], + page_size=page_size, + ), + total_tokens, + ) + + flashinfer_output_context = torch.ops.auto_deploy.flashinfer_mla_with_cache( + q_nope_flat, + q_pe_flat, + compressed_kv_flat, + kpe_flat, + kv_b_proj_weight, + flashinfer_meta["batch_info_host"], + flashinfer_meta["cu_seqlen_host"], + flashinfer_meta["cu_num_pages"], + flashinfer_meta["cu_num_pages_host"], + flashinfer_meta["cache_loc"], + flashinfer_meta["last_page_len"], + flashinfer_meta["last_page_len_host"], + flashinfer_meta["seq_len_with_cache_host"], + batch_indices, + positions, + flashinfer_meta["ckv_cache"], + flashinfer_meta["kpe_cache"], + None, + kv_lora_rank, + ) + + # Verify context outputs match + assert torch.allclose( + flashinfer_output_context.cpu().to(torch.float32), + torch_output_context.cpu().to(torch.float32), + atol=0.05, + rtol=0.02, + ), ( + f"Context phase outputs don't match. " + f"Max diff: {(flashinfer_output_context - torch_output_context).abs().max():.6f}" + ) + + # ========================================================================= + # Multiple decode steps + # ========================================================================= + current_positions = list(seq_lengths) # Track current position for each sequence + + for decode_step in range(num_decode_steps): + # Create decode inputs for this step + inputs_decode = _create_mla_inputs( + batch_size, + 1, # seq_length = 1 for decode + num_heads, + qk_nope_head_dim, + qk_rope_head_dim, + kv_lora_rank, + v_head_dim, + dtype, + device, + ) + + seq_lengths_decode = [1] * batch_size + input_positions_decode = current_positions.copy() + + # Update torch metadata for decode + torch_meta_decode = _create_unpaged_cache_and_metadata( + batch_size, + max_seq_len, + kv_lora_rank, + qk_rope_head_dim, + dtype, + device, + seq_lengths_decode, + input_positions_decode, + ) + # Use the same cache (accumulated from context and previous decode steps) + torch_meta_decode["mla_cache"] = torch_meta["mla_cache"] + + # Update flashinfer metadata for decode + flashinfer_meta_decode = _create_paged_cache_and_metadata( + batch_size, + max_num_pages, + page_size, + kv_lora_rank, + qk_rope_head_dim, + dtype, + device, + seq_lengths_decode, + input_positions_decode, + ) + # Use the same caches + flashinfer_meta_decode["ckv_cache"] = flashinfer_meta["ckv_cache"] + flashinfer_meta_decode["kpe_cache"] = flashinfer_meta["kpe_cache"] + + # Run torch backend decode + torch_output_decode = torch.ops.auto_deploy.torch_cached_mla_with_cache( + inputs_decode["q_nope"], + inputs_decode["q_pe"], + inputs_decode["compressed_kv"], + inputs_decode["kpe"], + kv_b_proj_weight, + torch_meta_decode["batch_info_host"], + torch_meta_decode["seq_len"], + torch_meta_decode["input_pos"], + torch_meta_decode["slot_idx"], + torch_meta_decode["cu_seqlen"], + torch_meta_decode["mla_cache"], + None, + kv_lora_rank, + ) + + # Run FlashInfer decode + _GlobalFlashInferMLAPlanner.reset(torch.device(device)) + + qo_indptr_decode = torch.zeros(batch_size + 1, dtype=torch.int32, device=device) + qo_indptr_decode[1:] = torch.cumsum( + torch.tensor(seq_lengths_decode, device=device), dim=0 + ).int() + + batch_indices_decode, positions_decode = flashinfer.get_batch_indices_positions( + qo_indptr_decode, + flashinfer.get_seq_lens( + flashinfer_meta_decode["cu_num_pages"], + flashinfer_meta_decode["last_page_len"], + page_size=page_size, + ), + batch_size, + ) + + flashinfer_output_decode = torch.ops.auto_deploy.flashinfer_mla_with_cache( + inputs_decode["q_nope"], + inputs_decode["q_pe"], + inputs_decode["compressed_kv"], + inputs_decode["kpe"], + kv_b_proj_weight, + flashinfer_meta_decode["batch_info_host"], + flashinfer_meta_decode["cu_seqlen_host"], + flashinfer_meta_decode["cu_num_pages"], + flashinfer_meta_decode["cu_num_pages_host"], + flashinfer_meta_decode["cache_loc"], + flashinfer_meta_decode["last_page_len"], + flashinfer_meta_decode["last_page_len_host"], + flashinfer_meta_decode["seq_len_with_cache_host"], + batch_indices_decode, + positions_decode, + flashinfer_meta_decode["ckv_cache"], + flashinfer_meta_decode["kpe_cache"], + None, + kv_lora_rank, + ) + + # Verify decode outputs match + assert torch.allclose( + flashinfer_output_decode.cpu().to(torch.float32), + torch_output_decode.cpu().to(torch.float32), + atol=0.05, + rtol=0.02, + ), ( + f"Decode step {decode_step + 1} outputs don't match. " + f"Max diff: {(flashinfer_output_decode - torch_output_decode).abs().max():.6f}" + ) + # Update positions for next decode step + current_positions = [pos + 1 for pos in current_positions] + + # Final verification: all outputs should be finite + assert torch.isfinite(flashinfer_output_decode).all(), ( + "Final decode output contains NaN or Inf values" + ) + + +@pytest.mark.parametrize( + "chunk_config", + [ + # Each config has list of chunk sizes per sequence + # e.g., [[32, 16, 8], [64, 32, 16]] means 2 sequences with 3 chunks each + {"chunks_per_seq": [[32, 16], [64, 32]]}, # 2 sequences, 2 chunks each + {"chunks_per_seq": [[32, 16, 8], [64, 32, 16]]}, # 2 sequences, 3 chunks each + {"chunks_per_seq": [[64, 32, 16, 8]]}, # 1 sequence, 4 chunks + { + "chunks_per_seq": [[32, 32, 32], [48, 48, 48], [64, 64, 64]] + }, # 3 sequences, 3 chunks each + {"chunks_per_seq": [[16, 16, 16, 16, 16]]}, # 1 sequence, 5 chunks + ], +) +@pytest.mark.parametrize("num_heads", [8]) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@pytest.mark.parametrize("device", ["cuda"]) +def test_flashinfer_mla_chunked_prefill(chunk_config, num_heads, dtype, device): + """Test FlashInfer MLA chunked prefill (incremental prefill) with multiple chunks. + + This test verifies that chunked prefill works correctly when: + 1. First chunk is processed (input_pos == 0) - uses BatchPrefillWithRaggedKVCacheWrapper + 2. Subsequent chunks are processed (input_pos > 0) - uses BatchMLAPagedAttentionWrapper + + In chunked prefill, the Q tokens attend to all KV tokens (cached + current), + which is different from regular prefill where Q and KV lengths are equal. + + Compares flashinfer_mla_with_cache against torch_backend_mla_with_cache. + """ + chunks_per_seq = chunk_config["chunks_per_seq"] + batch_size = len(chunks_per_seq) + num_chunks = len(chunks_per_seq[0]) # Assume all sequences have same number of chunks + + # MLA dimensions + qk_nope_head_dim = 128 + qk_rope_head_dim = 64 # Must be 64 for FlashInfer MLA. + kv_lora_rank = 512 # Must be 512 for FlashInfer MLA. + v_head_dim = 128 + page_size = 32 + + # Calculate total sequence lengths + total_seq_lengths = [sum(chunks) for chunks in chunks_per_seq] + max_seq_len = max(total_seq_lengths) + 128 + max_num_pages = batch_size * (max_seq_len // page_size + 2) + + # Common kv_b_proj_weight + kv_head_dim = qk_nope_head_dim + v_head_dim + weight_scale = 1.0 / (kv_lora_rank**0.5) + kv_b_proj_weight = ( + torch.randn(num_heads * kv_head_dim, kv_lora_rank, dtype=dtype, device=device) + * weight_scale + ) + + # Initialize caches (will be reused across chunks) + torch_mla_cache = torch.zeros( + batch_size, max_seq_len, kv_lora_rank + qk_rope_head_dim, dtype=dtype, device=device + ) + flashinfer_ckv_cache = torch.zeros( + max_num_pages, page_size, kv_lora_rank, dtype=dtype, device=device + ) + flashinfer_kpe_cache = torch.zeros( + max_num_pages, page_size, qk_rope_head_dim, dtype=dtype, device=device + ) + + # Track cumulative positions per sequence + cumulative_positions = [0] * batch_size + + # Process each chunk + for chunk_idx in range(num_chunks): + # Get current chunk lengths + current_chunk_lengths = [ + chunks_per_seq[seq_idx][chunk_idx] for seq_idx in range(batch_size) + ] + total_tokens = sum(current_chunk_lengths) + input_positions = cumulative_positions.copy() + + # Create inputs for this chunk + q_nope_list = [] + q_pe_list = [] + compressed_kv_list = [] + kpe_list = [] + + for chunk_len in current_chunk_lengths: + inputs = _create_mla_inputs( + 1, + chunk_len, + num_heads, + qk_nope_head_dim, + qk_rope_head_dim, + kv_lora_rank, + v_head_dim, + dtype, + device, + ) + q_nope_list.append(inputs["q_nope"].squeeze(0)) + q_pe_list.append(inputs["q_pe"].squeeze(0)) + compressed_kv_list.append(inputs["compressed_kv"].squeeze(0)) + kpe_list.append(inputs["kpe"].squeeze(0)) + + q_nope_flat = torch.cat(q_nope_list, dim=0).unsqueeze(0) + q_pe_flat = torch.cat(q_pe_list, dim=0).unsqueeze(0) + compressed_kv_flat = torch.cat(compressed_kv_list, dim=0).unsqueeze(0) + kpe_flat = torch.cat(kpe_list, dim=0).unsqueeze(0) + + # ===================================================================== + # Torch backend + # ===================================================================== + torch_meta = _create_unpaged_cache_and_metadata( + batch_size, + max_seq_len, + kv_lora_rank, + qk_rope_head_dim, + dtype, + device, + current_chunk_lengths, + input_positions, + ) + torch_meta["mla_cache"] = torch_mla_cache # Use shared cache + + torch_output = torch.ops.auto_deploy.torch_cached_mla_with_cache( + q_nope_flat, + q_pe_flat, + compressed_kv_flat, + kpe_flat, + kv_b_proj_weight, + torch_meta["batch_info_host"], + torch_meta["seq_len"], + torch_meta["input_pos"], + torch_meta["slot_idx"], + torch_meta["cu_seqlen"], + torch_meta["mla_cache"], + None, + kv_lora_rank, + ) + + # ===================================================================== + # FlashInfer backend + # ===================================================================== + _GlobalFlashInferMLAPlanner.reset(torch.device(device)) + + flashinfer_meta = _create_paged_cache_and_metadata( + batch_size, + max_num_pages, + page_size, + kv_lora_rank, + qk_rope_head_dim, + dtype, + device, + current_chunk_lengths, + input_positions, + ) + # Use shared caches + flashinfer_meta["ckv_cache"] = flashinfer_ckv_cache + flashinfer_meta["kpe_cache"] = flashinfer_kpe_cache + + qo_indptr = torch.zeros(batch_size + 1, dtype=torch.int32, device=device) + qo_indptr[1:] = torch.cumsum( + torch.tensor(current_chunk_lengths, device=device), dim=0 + ).int() + + batch_indices, positions = flashinfer.get_batch_indices_positions( + qo_indptr, + flashinfer.get_seq_lens( + flashinfer_meta["cu_num_pages"], + flashinfer_meta["last_page_len"], + page_size=page_size, + ), + total_tokens, + ) + + flashinfer_output = torch.ops.auto_deploy.flashinfer_mla_with_cache( + q_nope_flat, + q_pe_flat, + compressed_kv_flat, + kpe_flat, + kv_b_proj_weight, + flashinfer_meta["batch_info_host"], + flashinfer_meta["cu_seqlen_host"], + flashinfer_meta["cu_num_pages"], + flashinfer_meta["cu_num_pages_host"], + flashinfer_meta["cache_loc"], + flashinfer_meta["last_page_len"], + flashinfer_meta["last_page_len_host"], + flashinfer_meta["seq_len_with_cache_host"], + batch_indices, + positions, + flashinfer_meta["ckv_cache"], + flashinfer_meta["kpe_cache"], + None, + kv_lora_rank, + ) + + # Verify outputs match + is_first_chunk = chunk_idx == 0 + chunk_type = "regular prefill" if is_first_chunk else "chunked prefill" + assert torch.allclose( + flashinfer_output.cpu().to(torch.float32), + torch_output.cpu().to(torch.float32), + atol=0.05, + rtol=0.02, + ), ( + f"Chunk {chunk_idx + 1}/{num_chunks} ({chunk_type}) outputs don't match. " + f"Max diff: {(flashinfer_output - torch_output).abs().max():.6f}" + ) + + # Verify outputs are finite + assert torch.isfinite(flashinfer_output).all(), ( + f"Chunk {chunk_idx + 1}/{num_chunks} ({chunk_type}) output contains NaN or Inf values" + ) + + # Update cumulative positions for next chunk + for seq_idx in range(batch_size): + cumulative_positions[seq_idx] += current_chunk_lengths[seq_idx] + + +# ============================================================================= +# CUDA Graph Tests +# ============================================================================= +# Tests for CUDA graph functionality of the FlashInfer MLA planner to verify +# that wrappers are correctly created and cached with use_cuda_graph=True. + + +@pytest.mark.parametrize("prefill_seq_length", [64, 128]) +@pytest.mark.parametrize("num_heads", [1, 8]) +@pytest.mark.parametrize("batch_size", [4, 16]) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@pytest.mark.parametrize("device", ["cuda"]) +def test_flashinfer_mla_cuda_graph_wrapper_creation( + prefill_seq_length, num_heads, batch_size, dtype, device +): + """Test that CUDA graph wrappers are created with use_cuda_graph=True during warm-up. + + This test verifies that: + 1. During CudaGraphWarmUpPhase, the planner creates a wrapper with use_cuda_graph=True + 2. The wrapper is cached in cached_cuda_graph_decode_wrappers + 3. The wrapper has correct buffer tensors attached + """ + # MLA dimensions + qk_nope_head_dim = 128 + qk_rope_head_dim = 64 # Must be 64 for FlashInfer MLA. + kv_lora_rank = 512 # Must be 512 for FlashInfer MLA. + v_head_dim = 128 + page_size = 64 + + seq_length = 1 # Decode phase + + max_seq_len = 256 + max_num_pages = batch_size * (max_seq_len // page_size + 1) + + # Create input tensors + inputs = _create_mla_inputs( + batch_size, + seq_length, + num_heads, + qk_nope_head_dim, + qk_rope_head_dim, + kv_lora_rank, + v_head_dim, + dtype, + device, + ) + + # Sequence lengths and positions + seq_lengths = [seq_length] * batch_size + input_positions = [prefill_seq_length] * batch_size + + # Create paged cache + flashinfer_meta = _create_paged_cache_and_metadata( + batch_size, + max_num_pages, + page_size, + kv_lora_rank, + qk_rope_head_dim, + dtype, + device, + seq_lengths, + input_positions, + ) + + # Pre-fill cache with random data + if prefill_seq_length > 0: + total_prefill_pages = (prefill_seq_length - 1) // page_size + 1 + for batch_idx in range(batch_size): + page_start = batch_idx * total_prefill_pages + for page_offset in range(total_prefill_pages): + page_idx = page_start + page_offset + tokens_in_page = min(page_size, prefill_seq_length - page_offset * page_size) + if tokens_in_page > 0: + flashinfer_meta["ckv_cache"][page_idx, :tokens_in_page] = torch.randn( + tokens_in_page, kv_lora_rank, dtype=dtype, device=device + ) / (kv_lora_rank**0.5) + flashinfer_meta["kpe_cache"][page_idx, :tokens_in_page] = torch.randn( + tokens_in_page, qk_rope_head_dim, dtype=dtype, device=device + ) / (qk_rope_head_dim**0.5) + + # Reset planner + _GlobalFlashInferMLAPlanner.workspace_buffer = None + _GlobalFlashInferMLAPlanner.cached_cuda_graph_decode_wrappers = {} + _GlobalFlashInferMLAPlanner.reset(torch.device(device)) + + qo_indptr = torch.zeros(batch_size + 1, dtype=torch.int32, device=device) + qo_indptr[1:] = torch.cumsum(torch.tensor(seq_lengths, device=device), dim=0).int() + + batch_indices, positions = flashinfer.get_batch_indices_positions( + qo_indptr, + flashinfer.get_seq_lens( + flashinfer_meta["cu_num_pages"], + flashinfer_meta["last_page_len"], + page_size=page_size, + ), + batch_size * seq_length, + ) + + # Verify no wrappers exist before warm-up + assert len(_GlobalFlashInferMLAPlanner.cached_cuda_graph_decode_wrappers) == 0, ( + "Expected no cached wrappers before warm-up" + ) + + # Warm-up phase: This triggers wrapper creation with use_cuda_graph=True + with CudaGraphWarmUpPhase(): + output = torch.ops.auto_deploy.flashinfer_mla_with_cache( + inputs["q_nope"], + inputs["q_pe"], + inputs["compressed_kv"], + inputs["kpe"], + inputs["kv_b_proj_weight"], + flashinfer_meta["batch_info_host"], + flashinfer_meta["cu_seqlen_host"], + flashinfer_meta["cu_num_pages"], + flashinfer_meta["cu_num_pages_host"], + flashinfer_meta["cache_loc"], + flashinfer_meta["last_page_len"], + flashinfer_meta["last_page_len_host"], + flashinfer_meta["seq_len_with_cache_host"], + batch_indices, + positions, + flashinfer_meta["ckv_cache"], + flashinfer_meta["kpe_cache"], + None, # scale + kv_lora_rank, + ) + + # Verify a CUDA graph wrapper was created + assert len(_GlobalFlashInferMLAPlanner.cached_cuda_graph_decode_wrappers) == 1, ( + f"Expected 1 cached wrapper after warm-up, " + f"got {len(_GlobalFlashInferMLAPlanner.cached_cuda_graph_decode_wrappers)}" + ) + + # Verify the wrapper has the correct plan params + for ( + plan_params, + wrapper, + ) in _GlobalFlashInferMLAPlanner.cached_cuda_graph_decode_wrappers.items(): + assert plan_params.num_seq == batch_size, ( + f"Plan params num_seq={plan_params.num_seq} doesn't match batch_size={batch_size}" + ) + assert plan_params.num_heads == num_heads, ( + f"Plan params num_heads={plan_params.num_heads} doesn't match num_heads={num_heads}" + ) + assert plan_params.kv_lora_rank == kv_lora_rank, ( + f"Plan params kv_lora_rank={plan_params.kv_lora_rank} doesn't match " + f"kv_lora_rank={kv_lora_rank}" + ) + assert plan_params.page_size == page_size, ( + f"Plan params page_size={plan_params.page_size} doesn't match page_size={page_size}" + ) + + # Verify wrapper is not None + assert wrapper is not None, "CUDA graph wrapper should not be None" + + # Verify output is valid + expected_shape = (batch_size, seq_length, num_heads, v_head_dim) + assert output.shape == expected_shape, f"Expected shape {expected_shape}, got {output.shape}" + assert torch.isfinite(output).all(), "Output contains NaN or Inf values" + + +@pytest.mark.parametrize("batch_size", [1, 4, 8, 16]) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@pytest.mark.parametrize("device", ["cuda"]) +def test_flashinfer_mla_cuda_graph_wrapper_caching_per_batch_size(batch_size, dtype, device): + """Test that CUDA graph wrappers are cached per batch size. + + This test verifies that: + 1. Each batch size gets its own cached wrapper + 2. Wrappers are keyed by MLADecodePlanParams which includes num_seq + """ + # MLA dimensions + num_heads = 8 + qk_nope_head_dim = 128 + qk_rope_head_dim = 64 + kv_lora_rank = 512 + v_head_dim = 128 + page_size = 64 + prefill_seq_length = 64 + + seq_length = 1 # Decode phase + + max_seq_len = 256 + max_num_pages = batch_size * (max_seq_len // page_size + 1) + + # Create inputs + inputs = _create_mla_inputs( + batch_size, + seq_length, + num_heads, + qk_nope_head_dim, + qk_rope_head_dim, + kv_lora_rank, + v_head_dim, + dtype, + device, + ) + + seq_lengths = [seq_length] * batch_size + input_positions = [prefill_seq_length] * batch_size + + flashinfer_meta = _create_paged_cache_and_metadata( + batch_size, + max_num_pages, + page_size, + kv_lora_rank, + qk_rope_head_dim, + dtype, + device, + seq_lengths, + input_positions, + ) + + # Pre-fill cache + if prefill_seq_length > 0: + total_prefill_pages = (prefill_seq_length - 1) // page_size + 1 + for batch_idx in range(batch_size): + page_start = batch_idx * total_prefill_pages + for page_offset in range(total_prefill_pages): + page_idx = page_start + page_offset + tokens_in_page = min(page_size, prefill_seq_length - page_offset * page_size) + if tokens_in_page > 0: + flashinfer_meta["ckv_cache"][page_idx, :tokens_in_page] = torch.randn( + tokens_in_page, kv_lora_rank, dtype=dtype, device=device + ) / (kv_lora_rank**0.5) + flashinfer_meta["kpe_cache"][page_idx, :tokens_in_page] = torch.randn( + tokens_in_page, qk_rope_head_dim, dtype=dtype, device=device + ) / (qk_rope_head_dim**0.5) + + qo_indptr = torch.zeros(batch_size + 1, dtype=torch.int32, device=device) + qo_indptr[1:] = torch.cumsum(torch.tensor(seq_lengths, device=device), dim=0).int() + + batch_indices, positions = flashinfer.get_batch_indices_positions( + qo_indptr, + flashinfer.get_seq_lens( + flashinfer_meta["cu_num_pages"], + flashinfer_meta["last_page_len"], + page_size=page_size, + ), + batch_size * seq_length, + ) + + # Reset planner + _GlobalFlashInferMLAPlanner.workspace_buffer = None + _GlobalFlashInferMLAPlanner.cached_cuda_graph_decode_wrappers = {} + _GlobalFlashInferMLAPlanner.reset(torch.device(device)) + + # Warm-up to create CUDA graph wrapper for this batch size + with CudaGraphWarmUpPhase(): + _ = torch.ops.auto_deploy.flashinfer_mla_with_cache( + inputs["q_nope"], + inputs["q_pe"], + inputs["compressed_kv"], + inputs["kpe"], + inputs["kv_b_proj_weight"], + flashinfer_meta["batch_info_host"], + flashinfer_meta["cu_seqlen_host"], + flashinfer_meta["cu_num_pages"], + flashinfer_meta["cu_num_pages_host"], + flashinfer_meta["cache_loc"], + flashinfer_meta["last_page_len"], + flashinfer_meta["last_page_len_host"], + flashinfer_meta["seq_len_with_cache_host"], + batch_indices, + positions, + flashinfer_meta["ckv_cache"], + flashinfer_meta["kpe_cache"], + None, + kv_lora_rank, + ) + + # Verify wrapper was created for this batch size + assert len(_GlobalFlashInferMLAPlanner.cached_cuda_graph_decode_wrappers) == 1, ( + f"Expected 1 cached wrapper for batch_size={batch_size}, " + f"got {len(_GlobalFlashInferMLAPlanner.cached_cuda_graph_decode_wrappers)}" + ) + + # Verify the wrapper has the correct num_seq + for plan_params in _GlobalFlashInferMLAPlanner.cached_cuda_graph_decode_wrappers: + assert plan_params.num_seq == batch_size, ( + f"Plan params num_seq={plan_params.num_seq} doesn't match batch_size={batch_size}" + ) + + +@pytest.mark.parametrize("prefill_seq_length", [64]) +@pytest.mark.parametrize("num_heads", [8]) +@pytest.mark.parametrize("batch_size", [4, 8]) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@pytest.mark.parametrize("device", ["cuda"]) +def test_flashinfer_mla_plan_generate_only( + prefill_seq_length, num_heads, batch_size, dtype, device +): + """Test plan_generate_only function for re-planning decode-only batches. + + This test verifies that: + 1. plan_generate_only can re-plan cached CUDA graph wrappers + 2. The wrappers can be used after re-planning + """ + # MLA dimensions + qk_nope_head_dim = 128 + qk_rope_head_dim = 64 + kv_lora_rank = 512 + v_head_dim = 128 + page_size = 64 + + seq_length = 1 # Decode phase + + max_seq_len = 256 + max_num_pages = batch_size * (max_seq_len // page_size + 1) + + # Create inputs + inputs = _create_mla_inputs( + batch_size, + seq_length, + num_heads, + qk_nope_head_dim, + qk_rope_head_dim, + kv_lora_rank, + v_head_dim, + dtype, + device, + ) + + seq_lengths = [seq_length] * batch_size + input_positions = [prefill_seq_length] * batch_size + + flashinfer_meta = _create_paged_cache_and_metadata( + batch_size, + max_num_pages, + page_size, + kv_lora_rank, + qk_rope_head_dim, + dtype, + device, + seq_lengths, + input_positions, + ) + + # Pre-fill cache + if prefill_seq_length > 0: + total_prefill_pages = (prefill_seq_length - 1) // page_size + 1 + for batch_idx in range(batch_size): + page_start = batch_idx * total_prefill_pages + for page_offset in range(total_prefill_pages): + page_idx = page_start + page_offset + tokens_in_page = min(page_size, prefill_seq_length - page_offset * page_size) + if tokens_in_page > 0: + flashinfer_meta["ckv_cache"][page_idx, :tokens_in_page] = torch.randn( + tokens_in_page, kv_lora_rank, dtype=dtype, device=device + ) / (kv_lora_rank**0.5) + flashinfer_meta["kpe_cache"][page_idx, :tokens_in_page] = torch.randn( + tokens_in_page, qk_rope_head_dim, dtype=dtype, device=device + ) / (qk_rope_head_dim**0.5) + + qo_indptr = torch.zeros(batch_size + 1, dtype=torch.int32, device=device) + qo_indptr[1:] = torch.cumsum(torch.tensor(seq_lengths, device=device), dim=0).int() + + batch_indices, positions = flashinfer.get_batch_indices_positions( + qo_indptr, + flashinfer.get_seq_lens( + flashinfer_meta["cu_num_pages"], + flashinfer_meta["last_page_len"], + page_size=page_size, + ), + batch_size * seq_length, + ) + + # Reset planner + _GlobalFlashInferMLAPlanner.workspace_buffer = None + _GlobalFlashInferMLAPlanner.cached_cuda_graph_decode_wrappers = {} + _GlobalFlashInferMLAPlanner.reset(torch.device(device)) + + # First warm-up to create the wrapper + with CudaGraphWarmUpPhase(): + output1 = torch.ops.auto_deploy.flashinfer_mla_with_cache( + inputs["q_nope"], + inputs["q_pe"], + inputs["compressed_kv"], + inputs["kpe"], + inputs["kv_b_proj_weight"], + flashinfer_meta["batch_info_host"], + flashinfer_meta["cu_seqlen_host"], + flashinfer_meta["cu_num_pages"], + flashinfer_meta["cu_num_pages_host"], + flashinfer_meta["cache_loc"], + flashinfer_meta["last_page_len"], + flashinfer_meta["last_page_len_host"], + flashinfer_meta["seq_len_with_cache_host"], + batch_indices, + positions, + flashinfer_meta["ckv_cache"], + flashinfer_meta["kpe_cache"], + None, + kv_lora_rank, + ) + + # Verify wrapper was created + assert len(_GlobalFlashInferMLAPlanner.cached_cuda_graph_decode_wrappers) == 1 + + # Now test plan_generate_only - this is called by the host-side preparation + # to re-plan the cached wrappers before graph replay + _GlobalFlashInferMLAPlanner.plan_generate_only( + batch_size, + flashinfer_meta["cu_num_pages"][: batch_size + 1], + flashinfer_meta["cache_loc"], + flashinfer_meta["last_page_len"][:batch_size], + ) + + # Run again (not in warm-up, so it should use cached wrapper) + # First, update the inputs to simulate new tokens + new_inputs = _create_mla_inputs( + batch_size, + seq_length, + num_heads, + qk_nope_head_dim, + qk_rope_head_dim, + kv_lora_rank, + v_head_dim, + dtype, + device, + ) + + # Create new metadata for position+1 + new_input_positions = [prefill_seq_length + 1] * batch_size + new_flashinfer_meta = _create_paged_cache_and_metadata( + batch_size, + max_num_pages, + page_size, + kv_lora_rank, + qk_rope_head_dim, + dtype, + device, + seq_lengths, + new_input_positions, + ) + # Reuse the same cache + new_flashinfer_meta["ckv_cache"] = flashinfer_meta["ckv_cache"] + new_flashinfer_meta["kpe_cache"] = flashinfer_meta["kpe_cache"] + + new_qo_indptr = torch.zeros(batch_size + 1, dtype=torch.int32, device=device) + new_qo_indptr[1:] = torch.cumsum(torch.tensor(seq_lengths, device=device), dim=0).int() + + new_batch_indices, new_positions = flashinfer.get_batch_indices_positions( + new_qo_indptr, + flashinfer.get_seq_lens( + new_flashinfer_meta["cu_num_pages"], + new_flashinfer_meta["last_page_len"], + page_size=page_size, + ), + batch_size * seq_length, + ) + + output2 = torch.ops.auto_deploy.flashinfer_mla_with_cache( + new_inputs["q_nope"], + new_inputs["q_pe"], + new_inputs["compressed_kv"], + new_inputs["kpe"], + inputs["kv_b_proj_weight"], # Use same weights + new_flashinfer_meta["batch_info_host"], + new_flashinfer_meta["cu_seqlen_host"], + new_flashinfer_meta["cu_num_pages"], + new_flashinfer_meta["cu_num_pages_host"], + new_flashinfer_meta["cache_loc"], + new_flashinfer_meta["last_page_len"], + new_flashinfer_meta["last_page_len_host"], + new_flashinfer_meta["seq_len_with_cache_host"], + new_batch_indices, + new_positions, + new_flashinfer_meta["ckv_cache"], + new_flashinfer_meta["kpe_cache"], + None, + kv_lora_rank, + ) + + # Verify output is valid + expected_shape = (batch_size, seq_length, num_heads, v_head_dim) + assert output2.shape == expected_shape, f"Expected shape {expected_shape}, got {output2.shape}" + assert torch.isfinite(output2).all(), "Output contains NaN or Inf values" + + # Outputs should be different since inputs are different + assert not torch.allclose(output1, output2, atol=1e-6), ( + "Outputs should differ since inputs are different" + ) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@pytest.mark.parametrize("device", ["cuda"]) +def test_flashinfer_mla_cuda_graph_multiple_batch_sizes(dtype, device): + """Test that multiple batch sizes can have their own CUDA graph wrappers. + + This test verifies that the planner correctly caches wrappers for + different batch sizes, which is important for supporting multiple + CUDA graph configurations. + """ + # MLA dimensions + num_heads = 8 + qk_nope_head_dim = 128 + qk_rope_head_dim = 64 + kv_lora_rank = 512 + v_head_dim = 128 + page_size = 64 + prefill_seq_length = 64 + + seq_length = 1 # Decode phase + + batch_sizes = [4, 8, 16] + + # Reset planner + _GlobalFlashInferMLAPlanner.workspace_buffer = None + _GlobalFlashInferMLAPlanner.cached_cuda_graph_decode_wrappers = {} + _GlobalFlashInferMLAPlanner.reset(torch.device(device)) + + for batch_size in batch_sizes: + max_num_pages = batch_size * (256 // page_size + 1) + + inputs = _create_mla_inputs( + batch_size, + seq_length, + num_heads, + qk_nope_head_dim, + qk_rope_head_dim, + kv_lora_rank, + v_head_dim, + dtype, + device, + ) + + seq_lengths = [seq_length] * batch_size + input_positions = [prefill_seq_length] * batch_size + + flashinfer_meta = _create_paged_cache_and_metadata( + batch_size, + max_num_pages, + page_size, + kv_lora_rank, + qk_rope_head_dim, + dtype, + device, + seq_lengths, + input_positions, + ) + + qo_indptr = torch.zeros(batch_size + 1, dtype=torch.int32, device=device) + qo_indptr[1:] = torch.cumsum(torch.tensor(seq_lengths, device=device), dim=0).int() + + batch_indices, positions = flashinfer.get_batch_indices_positions( + qo_indptr, + flashinfer.get_seq_lens( + flashinfer_meta["cu_num_pages"], + flashinfer_meta["last_page_len"], + page_size=page_size, + ), + batch_size * seq_length, + ) + + # Warm-up to create wrapper for this batch size + with CudaGraphWarmUpPhase(): + _ = torch.ops.auto_deploy.flashinfer_mla_with_cache( + inputs["q_nope"], + inputs["q_pe"], + inputs["compressed_kv"], + inputs["kpe"], + inputs["kv_b_proj_weight"], + flashinfer_meta["batch_info_host"], + flashinfer_meta["cu_seqlen_host"], + flashinfer_meta["cu_num_pages"], + flashinfer_meta["cu_num_pages_host"], + flashinfer_meta["cache_loc"], + flashinfer_meta["last_page_len"], + flashinfer_meta["last_page_len_host"], + flashinfer_meta["seq_len_with_cache_host"], + batch_indices, + positions, + flashinfer_meta["ckv_cache"], + flashinfer_meta["kpe_cache"], + None, + kv_lora_rank, + ) + + # Verify we have a wrapper for each batch size + assert len(_GlobalFlashInferMLAPlanner.cached_cuda_graph_decode_wrappers) == len(batch_sizes), ( + f"Expected {len(batch_sizes)} cached wrappers, " + f"got {len(_GlobalFlashInferMLAPlanner.cached_cuda_graph_decode_wrappers)}" + ) + + # Verify each batch size has a wrapper + cached_num_seqs = { + params.num_seq for params in _GlobalFlashInferMLAPlanner.cached_cuda_graph_decode_wrappers + } + assert cached_num_seqs == set(batch_sizes), ( + f"Expected wrappers for batch_sizes {batch_sizes}, got {cached_num_seqs}" + ) + + +@pytest.mark.parametrize("batch_size", [4]) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@pytest.mark.parametrize("device", ["cuda"]) +def test_flashinfer_mla_init_decode_wrapper_with_buffers(batch_size, dtype, device): + """Test that _init_decode_wrapper correctly passes buffer tensors with use_cuda_graph=True. + + This test directly tests the _init_decode_wrapper method to verify buffer handling. + """ + # Reset planner + _GlobalFlashInferMLAPlanner.workspace_buffer = None + _GlobalFlashInferMLAPlanner.cached_cuda_graph_decode_wrappers = {} + _GlobalFlashInferMLAPlanner.reset(torch.device(device)) + + # Create buffer tensors + qo_indptr = torch.arange(batch_size + 1, device=device, dtype=torch.int32) + kv_indptr = torch.arange(batch_size + 1, device=device, dtype=torch.int32) * 2 + kv_indices = torch.arange(batch_size * 2, device=device, dtype=torch.int32) + kv_len_arr = torch.ones(batch_size, device=device, dtype=torch.int32) * 64 + + # Test creating wrapper without CUDA graph (no buffers needed) + wrapper_no_cg = _GlobalFlashInferMLAPlanner._init_decode_wrapper(use_cuda_graph=False) + assert wrapper_no_cg is not None, "Should create wrapper without CUDA graph" + + # Test creating wrapper with CUDA graph (buffers required) + wrapper_with_cg = _GlobalFlashInferMLAPlanner._init_decode_wrapper( + use_cuda_graph=True, + qo_indptr=qo_indptr, + kv_indptr=kv_indptr, + kv_indices=kv_indices, + kv_len_arr=kv_len_arr, + ) + assert wrapper_with_cg is not None, "Should create wrapper with CUDA graph" + + # Both wrappers should be valid BatchMLAPagedAttentionWrapper instances + assert isinstance(wrapper_no_cg, flashinfer.mla.BatchMLAPagedAttentionWrapper), ( + "Should be BatchMLAPagedAttentionWrapper" + ) + assert isinstance(wrapper_with_cg, flashinfer.mla.BatchMLAPagedAttentionWrapper), ( + "Should be BatchMLAPagedAttentionWrapper" + ) diff --git a/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_torch_mla_op.py b/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_torch_mla_op.py new file mode 100644 index 000000000000..369f2b181f2a --- /dev/null +++ b/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_torch_mla_op.py @@ -0,0 +1,780 @@ +"""Comprehensive test suite for torch MLA backend operations. + +Tests the torch_mla source op and torch_backend_mla_with_cache cached op +with FlashInfer-compatible compressed cache layout. + +Key features: +- 5 tensor arguments: q_nope, q_pe, compressed_kv, kpe, kv_b_proj_weight +- Compressed cache: [max_batch, max_seq, kv_lora_rank + qk_rope_head_dim] +- Prefill: Expand compressed_kv, compute normal attention +- Generate: Weight absorption for efficiency +""" + +import math + +import numpy as np +import pytest +import torch + +import tensorrt_llm._torch.auto_deploy # noqa: F401 + + +def numpy_mla_reference_with_expansion( + q_nope: np.ndarray, + q_pe: np.ndarray, + compressed_kv: np.ndarray, + kpe: np.ndarray, + kv_b_proj_weight: np.ndarray, + mla_cache: np.ndarray, + seq_len: np.ndarray, + input_pos: np.ndarray, + cache_loc: np.ndarray, + seq_start: np.ndarray, + scale: float = None, + kv_lora_rank: int = None, + is_generate: bool = False, +): + """Numpy reference implementation of MLA attention with FlashInfer cache layout. + + This expands compressed_kv using kv_b_proj_weight for attention computation. + """ + # Get dimensions + if is_generate: + batch_size = q_nope.shape[0] + num_heads = q_nope.shape[2] + qk_nope_head_dim = q_nope.shape[3] + qk_rope_head_dim = q_pe.shape[3] + else: + batch_size = len(seq_len) + num_heads = q_nope.shape[2] + qk_nope_head_dim = q_nope.shape[3] + qk_rope_head_dim = q_pe.shape[3] + + qk_head_dim = qk_nope_head_dim + qk_rope_head_dim + + if kv_lora_rank is None: + kv_lora_rank = compressed_kv.shape[-1] + + # Infer v_head_dim from kv_b_proj_weight + out_features = kv_b_proj_weight.shape[0] + kv_head_dim = out_features // num_heads + v_head_dim = kv_head_dim - qk_nope_head_dim + + if scale is None: + scale = 1.0 / math.sqrt(qk_head_dim) + + # Update MLA cache first + if is_generate: + for i in range(batch_size): + cache_idx = cache_loc[i] + pos = input_pos[i] + mla_cache[cache_idx, pos, :kv_lora_rank] = compressed_kv[i, 0] + mla_cache[cache_idx, pos, kv_lora_rank:] = kpe[i, 0, 0] + else: + for i in range(batch_size): + cache_idx = cache_loc[i] + pos = input_pos[i] + seq_len_i = seq_len[i] + seq_start_i = seq_start[i] + for j in range(seq_len_i): + mla_cache[cache_idx, pos + j, :kv_lora_rank] = compressed_kv[seq_start_i + j] + mla_cache[cache_idx, pos + j, kv_lora_rank:] = kpe[seq_start_i + j, 0] + + # Compute attention for each sequence + outputs = [] + + for i in range(batch_size): + cache_idx = cache_loc[i] + pos = input_pos[i] + seq_len_i = seq_len[i] + seq_start_i = seq_start[i] + + if seq_len_i == 0: + continue + + # Get query for this sequence + if is_generate: + q_nope_seq = q_nope[i, 0] # [N, qk_nope_head_dim] + q_pe_seq = q_pe[i, 0] # [N, qk_rope_head_dim] + else: + q_nope_seq = q_nope[seq_start_i : seq_start_i + seq_len_i] + q_pe_seq = q_pe[seq_start_i : seq_start_i + seq_len_i] + + # Get cached compressed_kv and kpe + kv_seq_len = pos + seq_len_i + cached_data = mla_cache[cache_idx, :kv_seq_len] + compressed_kv_cached = cached_data[:, :kv_lora_rank] + kpe_cached = cached_data[:, kv_lora_rank:] + + # Expand compressed_kv using kv_b_proj_weight + # compressed_kv_cached: [kv_seq_len, kv_lora_rank] + # kv_b_proj_weight: [N * kv_head_dim, kv_lora_rank] + kv_expanded = np.matmul(compressed_kv_cached, kv_b_proj_weight.T) + kv_expanded = kv_expanded.reshape(kv_seq_len, num_heads, kv_head_dim) + + k_nope = kv_expanded[:, :, :qk_nope_head_dim] + v = kv_expanded[:, :, qk_nope_head_dim:] + + # Expand kpe to all heads + kpe_expanded = np.broadcast_to( + kpe_cached[:, None, :], (kv_seq_len, num_heads, qk_rope_head_dim) + ) + + # Construct full query and key + if is_generate: + query_full = np.concatenate([q_nope_seq, q_pe_seq], axis=-1) + else: + query_full = np.concatenate([q_nope_seq, q_pe_seq], axis=-1) + + key_full = np.concatenate([k_nope, kpe_expanded], axis=-1) + + # Compute attention scores + if is_generate: + attn_scores = np.einsum("nh,knh->nk", query_full, key_full) * scale + else: + attn_scores = np.einsum("snh,knh->snk", query_full, key_full) * scale + causal_mask = np.triu(np.ones((seq_len_i, kv_seq_len)), k=kv_seq_len - seq_len_i + 1) + attn_scores = np.where(causal_mask[:, None, :], -np.inf, attn_scores) + + # Apply softmax + attn_scores_max = np.max(attn_scores, axis=-1, keepdims=True) + attn_scores_exp = np.exp(attn_scores - attn_scores_max) + attn_weights = attn_scores_exp / np.sum(attn_scores_exp, axis=-1, keepdims=True) + + # Compute output + if is_generate: + attn_out = np.einsum("nk,knh->nh", attn_weights, v) + else: + attn_out = np.einsum("snk,knh->snh", attn_weights, v) + + outputs.append(attn_out) + + # Concatenate outputs + if len(outputs) == 0: + return np.zeros((1, 0, num_heads, v_head_dim), dtype=np.float32) + elif is_generate: + result = np.stack(outputs, axis=0) + return result[:, None, :, :] + else: + result = np.concatenate(outputs, axis=0) + return result[None, :, :, :] + + +class TestTorchMLASourceOp: + """Test torch_mla source op (without cache).""" + + @pytest.fixture(autouse=True) + def setup_method(self): + """Setup test configuration.""" + self.device = "cuda" + self.dtype = torch.bfloat16 + self.atol = 1e-2 + self.rtol = 1e-2 + + torch.cuda.empty_cache() + torch.manual_seed(42) + np.random.seed(42) + + def _create_mla_data( + self, + batch_size: int, + seq_len: int, + num_heads: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + kv_lora_rank: int, + v_head_dim: int, + layout: str = "bsnd", + ): + """Create test data for MLA source op with compressed_kv.""" + kv_head_dim = qk_nope_head_dim + v_head_dim + + if layout == "bsnd": + q_nope = torch.randn( + batch_size, + seq_len, + num_heads, + qk_nope_head_dim, + dtype=self.dtype, + device=self.device, + ) + q_pe = torch.randn( + batch_size, + seq_len, + num_heads, + qk_rope_head_dim, + dtype=self.dtype, + device=self.device, + ) + compressed_kv = torch.randn( + batch_size, seq_len, kv_lora_rank, dtype=self.dtype, device=self.device + ) + kpe = torch.randn( + batch_size, seq_len, 1, qk_rope_head_dim, dtype=self.dtype, device=self.device + ) + else: # bnsd + q_nope = torch.randn( + batch_size, + num_heads, + seq_len, + qk_nope_head_dim, + dtype=self.dtype, + device=self.device, + ) + q_pe = torch.randn( + batch_size, + num_heads, + seq_len, + qk_rope_head_dim, + dtype=self.dtype, + device=self.device, + ) + compressed_kv = torch.randn( + batch_size, seq_len, kv_lora_rank, dtype=self.dtype, device=self.device + ) + kpe = torch.randn( + batch_size, 1, seq_len, qk_rope_head_dim, dtype=self.dtype, device=self.device + ) + + # kv_b_proj_weight: [num_heads * kv_head_dim, kv_lora_rank] + kv_b_proj_weight = torch.randn( + num_heads * kv_head_dim, kv_lora_rank, dtype=self.dtype, device=self.device + ) + + return { + "q_nope": q_nope, + "q_pe": q_pe, + "compressed_kv": compressed_kv, + "kpe": kpe, + "kv_b_proj_weight": kv_b_proj_weight, + } + + def test_basic_functionality(self): + """Test basic MLA source op functionality.""" + batch_size, seq_len, num_heads = 2, 4, 8 + qk_nope_head_dim, qk_rope_head_dim = 128, 64 + kv_lora_rank = 512 + v_head_dim = 128 + + data = self._create_mla_data( + batch_size, + seq_len, + num_heads, + qk_nope_head_dim, + qk_rope_head_dim, + kv_lora_rank, + v_head_dim, + ) + + output = torch.ops.auto_deploy.torch_mla( + data["q_nope"], + data["q_pe"], + data["compressed_kv"], + data["kpe"], + data["kv_b_proj_weight"], + True, # is_causal + None, # scale + "bsnd", # layout + ) + + # Verify output shape: [B, S, N, v_head_dim] + expected_shape = (batch_size, seq_len, num_heads, v_head_dim) + assert output.shape == expected_shape, f"Expected {expected_shape}, got {output.shape}" + + # Verify output is finite + assert torch.isfinite(output).all(), "Output contains NaN or Inf values" + + def test_both_layouts(self): + """Test MLA source op with both bsnd and bnsd layouts.""" + batch_size, seq_len, num_heads = 2, 4, 8 + qk_nope_head_dim, qk_rope_head_dim = 64, 32 + kv_lora_rank = 256 + v_head_dim = 64 + + for layout in ["bsnd", "bnsd"]: + data = self._create_mla_data( + batch_size, + seq_len, + num_heads, + qk_nope_head_dim, + qk_rope_head_dim, + kv_lora_rank, + v_head_dim, + layout, + ) + + output = torch.ops.auto_deploy.torch_mla( + data["q_nope"], + data["q_pe"], + data["compressed_kv"], + data["kpe"], + data["kv_b_proj_weight"], + True, + None, + layout, + ) + + if layout == "bsnd": + expected_shape = (batch_size, seq_len, num_heads, v_head_dim) + else: + expected_shape = (batch_size, num_heads, seq_len, v_head_dim) + + assert output.shape == expected_shape, ( + f"Layout {layout}: Expected {expected_shape}, got {output.shape}" + ) + + def test_custom_scale(self): + """Test MLA source op with custom scale.""" + batch_size, seq_len, num_heads = 1, 2, 4 + qk_nope_head_dim, qk_rope_head_dim = 32, 16 + kv_lora_rank = 128 + v_head_dim = 32 + + data = self._create_mla_data( + batch_size, + seq_len, + num_heads, + qk_nope_head_dim, + qk_rope_head_dim, + kv_lora_rank, + v_head_dim, + ) + + # Test with default scale + output_default = torch.ops.auto_deploy.torch_mla( + data["q_nope"], + data["q_pe"], + data["compressed_kv"], + data["kpe"], + data["kv_b_proj_weight"], + True, + None, + "bsnd", + ) + + # Test with custom scale + custom_scale = 0.5 + output_custom = torch.ops.auto_deploy.torch_mla( + data["q_nope"], + data["q_pe"], + data["compressed_kv"], + data["kpe"], + data["kv_b_proj_weight"], + True, + custom_scale, + "bsnd", + ) + + # Outputs should be different + assert not torch.allclose(output_default, output_custom, atol=1e-3), ( + "Custom scale should affect output" + ) + + +class TestTorchBackendMLAWithCache: + """Test torch_backend_mla_with_cache cached op.""" + + @pytest.fixture(autouse=True) + def setup_method(self): + """Setup test configuration.""" + self.device = "cuda" + self.dtype = torch.bfloat16 + self.atol = 5e-2 + self.rtol = 5e-2 + + torch.cuda.empty_cache() + torch.manual_seed(42) + np.random.seed(42) + + def _create_cached_mla_data( + self, + batch_size: int, + seq_len: int, + num_heads: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + kv_lora_rank: int, + v_head_dim: int, + max_seq_len: int, + cache_offset: int = 0, + ): + """Create test data for cached MLA op with FlashInfer layout.""" + kv_head_dim = qk_nope_head_dim + v_head_dim + + # Create input tensors (BSND layout) + q_nope = torch.randn( + batch_size, seq_len, num_heads, qk_nope_head_dim, dtype=self.dtype, device=self.device + ) + q_pe = torch.randn( + batch_size, seq_len, num_heads, qk_rope_head_dim, dtype=self.dtype, device=self.device + ) + compressed_kv = torch.randn( + batch_size, seq_len, kv_lora_rank, dtype=self.dtype, device=self.device + ) + kpe = torch.randn( + batch_size, seq_len, 1, qk_rope_head_dim, dtype=self.dtype, device=self.device + ) + + # kv_b_proj_weight: [num_heads * kv_head_dim, kv_lora_rank] + kv_b_proj_weight = torch.randn( + num_heads * kv_head_dim, kv_lora_rank, dtype=self.dtype, device=self.device + ) + + # Create FlashInfer MLA cache: [max_batch, max_seq, kv_lora_rank + qk_rope_head_dim] + mla_cache = torch.zeros( + batch_size, + max_seq_len, + kv_lora_rank + qk_rope_head_dim, + dtype=self.dtype, + device=self.device, + ) + + # Pre-fill cache with random data if cache_offset > 0 + if cache_offset > 0: + mla_cache[:, :cache_offset, :] = torch.randn( + batch_size, + cache_offset, + kv_lora_rank + qk_rope_head_dim, + dtype=self.dtype, + device=self.device, + ) + + # Setup metadata + seq_len_tensor = torch.full((batch_size,), seq_len, device=self.device, dtype=torch.int32) + input_pos = torch.full((batch_size,), cache_offset, device=self.device, dtype=torch.int32) + cache_loc = torch.arange(batch_size, device=self.device, dtype=torch.int32) + + if seq_len == 1: + # Generate phase + batch_info_host = torch.tensor( + [0, 0, batch_size], device=self.device, dtype=torch.int32 + ) + cu_seqlen = torch.arange(batch_size, device=self.device, dtype=torch.int32) + else: + # Context phase + batch_info_host = torch.tensor( + [batch_size, batch_size * seq_len, 0], device=self.device, dtype=torch.int32 + ) + cu_seqlen = torch.arange( + 0, batch_size * seq_len, seq_len, device=self.device, dtype=torch.int32 + ) + # Flatten inputs for context phase + q_nope = q_nope.view(1, batch_size * seq_len, num_heads, qk_nope_head_dim) + q_pe = q_pe.view(1, batch_size * seq_len, num_heads, qk_rope_head_dim) + compressed_kv = compressed_kv.view(1, batch_size * seq_len, kv_lora_rank) + kpe = kpe.view(1, batch_size * seq_len, 1, qk_rope_head_dim) + + return { + "q_nope": q_nope, + "q_pe": q_pe, + "compressed_kv": compressed_kv, + "kpe": kpe, + "kv_b_proj_weight": kv_b_proj_weight, + "batch_info_host": batch_info_host, + "seq_len": seq_len_tensor, + "input_pos": input_pos, + "cache_loc": cache_loc, + "cu_seqlen": cu_seqlen, + "mla_cache": mla_cache, + "kv_lora_rank": kv_lora_rank, + } + + def _run_cached_mla(self, data, scale=None): + """Run cached MLA operation.""" + return torch.ops.auto_deploy.torch_cached_mla_with_cache( + data["q_nope"], + data["q_pe"], + data["compressed_kv"], + data["kpe"], + data["kv_b_proj_weight"], + data["batch_info_host"], + data["seq_len"], + data["input_pos"], + data["cache_loc"], + data["cu_seqlen"], + data["mla_cache"], + scale, + data["kv_lora_rank"], + ) + + def test_generate_phase_basic(self): + """Test generate phase (single token) basic functionality.""" + batch_size, seq_len, num_heads = 2, 1, 8 + qk_nope_head_dim, qk_rope_head_dim = 64, 32 + kv_lora_rank = 256 + v_head_dim = 64 + max_seq_len = 128 + cache_offset = 5 + + data = self._create_cached_mla_data( + batch_size, + seq_len, + num_heads, + qk_nope_head_dim, + qk_rope_head_dim, + kv_lora_rank, + v_head_dim, + max_seq_len, + cache_offset, + ) + + output = self._run_cached_mla(data) + + # Verify output shape + expected_shape = (batch_size, seq_len, num_heads, v_head_dim) + assert output.shape == expected_shape, f"Expected {expected_shape}, got {output.shape}" + + # Verify output is finite + assert torch.isfinite(output).all(), "Output contains NaN or Inf values" + + def test_context_phase_basic(self): + """Test context phase (multi-token) basic functionality.""" + batch_size, seq_len, num_heads = 2, 4, 8 + qk_nope_head_dim, qk_rope_head_dim = 64, 32 + kv_lora_rank = 256 + v_head_dim = 64 + max_seq_len = 128 + + data = self._create_cached_mla_data( + batch_size, + seq_len, + num_heads, + qk_nope_head_dim, + qk_rope_head_dim, + kv_lora_rank, + v_head_dim, + max_seq_len, + ) + + output = self._run_cached_mla(data) + + # Verify output shape + expected_shape = (1, batch_size * seq_len, num_heads, v_head_dim) + assert output.shape == expected_shape, f"Expected {expected_shape}, got {output.shape}" + + # Verify output is finite + assert torch.isfinite(output).all(), "Output contains NaN or Inf values" + + def test_cache_update_correctness(self): + """Test that cache is updated correctly during forward pass.""" + batch_size, seq_len, num_heads = 1, 1, 4 + qk_nope_head_dim, qk_rope_head_dim = 32, 16 + kv_lora_rank = 128 + v_head_dim = 32 + max_seq_len = 32 + cache_offset = 5 + + data = self._create_cached_mla_data( + batch_size, + seq_len, + num_heads, + qk_nope_head_dim, + qk_rope_head_dim, + kv_lora_rank, + v_head_dim, + max_seq_len, + cache_offset, + ) + + # Store original cache values at target position + original_cache_at_pos = data["mla_cache"][0, cache_offset].clone() + + # Run forward pass + _ = self._run_cached_mla(data) + + # Check cache was updated at the correct position + updated_cache_at_pos = data["mla_cache"][0, cache_offset] + + # The cache should have been updated + assert not torch.allclose(original_cache_at_pos, updated_cache_at_pos, atol=1e-6), ( + "Cache should have been updated at the target position" + ) + + def test_cache_layout_flashinfer_compatible(self): + """Test that cache layout matches FlashInfer spec (no num_heads dimension).""" + batch_size, seq_len, num_heads = 2, 1, 4 + qk_nope_head_dim, qk_rope_head_dim = 64, 32 + kv_lora_rank = 512 # DeepSeek-style + v_head_dim = 128 + max_seq_len = 64 + + data = self._create_cached_mla_data( + batch_size, + seq_len, + num_heads, + qk_nope_head_dim, + qk_rope_head_dim, + kv_lora_rank, + v_head_dim, + max_seq_len, + ) + + # Verify cache shape matches FlashInfer layout: [batch, seq, kv_lora_rank + rope_dim] + expected_cache_shape = (batch_size, max_seq_len, kv_lora_rank + qk_rope_head_dim) + assert data["mla_cache"].shape == expected_cache_shape, ( + f"Cache shape {data['mla_cache'].shape} doesn't match FlashInfer layout {expected_cache_shape}" + ) + + # Verify zero-copy slicing works + compressed_kv_slice = data["mla_cache"][:, :, :kv_lora_rank] + kpe_slice = data["mla_cache"][:, :, kv_lora_rank:] + + assert compressed_kv_slice.shape == (batch_size, max_seq_len, kv_lora_rank) + assert kpe_slice.shape == (batch_size, max_seq_len, qk_rope_head_dim) + + # Verify slices share memory (zero-copy) + assert compressed_kv_slice.data_ptr() == data["mla_cache"].data_ptr(), ( + "compressed_kv slice should be zero-copy" + ) + + def test_generate_with_reference(self): + """Test generate phase against numpy reference.""" + batch_size, seq_len, num_heads = 2, 1, 4 + qk_nope_head_dim, qk_rope_head_dim = 32, 16 + kv_lora_rank = 64 + v_head_dim = 32 + max_seq_len = 64 + cache_offset = 3 + + data = self._create_cached_mla_data( + batch_size, + seq_len, + num_heads, + qk_nope_head_dim, + qk_rope_head_dim, + kv_lora_rank, + v_head_dim, + max_seq_len, + cache_offset, + ) + + # Run backend + output = self._run_cached_mla(data) + + # Run numpy reference + reference = numpy_mla_reference_with_expansion( + data["q_nope"].cpu().float().numpy(), + data["q_pe"].cpu().float().numpy(), + data["compressed_kv"].cpu().float().numpy(), + data["kpe"].cpu().float().numpy(), + data["kv_b_proj_weight"].cpu().float().numpy(), + data["mla_cache"].cpu().float().numpy(), + data["seq_len"].cpu().numpy(), + data["input_pos"].cpu().numpy(), + data["cache_loc"].cpu().numpy(), + data["cu_seqlen"].cpu().numpy(), + None, + kv_lora_rank, + is_generate=True, + ) + + reference_torch = torch.from_numpy(reference).to(output.device, output.dtype) + assert torch.allclose(output, reference_torch, atol=self.atol, rtol=self.rtol), ( + f"Generate phase output doesn't match reference. " + f"Max diff: {(output - reference_torch).abs().max():.6f}" + ) + + def test_dtype_preservation(self): + """Test that output dtype matches input dtype.""" + batch_size, seq_len, num_heads = 1, 1, 4 + qk_nope_head_dim, qk_rope_head_dim = 32, 16 + kv_lora_rank = 64 + v_head_dim = 32 + max_seq_len = 32 + + for dtype in [torch.float16, torch.bfloat16]: + self.dtype = dtype + data = self._create_cached_mla_data( + batch_size, + seq_len, + num_heads, + qk_nope_head_dim, + qk_rope_head_dim, + kv_lora_rank, + v_head_dim, + max_seq_len, + ) + + output = self._run_cached_mla(data) + assert output.dtype == dtype, f"Expected dtype {dtype}, got {output.dtype}" + + def test_memory_efficiency(self): + """Test that cache uses compressed dimensions (no num_heads).""" + batch_size = 1 + max_seq_len = 1024 + kv_lora_rank = 512 + qk_rope_head_dim = 64 + num_heads = 128 # DeepSeek V3 + + # FlashInfer compressed cache size + compressed_cache_size = batch_size * max_seq_len * (kv_lora_rank + qk_rope_head_dim) + + # Expanded per-head cache size (what we avoid) + qk_nope_head_dim = 128 + v_head_dim = 128 + expanded_cache_size = ( + batch_size + * max_seq_len + * num_heads + * (qk_nope_head_dim + v_head_dim + qk_rope_head_dim) + ) + + # Verify compression ratio + compression_ratio = expanded_cache_size / compressed_cache_size + assert compression_ratio > 50, f"Expected >50x compression, got {compression_ratio:.1f}x" + + +class TestMLADescriptor: + """Test MultiHeadLatentAttention descriptor configuration.""" + + def _get_mla_descriptor(self): + """Get MLA descriptor from registry.""" + from tensorrt_llm._torch.auto_deploy.custom_ops.attention_interface import AttentionRegistry + + return AttentionRegistry.get("torch_mla") + + def test_descriptor_registration(self): + """Test that MLA descriptor is properly registered.""" + from tensorrt_llm._torch.auto_deploy.custom_ops.attention_interface import AttentionRegistry + + assert AttentionRegistry.has("torch_mla"), "torch_mla should be registered" + + def test_descriptor_layout(self): + """Test that MLA descriptor uses correct layout.""" + mla_descriptor = self._get_mla_descriptor() + + assert mla_descriptor.get_attention_layout() == "bsnd", "MLA should use bsnd layout" + + def test_descriptor_num_qkv_args(self): + """Test that MLA descriptor expects 5 tensor args.""" + mla_descriptor = self._get_mla_descriptor() + + assert mla_descriptor.get_num_qkv_args() == 5, ( + "MLA should expect 5 tensor args (q_nope, q_pe, compressed_kv, kpe, kv_b_proj_weight)" + ) + + def test_descriptor_source_op(self): + """Test that MLA descriptor points to correct source op.""" + mla_descriptor = self._get_mla_descriptor() + + source_op = mla_descriptor.get_source_attention_op() + assert source_op == torch.ops.auto_deploy.torch_mla, "MLA should use torch_mla as source op" + + def test_descriptor_cached_op(self): + """Test that MLA descriptor points to correct cached op.""" + mla_descriptor = self._get_mla_descriptor() + + cached_op = mla_descriptor.get_cached_attention_op() + assert cached_op == torch.ops.auto_deploy.torch_cached_mla_with_cache.default, ( + "MLA should use torch_cached_mla_with_cache as cached op" + ) + + def test_descriptor_standard_metadata(self): + """Test that MLA descriptor uses standard metadata args.""" + mla_descriptor = self._get_mla_descriptor() + + expected_args = ["batch_info_host", "seq_len", "input_pos", "slot_idx", "cu_seqlen"] + actual_args = mla_descriptor.get_standard_metadata_args() + assert actual_args == expected_args, ( + f"Expected standard metadata {expected_args}, got {actual_args}" + ) diff --git a/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_update_kv_cache.py b/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_update_kv_cache.py index 34baf42d43c9..b9a635f609e7 100644 --- a/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_update_kv_cache.py +++ b/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_update_kv_cache.py @@ -1,6 +1,6 @@ import torch -from tensorrt_llm._torch.auto_deploy.custom_ops.attention.torch_attention import update_kv_cache +from tensorrt_llm._torch.auto_deploy.custom_ops.attention.torch_backend_attention import _update_kv_cache def test_update_kv_cache(): @@ -26,7 +26,7 @@ def test_update_kv_cache(): print("slot_idx: " + str(torch.tensor([0, 1]))) print("seq_start: " + str(torch.tensor([0, 3]))) - update_kv_cache( + _update_kv_cache( k.view(batch_size * seq_length, n_heads, K_D_HEAD), v.view(batch_size * seq_length, n_heads, V_D_HEAD), k_cache, diff --git a/tests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_deepseek_custom.py b/tests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_deepseek_custom.py new file mode 100644 index 000000000000..7e8779941cf8 --- /dev/null +++ b/tests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_deepseek_custom.py @@ -0,0 +1,494 @@ +"""Testing custom DeepSeekV3 model implementation for auto_deploy export.""" + +import pytest +import torch +from transformers import PretrainedConfig + +from tensorrt_llm._torch.auto_deploy.models.custom.modeling_deepseek import ( + DeepSeekV3Attention, + DeepSeekV3DecoderLayer, + DeepSeekV3ForCausalLM, + DeepSeekV3MLP, + DeepSeekV3Model, + DeepSeekV3MoE, + DeepSeekV3RMSNorm, + DeepSeekV3RotaryEmbedding, + DeepSeekV3YarnRotaryEmbedding, +) + + +class MockDeepSeekConfig(PretrainedConfig): + """Mock DeepSeek config for testing the custom model components.""" + + model_type = "deepseek_v3" + + def __init__(self): + super().__init__() + # Attention config + self.num_attention_heads = 8 + self.qk_nope_head_dim = 64 + self.qk_rope_head_dim = 32 + self.v_head_dim = 64 + self.kv_lora_rank = 128 + self.q_lora_rank = None # No LoRA for Q in tests + self.hidden_size = 256 + self.rope_theta = 10000.0 + self.max_position_embeddings = 512 + self.attention_bias = False + self.rope_scaling = None + self.rms_norm_eps = 1e-6 + + # MLP config + self.intermediate_size = 512 + self.hidden_act = "silu" + + # MoE config + self.n_routed_experts = 4 + self.num_experts_per_tok = 2 + self.moe_intermediate_size = 256 + self.n_shared_experts = 1 + self.routed_scaling_factor = 1.0 + self.n_group = 1 + self.topk_group = 1 + + # Model config + self.num_hidden_layers = 2 + self.first_k_dense_replace = 1 # First layer is dense, second is MoE + self.moe_layer_freq = 1 + self.vocab_size = 1000 + self.pad_token_id = 0 + self.initializer_range = 0.02 + + +class TestDeepSeekV3RMSNorm: + """Test DeepSeekV3RMSNorm implementation.""" + + @pytest.fixture(autouse=True) + def setup_method(self): + self.device = "cuda" if torch.cuda.is_available() else "cpu" + self.dtype = torch.bfloat16 + torch.manual_seed(42) + + def test_forward_shape(self): + """Test that RMSNorm preserves input shape.""" + hidden_size = 256 + norm = DeepSeekV3RMSNorm(hidden_size).to(self.device, self.dtype) + + x = torch.randn(2, 4, hidden_size, dtype=self.dtype, device=self.device) + output = norm(x) + + assert output.shape == x.shape + assert torch.isfinite(output).all() + + def test_output_normalized(self): + """Test that output has approximately unit variance.""" + hidden_size = 256 + norm = DeepSeekV3RMSNorm(hidden_size).to(self.device, torch.float32) + + x = torch.randn(2, 4, hidden_size, dtype=torch.float32, device=self.device) + output = norm(x) + + # RMS should be close to 1 after normalization (scaled by weight) + rms = torch.sqrt((output**2).mean(-1)) + assert torch.allclose(rms, torch.ones_like(rms), atol=0.1) + + +class TestDeepSeekV3RotaryEmbedding: + """Test DeepSeekV3 Rotary Embedding implementations.""" + + @pytest.fixture(autouse=True) + def setup_method(self): + self.device = "cuda" if torch.cuda.is_available() else "cpu" + self.dtype = torch.bfloat16 + torch.manual_seed(42) + + def test_base_rope_shape(self): + """Test base rotary embedding output shape.""" + dim = 32 + max_pos = 512 + rope = DeepSeekV3RotaryEmbedding(dim, max_pos).to(self.device) + + x = torch.randn(2, 4, 8, dim, dtype=self.dtype, device=self.device) + cos, sin = rope(x) + + # Should return full cached values + assert cos.shape == (max_pos, dim) + assert sin.shape == (max_pos, dim) + + def test_yarn_rope_shape(self): + """Test YaRN rotary embedding output shape.""" + dim = 32 + max_pos = 512 + rope = DeepSeekV3YarnRotaryEmbedding( + dim, + max_pos, + scaling_factor=2.0, + original_max_position_embeddings=256, + ).to(self.device) + + x = torch.randn(2, 4, 8, dim, dtype=self.dtype, device=self.device) + cos, sin = rope(x) + + assert cos.shape == (max_pos, dim) + assert sin.shape == (max_pos, dim) + + +class TestDeepSeekV3MLP: + """Test DeepSeekV3MLP implementation.""" + + @pytest.fixture(autouse=True) + def setup_method(self): + self.device = "cuda" if torch.cuda.is_available() else "cpu" + self.dtype = torch.bfloat16 + torch.manual_seed(42) + + def test_forward_shape(self): + """Test MLP output shape.""" + config = MockDeepSeekConfig() + mlp = DeepSeekV3MLP(config).to(self.device, self.dtype) + + x = torch.randn(2, 4, config.hidden_size, dtype=self.dtype, device=self.device) + output = mlp(x) + + assert output.shape == x.shape + assert torch.isfinite(output).all() + + +class TestDeepSeekV3Attention: + """Test DeepSeekV3Attention (MLA) implementation.""" + + @pytest.fixture(autouse=True) + def setup_method(self): + self.device = "cuda" if torch.cuda.is_available() else "cpu" + self.dtype = torch.bfloat16 + torch.manual_seed(42) + + def test_forward_shape(self): + """Test attention output shape.""" + config = MockDeepSeekConfig() + attn = DeepSeekV3Attention(config, layer_idx=0).to(self.device, self.dtype) + + batch_size, seq_len = 2, 4 + hidden_states = torch.randn( + batch_size, seq_len, config.hidden_size, dtype=self.dtype, device=self.device + ) + position_ids = torch.arange(seq_len, device=self.device).unsqueeze(0).expand(batch_size, -1) + + output = attn(hidden_states, position_ids) + + assert output.shape == hidden_states.shape + assert torch.isfinite(output).all() + + def test_different_batch_sizes(self): + """Test attention with different batch sizes.""" + config = MockDeepSeekConfig() + attn = DeepSeekV3Attention(config, layer_idx=0).to(self.device, self.dtype) + + for batch_size in [1, 2, 4]: + seq_len = 4 + hidden_states = torch.randn( + batch_size, seq_len, config.hidden_size, dtype=self.dtype, device=self.device + ) + position_ids = ( + torch.arange(seq_len, device=self.device).unsqueeze(0).expand(batch_size, -1) + ) + + output = attn(hidden_states, position_ids) + assert output.shape == (batch_size, seq_len, config.hidden_size) + + def test_different_sequence_lengths(self): + """Test attention with different sequence lengths.""" + config = MockDeepSeekConfig() + attn = DeepSeekV3Attention(config, layer_idx=0).to(self.device, self.dtype) + + for seq_len in [1, 4, 16]: + batch_size = 2 + hidden_states = torch.randn( + batch_size, seq_len, config.hidden_size, dtype=self.dtype, device=self.device + ) + position_ids = ( + torch.arange(seq_len, device=self.device).unsqueeze(0).expand(batch_size, -1) + ) + + output = attn(hidden_states, position_ids) + assert output.shape == (batch_size, seq_len, config.hidden_size) + + +class TestDeepSeekV3MoE: + """Test DeepSeekV3MoE implementation.""" + + @pytest.fixture(autouse=True) + def setup_method(self): + self.device = "cuda" if torch.cuda.is_available() else "cpu" + self.dtype = torch.bfloat16 + torch.manual_seed(42) + + def test_forward_shape(self): + """Test MoE output shape.""" + config = MockDeepSeekConfig() + moe = DeepSeekV3MoE(config).to(self.device, self.dtype) + + x = torch.randn(2, 4, config.hidden_size, dtype=self.dtype, device=self.device) + output = moe(x) + + assert output.shape == x.shape + assert torch.isfinite(output).all() + + def test_with_shared_experts(self): + """Test MoE with shared experts.""" + config = MockDeepSeekConfig() + config.n_shared_experts = 2 + moe = DeepSeekV3MoE(config).to(self.device, self.dtype) + + assert moe.shared_experts is not None + + x = torch.randn(2, 4, config.hidden_size, dtype=self.dtype, device=self.device) + output = moe(x) + + assert output.shape == x.shape + + def test_without_shared_experts(self): + """Test MoE without shared experts.""" + config = MockDeepSeekConfig() + config.n_shared_experts = None + moe = DeepSeekV3MoE(config).to(self.device, self.dtype) + + assert moe.shared_experts is None + + x = torch.randn(2, 4, config.hidden_size, dtype=self.dtype, device=self.device) + output = moe(x) + + assert output.shape == x.shape + + +class TestDeepSeekV3DecoderLayer: + """Test DeepSeekV3DecoderLayer implementation.""" + + @pytest.fixture(autouse=True) + def setup_method(self): + self.device = "cuda" if torch.cuda.is_available() else "cpu" + self.dtype = torch.bfloat16 + torch.manual_seed(42) + + def test_dense_layer(self): + """Test decoder layer with dense MLP.""" + config = MockDeepSeekConfig() + # Layer 0 should be dense (before first_k_dense_replace) + layer = DeepSeekV3DecoderLayer(config, layer_idx=0).to(self.device, self.dtype) + + assert isinstance(layer.mlp, DeepSeekV3MLP) + + batch_size, seq_len = 2, 4 + hidden_states = torch.randn( + batch_size, seq_len, config.hidden_size, dtype=self.dtype, device=self.device + ) + position_ids = torch.arange(seq_len, device=self.device).unsqueeze(0).expand(batch_size, -1) + + output = layer(hidden_states, position_ids) + + assert output.shape == hidden_states.shape + + def test_moe_layer(self): + """Test decoder layer with MoE.""" + config = MockDeepSeekConfig() + # Layer 1 should be MoE (at first_k_dense_replace) + layer = DeepSeekV3DecoderLayer(config, layer_idx=1).to(self.device, self.dtype) + + assert isinstance(layer.mlp, DeepSeekV3MoE) + + batch_size, seq_len = 2, 4 + hidden_states = torch.randn( + batch_size, seq_len, config.hidden_size, dtype=self.dtype, device=self.device + ) + position_ids = torch.arange(seq_len, device=self.device).unsqueeze(0).expand(batch_size, -1) + + output = layer(hidden_states, position_ids) + + assert output.shape == hidden_states.shape + + +class TestDeepSeekV3Model: + """Test DeepSeekV3Model implementation.""" + + @pytest.fixture(autouse=True) + def setup_method(self): + self.device = "cuda" if torch.cuda.is_available() else "cpu" + self.dtype = torch.bfloat16 + torch.manual_seed(42) + + def test_forward_with_input_ids(self): + """Test model forward with input_ids.""" + config = MockDeepSeekConfig() + model = DeepSeekV3Model(config).to(self.device, self.dtype) + + batch_size, seq_len = 2, 4 + input_ids = torch.randint(0, config.vocab_size, (batch_size, seq_len), device=self.device) + + output = model(input_ids=input_ids) + + assert output.last_hidden_state.shape == (batch_size, seq_len, config.hidden_size) + + def test_forward_with_inputs_embeds(self): + """Test model forward with inputs_embeds.""" + config = MockDeepSeekConfig() + model = DeepSeekV3Model(config).to(self.device, self.dtype) + + batch_size, seq_len = 2, 4 + inputs_embeds = torch.randn( + batch_size, seq_len, config.hidden_size, dtype=self.dtype, device=self.device + ) + + output = model(inputs_embeds=inputs_embeds) + + assert output.last_hidden_state.shape == inputs_embeds.shape + + +class TestDeepSeekV3ForCausalLM: + """Test DeepSeekV3ForCausalLM implementation.""" + + @pytest.fixture(autouse=True) + def setup_method(self): + self.device = "cuda" if torch.cuda.is_available() else "cpu" + self.dtype = torch.bfloat16 + torch.manual_seed(42) + + def test_forward(self): + """Test causal LM forward pass.""" + config = MockDeepSeekConfig() + model = DeepSeekV3ForCausalLM(config).to(self.device, self.dtype) + + batch_size, seq_len = 2, 4 + input_ids = torch.randint(0, config.vocab_size, (batch_size, seq_len), device=self.device) + + output = model(input_ids=input_ids) + + assert output.logits.shape == (batch_size, seq_len, config.vocab_size) + + def test_output_dtype(self): + """Test that logits are float32.""" + config = MockDeepSeekConfig() + model = DeepSeekV3ForCausalLM(config).to(self.device, self.dtype) + + batch_size, seq_len = 2, 4 + input_ids = torch.randint(0, config.vocab_size, (batch_size, seq_len), device=self.device) + + output = model(input_ids=input_ids) + + # Logits should be float32 for numerical stability + assert output.logits.dtype == torch.float32 + + +class TestMLAOpRegistration: + """Test that MLA ops are properly registered.""" + + def test_torch_mla_registered(self): + """Test that torch_mla op is registered.""" + assert hasattr(torch.ops.auto_deploy, "torch_mla"), "torch_mla op should be registered" + + def test_torch_cached_mla_registered(self): + """Test that torch_cached_mla_with_cache op is registered.""" + assert hasattr(torch.ops.auto_deploy, "torch_cached_mla_with_cache"), ( + "torch_cached_mla_with_cache op should be registered" + ) + + def test_torch_mla_callable(self): + """Test that torch_mla op is callable.""" + device = "cuda" if torch.cuda.is_available() else "cpu" + dtype = torch.bfloat16 + + batch_size, seq_len, num_heads = 1, 2, 4 + qk_nope_head_dim, qk_rope_head_dim = 64, 32 + kv_lora_rank = 128 + v_head_dim = 64 + + q_nope = torch.randn( + batch_size, seq_len, num_heads, qk_nope_head_dim, dtype=dtype, device=device + ) + q_pe = torch.randn( + batch_size, seq_len, num_heads, qk_rope_head_dim, dtype=dtype, device=device + ) + compressed_kv = torch.randn(batch_size, seq_len, kv_lora_rank, dtype=dtype, device=device) + kpe = torch.randn(batch_size, seq_len, 1, qk_rope_head_dim, dtype=dtype, device=device) + kv_b_proj_weight = torch.randn( + num_heads * (qk_nope_head_dim + v_head_dim), kv_lora_rank, dtype=dtype, device=device + ) + + # Should not raise + output = torch.ops.auto_deploy.torch_mla( + q_nope, q_pe, compressed_kv, kpe, kv_b_proj_weight, True, None, "bsnd" + ) + assert output is not None + + def test_torch_cached_mla_callable(self): + """Test that torch_cached_mla_with_cache op is callable.""" + device = "cuda" if torch.cuda.is_available() else "cpu" + dtype = torch.bfloat16 + + batch_size, seq_len, num_heads = 1, 1, 4 + qk_nope_head_dim, qk_rope_head_dim = 32, 16 + kv_lora_rank = 64 + v_head_dim = 32 + max_seq_len = 32 + + q_nope = torch.randn( + batch_size, seq_len, num_heads, qk_nope_head_dim, dtype=dtype, device=device + ) + q_pe = torch.randn( + batch_size, seq_len, num_heads, qk_rope_head_dim, dtype=dtype, device=device + ) + compressed_kv = torch.randn(batch_size, seq_len, kv_lora_rank, dtype=dtype, device=device) + kpe = torch.randn(batch_size, seq_len, 1, qk_rope_head_dim, dtype=dtype, device=device) + kv_b_proj_weight = torch.randn( + num_heads * (qk_nope_head_dim + v_head_dim), kv_lora_rank, dtype=dtype, device=device + ) + + batch_info_host = torch.tensor([0, 0, batch_size], dtype=torch.int32, device=device) + seq_len_tensor = torch.tensor([seq_len], dtype=torch.int32, device=device) + input_pos = torch.tensor([0], dtype=torch.int32, device=device) + cache_loc = torch.tensor([0], dtype=torch.int32, device=device) + cu_seqlen = torch.tensor([0], dtype=torch.int32, device=device) + + mla_cache = torch.zeros( + batch_size, max_seq_len, kv_lora_rank + qk_rope_head_dim, dtype=dtype, device=device + ) + + # Should not raise + output = torch.ops.auto_deploy.torch_cached_mla_with_cache( + q_nope, + q_pe, + compressed_kv, + kpe, + kv_b_proj_weight, + batch_info_host, + seq_len_tensor, + input_pos, + cache_loc, + cu_seqlen, + mla_cache, + None, + kv_lora_rank, + ) + assert output is not None + + +class TestMoEOpRegistration: + """Test that MoE ops are properly registered.""" + + def test_torch_moe_registered(self): + """Test that torch_moe op is registered.""" + assert hasattr(torch.ops.auto_deploy, "torch_moe"), "torch_moe op should be registered" + + +class TestCustomModelRegistration: + """Test that custom model is properly registered.""" + + def test_model_registered(self): + """Test that DeepSeekV3ForCausalLM is registered with the factory.""" + from tensorrt_llm._torch.auto_deploy.models.hf import AutoModelForCausalLMFactory + + # Check that the model is registered in the custom model mapping + assert "DeepseekV3Config" in AutoModelForCausalLMFactory._custom_model_mapping + assert ( + AutoModelForCausalLMFactory._custom_model_mapping["DeepseekV3Config"] + == DeepSeekV3ForCausalLM + ) diff --git a/tests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_deepseek_patches.py b/tests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_deepseek_patches.py deleted file mode 100644 index 070d85a5984c..000000000000 --- a/tests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_deepseek_patches.py +++ /dev/null @@ -1,110 +0,0 @@ -"""Testing module patches that enable export of deepseek model.""" - -import types - -import pytest -import torch -from test_common.llm_data import hf_id_to_local_model_dir -from transformers import AutoConfig, AutoModelForCausalLM - -from tensorrt_llm._torch.auto_deploy.models.patches.deepseek import ( - deepseek_v3_attention, - deepseek_v3_moe_exact, -) - - -def _load_layer_from_model(model_name_or_path, layer_name): - """ - Loads a specific layer/module from a model without loading the entire model. - - Parameters: - model_name_or_path (str): Path or name of the pretrained model. - layer_name (str): Name of the layer to extract. - - Returns: - module: The specified layer/module if available, otherwise None. - """ - try: - # Load only the model configuration - config = AutoConfig.from_pretrained(model_name_or_path, trust_remote_code=True) - - # Load a subset of layers of the model and configure yarn - config.num_hidden_layers = 1 - config.use_cache = False - config.first_k_dense_replace = 0 - config.n_routed_experts = 2 - config.num_experts_per_tok = 1 - config.n_group = 1 - config.topk_group = 1 - config.hidden_size = 8 - config.moe_intermediate_size = 8 - config.num_attention_heads = 2 - config.num_key_value_heads = 2 - config.qk_nope_head_dim = 4 - config.qk_rope_head_dim = 2 - config.v_head_dim = 4 - config.intermediate_size = 8 - config.max_position_embeddings = 7 - - config.rope_scaling = None - - # Build the model architecture (no weights loaded yet) - model = AutoModelForCausalLM.from_config(config, trust_remote_code=True) - model.eval() - - # Access the specific layer by its name - module = dict(model.named_modules()).get(layer_name) - if module is None: - print(f"Layer '{layer_name}' not found in the model.") - else: - print(f"Successfully extracted layer '{layer_name}'.") - return module - except Exception as e: - print(f"Error extracting layer: {e}") - return None - - -def _generate_ds_attention_mask(b, s): - return torch.where( - torch.tril(torch.full((s, s), float("-inf"))).unsqueeze(0).unsqueeze(0).expand(b, 1, s, s) - == float("-inf"), - torch.tensor(0.0), - torch.tensor(float(-3.4028e38)), - ) - - -@pytest.mark.parametrize( - "model_name, module_name, patch, inputs", - [ - pytest.param( - hf_id_to_local_model_dir("deepseek-ai/DeepSeek-R1"), - "model.layers.0.self_attn", - deepseek_v3_attention, - [ - torch.randn(2, 6, 8, dtype=torch.bfloat16), - _generate_ds_attention_mask(2, 6), - torch.tensor([[0, 1, 2, 3, 4, 5]]), - ], - ), # attention requires inputs [hidden_states, attention_mask, position_ids] - pytest.param( - hf_id_to_local_model_dir("deepseek-ai/DeepSeek-R1"), - "model.layers.0.mlp", - deepseek_v3_moe_exact, - [torch.randn(2, 6, 8, dtype=torch.bfloat16)], - ), # moe requires inputs [hidden_states] - ], -) -def test_module_patches(model_name, module_name, patch, inputs): - # Get module - module = _load_layer_from_model(model_name, module_name) - - # Pass test inputs to generate reference - ref, *_ = module(*inputs) - - # Patch layer - module.forward = types.MethodType(patch, module) - - # Generate test output - test, *_ = module(*inputs) - - torch.allclose(ref, test, atol=0, rtol=0) diff --git a/tests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_glm4_moe_lite_modeling.py b/tests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_glm4_moe_lite_modeling.py new file mode 100644 index 000000000000..486ffe7d3137 --- /dev/null +++ b/tests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_glm4_moe_lite_modeling.py @@ -0,0 +1,652 @@ +"""Tests for GLM4 MoE Lite custom model implementation. + +This module tests the custom GLM4 MoE Lite model implementation which uses +auto_deploy custom ops (torch_mla, torch_moe, etc.) for export compatibility. +""" + +import pytest +import torch +from torch.export import Dim + +from tensorrt_llm._torch.auto_deploy.export import torch_export_to_gm +from tensorrt_llm._torch.auto_deploy.models.custom.modeling_glm4_moe_lite import ( + Glm4MoeLiteConfig, + Glm4MoeLiteForCausalLM, + Glm4MoeLiteMLP, + Glm4MoeLiteMoE, +) +from tensorrt_llm._torch.auto_deploy.utils._graph import move_to_device + +_BATCH_AND_SEQUENCE_TEST_CASES = ((2, 6), (1, 8)) + + +@pytest.fixture(scope="function", autouse=True) +def set_seed(): + torch.manual_seed(42) + + +def _create_small_config() -> Glm4MoeLiteConfig: + """Create a small GLM4 MoE Lite config for testing.""" + return Glm4MoeLiteConfig( + vocab_size=1000, + hidden_size=64, + intermediate_size=128, + num_hidden_layers=3, # Layer 0 dense, layers 1-2 MoE + num_attention_heads=4, + num_key_value_heads=4, + hidden_act="silu", + max_position_embeddings=512, + rms_norm_eps=1e-5, + # MLA params (scaled down) + q_lora_rank=32, + kv_lora_rank=32, + qk_nope_head_dim=8, + qk_rope_head_dim=8, + v_head_dim=16, + # MoE params (scaled down) + n_routed_experts=4, + n_shared_experts=1, + num_experts_per_tok=2, + moe_intermediate_size=32, + n_group=1, + topk_group=1, + routed_scaling_factor=1.0, + norm_topk_prob=True, + first_k_dense_replace=1, + # RoPE + rope_theta=10000.0, + rope_scaling=None, + # Other + attention_bias=False, + attention_dropout=0.0, + pad_token_id=0, + ) + + +def _create_moe_layer(config: Glm4MoeLiteConfig) -> Glm4MoeLiteMoE: + """Create a MoE layer from config.""" + moe = Glm4MoeLiteMoE(config) + # Initialize gate weights with randn for reproducibility + # (gate weight is initialized with torch.empty which isn't seeded) + moe.gate.weight = torch.nn.Parameter(torch.randn_like(moe.gate.weight)) + return moe + + +@pytest.mark.parametrize("B,S", _BATCH_AND_SEQUENCE_TEST_CASES) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@torch.no_grad() +def test_glm4_moe_lite_moe_layer(B, S, dtype): + """Test that MoE layer produces valid output.""" + device = "cuda" + config = _create_small_config() + + moe = _create_moe_layer(config) + moe.to(device=device, dtype=dtype) + moe.eval() + + H = config.hidden_size + x = torch.randn(B, S, H, device=device, dtype=dtype) + + output = moe(x) + + # Check output shape matches input shape + assert output.shape == x.shape, f"Expected shape {x.shape}, got {output.shape}" + + # Check output is not all zeros (MoE should transform the input) + assert not torch.allclose(output, torch.zeros_like(output)), "Output should not be all zeros" + + # Check output doesn't have NaN or Inf values + assert not torch.isnan(output).any(), "Output contains NaN values" + assert not torch.isinf(output).any(), "Output contains Inf values" + + +@pytest.mark.parametrize("B,S", _BATCH_AND_SEQUENCE_TEST_CASES) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@torch.no_grad() +def test_glm4_moe_lite_full_model(B, S, dtype): + """Test that full model produces valid output.""" + device = "cuda" + config = _create_small_config() + + model = Glm4MoeLiteForCausalLM(config) + model.to(device=device, dtype=dtype) + model.eval() + + # Create input + input_ids = torch.randint(0, config.vocab_size, (B, S), device=device) + position_ids = torch.arange(S, device=device).unsqueeze(0).expand(B, -1) + + output = model(input_ids=input_ids, position_ids=position_ids) + + # Check output shape + assert output.logits.shape == (B, S, config.vocab_size), ( + f"Expected logits shape {(B, S, config.vocab_size)}, got {output.logits.shape}" + ) + + # Check output doesn't have NaN or Inf values + assert not torch.isnan(output.logits).any(), "Logits contain NaN values" + assert not torch.isinf(output.logits).any(), "Logits contain Inf values" + + +def test_glm4_moe_lite_model_can_be_exported(): + """Test that the custom model can be exported with torch_export_to_gm. + + This test verifies: + 1. The model exports successfully without graph breaks + 2. The exported graph module produces outputs with correct shape + 3. The outputs contain finite values (no NaN/Inf) + + Note: We don't test numerical equivalence between original and exported model + here because torch.export lifts parameters into the graph, creating a different + parameter structure that doesn't match load_state_dict. The numerical correctness + of the model itself is already validated by test_glm4_moe_lite_full_model. + """ + device = "cuda" + dtype = torch.bfloat16 + config = _create_small_config() + + model = Glm4MoeLiteForCausalLM(config) + model.to(device=device, dtype=dtype) + model.eval() + + # Create input + B, S = 2, 8 + input_ids = torch.randint(0, config.vocab_size, (B, S), device=device) + position_ids = torch.arange(S, device=device).unsqueeze(0).expand(B, -1) + + # Define dynamic shapes + batch_size_dynamic = Dim.DYNAMIC + seq_len_dynamic = Dim.DYNAMIC + dynamic_shapes = ( + {0: batch_size_dynamic, 1: seq_len_dynamic}, + {0: batch_size_dynamic, 1: seq_len_dynamic}, + ) + + # Export the model - this is the main test: verify no graph breaks + gm = torch_export_to_gm( + model, + args=tuple(), + kwargs={"input_ids": input_ids, "position_ids": position_ids}, + dynamic_shapes=dynamic_shapes, + ) + + # Move graph module to device + move_to_device(gm, device) + + # Verify the exported model produces valid output + with torch.inference_mode(): + out_gm = gm(input_ids=input_ids, position_ids=position_ids) + + # Check output structure and shape + assert "logits" in out_gm, "Output should contain 'logits' key" + logits = out_gm["logits"] + assert logits.shape == (B, S, config.vocab_size), ( + f"Expected shape {(B, S, config.vocab_size)}, got {logits.shape}" + ) + assert torch.isfinite(logits).all(), "Logits should not contain NaN or Inf" + + # Test with different input shape to verify dynamic shapes work + B2, S2 = 1, 4 + input_ids2 = torch.randint(0, config.vocab_size, (B2, S2), device=device) + position_ids2 = torch.arange(S2, device=device).unsqueeze(0).expand(B2, -1) + + with torch.inference_mode(): + out_gm2 = gm(input_ids=input_ids2, position_ids=position_ids2) + + logits2 = out_gm2["logits"] + expected_shape = (B2, S2, config.vocab_size) + assert logits2.shape == expected_shape, ( + f"Dynamic shape test failed: expected {expected_shape}, got {logits2.shape}" + ) + assert torch.isfinite(logits2).all(), "Logits should not contain NaN or Inf" + + +def test_glm4_moe_lite_config_registration(): + """Test that the config is properly registered or model_type is correct.""" + # Create a config and verify model_type + config = _create_small_config() + assert config.model_type == "glm4_moe_lite" + + # Verify our config class can be instantiated with expected attributes + assert hasattr(config, "hidden_size") + assert hasattr(config, "num_attention_heads") + assert hasattr(config, "n_routed_experts") + assert hasattr(config, "kv_lora_rank") + assert hasattr(config, "qk_rope_head_dim") + + +def test_glm4_moe_lite_layer_types(): + """Test that layer 0 uses dense MLP and later layers use MoE.""" + config = _create_small_config() + model = Glm4MoeLiteForCausalLM(config) + + # Check layer 0 (should be dense MLP, not MoE) + layer0_mlp = model.model.layers[0].mlp + assert type(layer0_mlp).__name__ == "Glm4MoeLiteMLP", ( + f"Layer 0 should use Glm4MoeLiteMLP, got {type(layer0_mlp).__name__}" + ) + + # Check layer 1+ (should be MoE) + for i in range(1, config.num_hidden_layers): + layer_mlp = model.model.layers[i].mlp + assert type(layer_mlp).__name__ == "Glm4MoeLiteMoE", ( + f"Layer {i} should use Glm4MoeLiteMoE, got {type(layer_mlp).__name__}" + ) + + +def test_glm4_moe_lite_expert_structure(): + """Test that experts have correct structure for checkpoint loading.""" + config = _create_small_config() + moe = Glm4MoeLiteMoE(config) + + # Check that experts is a ModuleList + assert isinstance(moe.experts, torch.nn.ModuleList), "experts should be nn.ModuleList" + + # Check number of experts + assert len(moe.experts) == config.n_routed_experts, ( + f"Expected {config.n_routed_experts} experts, got {len(moe.experts)}" + ) + + # Check each expert has the correct structure + for i, expert in enumerate(moe.experts): + assert hasattr(expert, "gate_proj"), f"Expert {i} missing gate_proj" + assert hasattr(expert, "up_proj"), f"Expert {i} missing up_proj" + assert hasattr(expert, "down_proj"), f"Expert {i} missing down_proj" + + # Check state_dict keys match expected checkpoint format + state_dict = moe.state_dict() + expected_keys = [ + "experts.0.gate_proj.weight", + "experts.0.up_proj.weight", + "experts.0.down_proj.weight", + ] + for key in expected_keys: + assert key in state_dict, ( + f"Expected key '{key}' in state_dict, got keys: {list(state_dict.keys())[:10]}..." + ) + + +# ============================================================================= +# Numerical Equivalence Tests +# These tests compare our custom implementation against the HF implementation +# ============================================================================= + + +def _get_hf_model_class(): + """Get the HF model class for GLM4 MoE Lite. + + Returns None if transformers doesn't have glm4_moe_lite (older versions). + """ + try: + from transformers.models.glm4_moe_lite.modeling_glm4_moe_lite import ( + Glm4MoeLiteForCausalLM as HFGlm4MoeLiteForCausalLM, + ) + + return HFGlm4MoeLiteForCausalLM + except ImportError: + return None + + +def _get_hf_moe_class(): + """Get the HF MoE class for GLM4 MoE Lite.""" + try: + from transformers.models.glm4_moe_lite.modeling_glm4_moe_lite import ( + Glm4MoeLiteMoE as HFGlm4MoeLiteMoE, + ) + + return HFGlm4MoeLiteMoE + except ImportError: + return None + + +def _get_hf_attention_class(): + """Get the HF Attention class for GLM4 MoE Lite.""" + try: + from transformers.models.glm4_moe_lite.modeling_glm4_moe_lite import ( + Glm4MoeLiteAttention as HFGlm4MoeLiteAttention, + ) + + return HFGlm4MoeLiteAttention + except ImportError: + return None + + +def _get_hf_mlp_class(): + """Get the HF MLP class for GLM4 MoE Lite.""" + try: + from transformers.models.glm4_moe_lite.modeling_glm4_moe_lite import ( + Glm4MoeLiteMLP as HFGlm4MoeLiteMLP, + ) + + return HFGlm4MoeLiteMLP + except ImportError: + return None + + +def _get_hf_config_class(): + """Get the HF Config class for GLM4 MoE Lite.""" + try: + from transformers.models.glm4_moe_lite.configuration_glm4_moe_lite import ( + Glm4MoeLiteConfig as HFGlm4MoeLiteConfig, + ) + + return HFGlm4MoeLiteConfig + except ImportError: + return None + + +def _convert_hf_moe_state_dict_to_custom(hf_state_dict: dict, n_experts: int) -> dict: + """Convert HF MoE state dict to our custom format. + + HF format (stacked): + experts.gate_up_proj: [n_experts, 2 * intermediate_size, hidden_size] + experts.down_proj: [n_experts, hidden_size, intermediate_size] + + Custom format (per-expert): + experts.0.gate_proj.weight: [intermediate_size, hidden_size] + experts.0.up_proj.weight: [intermediate_size, hidden_size] + experts.0.down_proj.weight: [hidden_size, intermediate_size] + """ + custom_state_dict = {} + + for key, value in hf_state_dict.items(): + if key == "experts.gate_up_proj": + # Split stacked gate_up into individual gate and up per expert + # Shape: [n_experts, 2 * intermediate_size, hidden_size] + intermediate_size = value.shape[1] // 2 + for i in range(n_experts): + gate_up = value[i] # [2 * intermediate_size, hidden_size] + gate_weight = gate_up[:intermediate_size] # [intermediate_size, hidden_size] + up_weight = gate_up[intermediate_size:] # [intermediate_size, hidden_size] + custom_state_dict[f"experts.{i}.gate_proj.weight"] = gate_weight + custom_state_dict[f"experts.{i}.up_proj.weight"] = up_weight + elif key == "experts.down_proj": + # Split stacked down into individual down per expert + # Shape: [n_experts, hidden_size, intermediate_size] + for i in range(n_experts): + custom_state_dict[f"experts.{i}.down_proj.weight"] = value[i] + else: + # Copy other keys as-is + custom_state_dict[key] = value + + return custom_state_dict + + +def _convert_hf_full_model_state_dict_to_custom(hf_state_dict: dict, config) -> dict: + """Convert full HF model state dict to custom format. + + Handles MoE expert weight conversion for all MoE layers. + """ + custom_state_dict = {} + n_experts = config.n_routed_experts + + for key, value in hf_state_dict.items(): + # Check if this is an MoE expert weight + if ".mlp.experts.gate_up_proj" in key: + # Extract layer prefix (e.g., "model.layers.1.mlp.") + prefix = key.replace("experts.gate_up_proj", "") + intermediate_size = value.shape[1] // 2 + for i in range(n_experts): + gate_up = value[i] + gate_weight = gate_up[:intermediate_size] + up_weight = gate_up[intermediate_size:] + custom_state_dict[f"{prefix}experts.{i}.gate_proj.weight"] = gate_weight + custom_state_dict[f"{prefix}experts.{i}.up_proj.weight"] = up_weight + elif ".mlp.experts.down_proj" in key: + prefix = key.replace("experts.down_proj", "") + for i in range(n_experts): + custom_state_dict[f"{prefix}experts.{i}.down_proj.weight"] = value[i] + else: + # Copy other keys as-is + custom_state_dict[key] = value + + return custom_state_dict + + +def _create_hf_config(): + """Create HF config that matches our test config.""" + HFConfig = _get_hf_config_class() + if HFConfig is None: + return None + + config = HFConfig( + vocab_size=1000, + hidden_size=64, + intermediate_size=128, + num_hidden_layers=3, + num_attention_heads=4, + num_key_value_heads=4, + hidden_act="silu", + max_position_embeddings=512, + rms_norm_eps=1e-5, + # MLA params + q_lora_rank=32, + kv_lora_rank=32, + qk_nope_head_dim=8, + qk_rope_head_dim=8, + v_head_dim=16, + # MoE params + n_routed_experts=4, + n_shared_experts=1, + num_experts_per_tok=2, + moe_intermediate_size=32, + n_group=1, + topk_group=1, + routed_scaling_factor=1.0, + norm_topk_prob=True, + first_k_dense_replace=1, + # RoPE + rope_theta=10000.0, + rope_scaling=None, + # Other + attention_bias=False, + attention_dropout=0.0, + pad_token_id=0, + ) + + # Set internal attributes needed by HF's MoE implementation + # _experts_implementation tells HF which expert forward function to use + config._experts_implementation = "eager" + + return config + + +@pytest.mark.parametrize("B,S", _BATCH_AND_SEQUENCE_TEST_CASES) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@torch.no_grad() +def test_glm4_moe_lite_moe_numerical_equivalence(B, S, dtype): + """Test MoE layer produces numerically equivalent output to HF implementation.""" + HFMoE = _get_hf_moe_class() + if HFMoE is None: + pytest.skip("transformers doesn't have glm4_moe_lite (requires v5.0+)") + + device = "cuda" + config = _create_small_config() + hf_config = _create_hf_config() + + # Create HF MoE + hf_moe = HFMoE(hf_config) + hf_moe.to(device=device, dtype=dtype) + hf_moe.eval() + + # Initialize weights with randn for reproducibility + # (HF uses torch.empty which may be zeros, causing NaN in computation) + hf_moe.gate.weight = torch.nn.Parameter(torch.randn_like(hf_moe.gate.weight)) + hf_moe.experts.gate_up_proj = torch.nn.Parameter(torch.randn_like(hf_moe.experts.gate_up_proj)) + hf_moe.experts.down_proj = torch.nn.Parameter(torch.randn_like(hf_moe.experts.down_proj)) + + # Create custom MoE and load converted weights + custom_moe = Glm4MoeLiteMoE(config) + custom_moe.to(device=device, dtype=dtype) + + # Convert HF stacked expert weights to our per-expert format + hf_state_dict = hf_moe.state_dict() + + # Debug: print state dict keys and shapes + print("\n=== HF MoE state_dict keys and shapes ===") + for k, v in hf_state_dict.items(): + print(f" {k}: {v.shape}") + + custom_state_dict = _convert_hf_moe_state_dict_to_custom(hf_state_dict, config.n_routed_experts) + + print("\n=== Converted custom state_dict keys and shapes ===") + for k, v in custom_state_dict.items(): + print(f" {k}: {v.shape}") + + print("\n=== Expected custom MoE state_dict keys ===") + for k, v in custom_moe.state_dict().items(): + print(f" {k}: {v.shape}") + + custom_moe.load_state_dict(custom_state_dict) + custom_moe.eval() + + # Sanity check: verify expert weights match after conversion + # HF has stacked weights: experts.gate_up_proj [n_experts, 2*intermediate, hidden] + # Our model has per-expert: experts.{i}.gate_proj.weight, experts.{i}.up_proj.weight + hf_gate_up = hf_moe.experts.gate_up_proj # [n_experts, 2*intermediate, hidden] + hf_down = hf_moe.experts.down_proj # [n_experts, hidden, intermediate] + intermediate_size = config.moe_intermediate_size + + print(f"\n=== Debug: intermediate_size = {intermediate_size} ===") + print(f"hf_gate_up shape: {hf_gate_up.shape}") + print(f"hf_gate_up[0, :2, :2]: {hf_gate_up[0, :2, :2]}") + + # Get the converted state dict values for comparison + converted_gate_0 = custom_state_dict["experts.0.gate_proj.weight"] + print(f"converted_gate_0 shape: {converted_gate_0.shape}") + print(f"converted_gate_0[:2, :2]: {converted_gate_0[:2, :2]}") + + # After load_state_dict + loaded_gate_0 = custom_moe.experts[0].gate_proj.weight + print(f"loaded_gate_0 shape: {loaded_gate_0.shape}") + print(f"loaded_gate_0[:2, :2]: {loaded_gate_0[:2, :2]}") + + for i in range(config.n_routed_experts): + # Check gate_proj + hf_gate = hf_gate_up[i, :intermediate_size, :] # [intermediate, hidden] + custom_gate = custom_moe.experts[i].gate_proj.weight + torch.testing.assert_close( + custom_gate, hf_gate, msg=f"Expert {i} gate_proj weights don't match" + ) + + # Check up_proj + hf_up = hf_gate_up[i, intermediate_size:, :] # [intermediate, hidden] + custom_up = custom_moe.experts[i].up_proj.weight + torch.testing.assert_close(custom_up, hf_up, msg=f"Expert {i} up_proj weights don't match") + + # Check down_proj + hf_down_i = hf_down[i] # [hidden, intermediate] + custom_down = custom_moe.experts[i].down_proj.weight + torch.testing.assert_close( + custom_down, hf_down_i, msg=f"Expert {i} down_proj weights don't match" + ) + + # Also verify gate weights match + torch.testing.assert_close( + custom_moe.gate.weight, hf_moe.gate.weight, msg="Gate weights don't match" + ) + + # Create input + H = config.hidden_size + x = torch.randn(B, S, H, device=device, dtype=dtype) + + # Run both + hf_out = hf_moe(x) + custom_out = custom_moe(x) + + # Handle tuple output from HF (output, router_logits) + if isinstance(hf_out, tuple): + hf_out = hf_out[0] + + # Compare + rtol, atol = 0.05, 0.05 + torch.testing.assert_close(custom_out, hf_out, rtol=rtol, atol=atol) + + +@pytest.mark.parametrize("B,S", _BATCH_AND_SEQUENCE_TEST_CASES) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@torch.no_grad() +def test_glm4_moe_lite_mlp_numerical_equivalence(B, S, dtype): + """Test MLP layer produces numerically equivalent output to HF implementation.""" + HFMLP = _get_hf_mlp_class() + if HFMLP is None: + pytest.skip("transformers doesn't have glm4_moe_lite (requires v5.0+)") + + device = "cuda" + config = _create_small_config() + hf_config = _create_hf_config() + + # Create HF MLP + hf_mlp = HFMLP(hf_config) + hf_mlp.to(device=device, dtype=dtype) + hf_mlp.eval() + + # Create custom MLP and load same weights + custom_mlp = Glm4MoeLiteMLP(config) + custom_mlp.to(device=device, dtype=dtype) + custom_mlp.load_state_dict(hf_mlp.state_dict()) + custom_mlp.eval() + + # Create input + H = config.hidden_size + x = torch.randn(B, S, H, device=device, dtype=dtype) + + # Run both + hf_out = hf_mlp(x) + custom_out = custom_mlp(x) + + # Compare + rtol, atol = 1e-3, 1e-3 + torch.testing.assert_close(custom_out, hf_out, rtol=rtol, atol=atol) + + +@pytest.mark.parametrize("B,S", _BATCH_AND_SEQUENCE_TEST_CASES) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@torch.no_grad() +def test_glm4_moe_lite_full_model_numerical_equivalence(B, S, dtype): + """Test full model produces numerically equivalent output to HF implementation.""" + HFModel = _get_hf_model_class() + if HFModel is None: + pytest.skip("transformers doesn't have glm4_moe_lite (requires v5.0+)") + + device = "cuda" + config = _create_small_config() + hf_config = _create_hf_config() + + # Create HF model + hf_model = HFModel(hf_config) + hf_model.to(device=device, dtype=dtype) + hf_model.eval() + + # Initialize all gate weights for reproducibility + for module in hf_model.modules(): + if hasattr(module, "gate") and hasattr(module.gate, "weight"): + module.gate.weight = torch.nn.Parameter(torch.randn_like(module.gate.weight)) + + # Create custom model and load converted weights + custom_model = Glm4MoeLiteForCausalLM(config) + custom_model.to(device=device, dtype=dtype) + + # Convert HF stacked expert weights to our per-expert format + hf_state_dict = hf_model.state_dict() + custom_state_dict = _convert_hf_full_model_state_dict_to_custom(hf_state_dict, config) + custom_model.load_state_dict(custom_state_dict) + custom_model.eval() + + # Create input + input_ids = torch.randint(0, config.vocab_size, (B, S), device=device) + position_ids = torch.arange(S, device=device).unsqueeze(0).expand(B, -1) + + # Run both + hf_out = hf_model(input_ids=input_ids, position_ids=position_ids) + custom_out = custom_model(input_ids=input_ids, position_ids=position_ids) + + # Compare logits - cast to same dtype for comparison + # (HF model may output float32 for numerical stability in lm_head) + rtol, atol = 0.05, 0.05 + torch.testing.assert_close( + custom_out.logits.float(), + hf_out.logits.float(), + rtol=rtol, + atol=atol, + ) diff --git a/tests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_sdpa_mla.py b/tests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_sdpa_mla.py deleted file mode 100644 index ffa2594d908f..000000000000 --- a/tests/unittest/_torch/auto_deploy/unit/singlegpu/models/test_sdpa_mla.py +++ /dev/null @@ -1,97 +0,0 @@ -import torch -import torch.nn as nn -import torch.nn.functional as F - -import tensorrt_llm._torch.auto_deploy # noqa: F401 - - -def naive_attn(q_nope, q_pe, compressed_kv, k_pe, wkv_b, softmax_scale): - bsz = q_nope.shape[0] - q_len = q_nope.shape[2] - num_heads = q_nope.shape[1] - qk_nope_head_dim = q_nope.shape[-1] - v_head_dim = wkv_b.weight.shape[-1] - qk_nope_head_dim - qk_head_dim = qk_nope_head_dim + q_pe.shape[-1] - k_pe = k_pe.view(bsz, q_len, 1, q_pe.shape[-1]).transpose(1, 2) - - # Up project compressed_kv - kv = ( - wkv_b(compressed_kv) - .view(bsz, q_len, num_heads, qk_nope_head_dim + v_head_dim) - .transpose(1, 2) - ) - - k_nope, value_states = torch.split(kv, [qk_nope_head_dim, v_head_dim], dim=-1) - - query_states = k_pe.new_empty(bsz, num_heads, q_len, qk_head_dim) - query_states[:, :, :, :qk_nope_head_dim] = q_nope - query_states[:, :, :, qk_nope_head_dim:] = q_pe - - key_states = k_pe.new_empty(bsz, num_heads, q_len, qk_head_dim) - key_states[:, :, :, :qk_nope_head_dim] = k_nope - key_states[:, :, :, qk_nope_head_dim:] = k_pe - - x = F.scaled_dot_product_attention( - query_states, - key_states, - value_states, - attn_mask=None, - is_causal=False, - dropout_p=0.0, - scale=softmax_scale, - ).transpose(1, 2) - - return x - - -def mla_attn(q_nope, q_pe, compressed_kv, k_pe, wkv_b, softmax_scale): - num_heads = q_nope.shape[1] - qk_nope_head_dim = q_nope.shape[-1] - v_head_dim = wkv_b.weight.shape[-1] - qk_nope_head_dim - kv_lora_rank = compressed_kv.shape[-1] - - # Down project q_nope - wkv_b_weight = wkv_b.weight.view(num_heads, -1, kv_lora_rank) - q_nope_proj = torch.einsum("bhsd,hdc->bhsc", q_nope, wkv_b_weight[:, :qk_nope_head_dim]) - - # MLA ref operation - x = torch.ops.auto_deploy.torch_attention_deepseek_mla( - q_nope_proj, q_pe, compressed_kv, k_pe, None, softmax_scale - ) - - # Up project attention scores - x = torch.einsum("bshc,hdc->bshd", x, wkv_b_weight[:, -v_head_dim:]) - return x - - -def test_attn(): - # Define test configurations - kv_lora_rank = 4 - bsz = 2 - q_len = 6 - v_head_dim = 2 - qk_nope_head_dim = 2 - qk_rope_head_dim = 1 - qk_head_dim = qk_nope_head_dim + qk_rope_head_dim - num_heads = 4 - softmax_scale = qk_head_dim**-0.5 - - # Generate inputs - q_nope = torch.randn(bsz, num_heads, q_len, qk_nope_head_dim) - q_pe = torch.randn(bsz, num_heads, q_len, qk_rope_head_dim) - compressed_kv = torch.randn(bsz, q_len, kv_lora_rank) - k_pe = torch.randn(bsz, q_len, qk_rope_head_dim) - - # Define w_kv_b projection matrix - wkv_b = nn.Linear( - kv_lora_rank, num_heads * (qk_head_dim - qk_rope_head_dim + v_head_dim), bias=False - ) - - # Compute naive attention - out_naive = naive_attn(q_nope, q_pe, compressed_kv, k_pe, wkv_b, softmax_scale) - - # Compute MLA attention - out_mla = mla_attn(q_nope, q_pe, compressed_kv, k_pe, wkv_b, softmax_scale) - - # Check if the two outputs are close - assert torch.allclose(out_naive, out_mla, rtol=1e-5, atol=1e-5) From 663c961713a26d94c2b087554b3ce4281255b5be Mon Sep 17 00:00:00 2001 From: Lucas Liebenwein <11156568+lucaslie@users.noreply.github.com> Date: Thu, 5 Feb 2026 19:25:46 -0800 Subject: [PATCH 02/17] removed custom masking in flashinfer_mla Signed-off-by: Lucas Liebenwein <11156568+lucaslie@users.noreply.github.com> --- .../custom_ops/mla/flashinfer_mla.py | 90 ++----------------- 1 file changed, 5 insertions(+), 85 deletions(-) diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashinfer_mla.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashinfer_mla.py index 3f9b353187bc..0171713aaee4 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashinfer_mla.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashinfer_mla.py @@ -166,47 +166,30 @@ def plan_prefill( qo_indptr_host: torch.Tensor, kv_indptr_host: torch.Tensor, plan_params: MLAPrefillPlanParams, - custom_mask: Optional[torch.Tensor] = None, ) -> flashinfer.prefill.BatchPrefillWithRaggedKVCacheWrapper: """Plan prefill using BatchPrefillWithRaggedKVCacheWrapper. For MLA prefill, we expand compressed_kv to get full K, V tensors - and use standard ragged KV cache attention. + and use standard ragged KV cache attention with causal masking. Args: qo_indptr_host: Cumulative query/output lengths on host. kv_indptr_host: Cumulative key/value lengths on host. plan_params: Parameters for planning (hashable, no tensors). - custom_mask: Optional custom attention mask tensor on CUDA. """ - # Always re-plan when custom_mask is provided since the mask changes per batch - needs_replan = (plan_params != self.plan_params_prefill) or (custom_mask is not None) - - if needs_replan: - # Use causal=True as fallback when no custom_mask is provided - use_causal = custom_mask is None - - # When using custom_mask, FlashInfer expects indptr tensors on CUDA - if custom_mask is not None: - qo_indptr = qo_indptr_host.to(custom_mask.device) - kv_indptr = kv_indptr_host.to(custom_mask.device) - else: - qo_indptr = qo_indptr_host - kv_indptr = kv_indptr_host - + if plan_params != self.plan_params_prefill: self.prefill_wrapper.plan( - qo_indptr, - kv_indptr, + qo_indptr_host, + kv_indptr_host, plan_params.num_heads, plan_params.num_kv_heads, plan_params.head_dim_qk, head_dim_vo=plan_params.head_dim_vo, use_fp16_qk_reduction=False, - causal=use_causal, + causal=True, q_data_type=plan_params.q_dtype, kv_data_type=plan_params.kv_dtype, sm_scale=plan_params.sm_scale, - custom_mask=custom_mask, ) self.plan_params_prefill = plan_params @@ -389,61 +372,6 @@ def plan_generate_only( _GlobalFlashInferMLAPlanner = _FlashInferMLAPlanner() -def _compute_ragged_causal_mask( - cu_seqlen_host: torch.Tensor, - num_seq: int, - device: torch.device, -) -> torch.Tensor: - """Compute a causal mask for FlashInfer ragged prefill (fully GPU-vectorized). - - For ragged attention with multiple sequences, this function creates a - causal mask where each sequence has its own lower triangular (causal) mask. - The masks are flattened and concatenated into a single 1D tensor. - - This implementation is fully vectorized on GPU with no Python loops. - - Args: - cu_seqlen_host: Cumulative sequence lengths on host [num_seq + 1]. - seq_len[i] = cu_seqlen_host[i+1] - cu_seqlen_host[i] - num_seq: Number of sequences in the batch. - device: Device to create the mask on. - - Returns: - A flattened boolean mask tensor where True = attend, False = mask out. - Shape: (sum(seq_len[i] * seq_len[i] for i in range(num_seq)),) - """ - # Move sequence lengths to GPU for vectorized computation - seq_lens = (cu_seqlen_host[1 : num_seq + 1] - cu_seqlen_host[:num_seq]).to(device) - - # Compute squared lengths (mask size per sequence) and total - sq_lens = seq_lens * seq_lens - total_mask_size = sq_lens.sum().item() - - # Cumulative squared lengths (offsets into flattened mask) - cu_sq_lens = torch.zeros(num_seq + 1, device=device, dtype=seq_lens.dtype) - cu_sq_lens[1:] = torch.cumsum(sq_lens, dim=0) - - # Create position indices for the entire mask - positions = torch.arange(total_mask_size, device=device) - - # Find which sequence each position belongs to using searchsorted - # right=True ensures positions at boundaries map to the next sequence - seq_ids = torch.searchsorted(cu_sq_lens[1:], positions, right=True) - - # Compute local position within each sequence's mask - local_pos = positions - cu_sq_lens[seq_ids] - - # Get sequence length for each position - seq_len_per_pos = seq_lens[seq_ids] - - # Compute row and column within each sequence's 2D mask - row = local_pos // seq_len_per_pos - col = local_pos % seq_len_per_pos - - # Causal mask: attend if row >= col (lower triangular) - return row >= col - - @torch.library.custom_op("auto_deploy::flashinfer_mla_prepare_metadata", mutates_args=()) def prepare_flashinfer_mla_metadata( position_ids: torch.Tensor, @@ -723,13 +651,6 @@ def flashinfer_mla_with_cache( # K: [tokens, N, qk_head_dim] k_prefill = torch.cat([k_nope_prefill, kpe_expanded], dim=-1).contiguous() - # Compute the causal mask for ragged prefill - custom_mask = _compute_ragged_causal_mask( - cu_seqlen_host=cu_seqlen_host, - num_seq=num_prefill, - device=q_nope.device, - ) - pp_prefill = MLAPrefillPlanParams( num_heads=num_heads, num_kv_heads=num_heads, # For MLA with expanded KV, same as num_heads @@ -745,7 +666,6 @@ def flashinfer_mla_with_cache( qo_indptr_host=cu_seqlen_host[: num_prefill + 1], kv_indptr_host=cu_seqlen_host[: num_prefill + 1], # Same as qo for self-attention plan_params=pp_prefill, - custom_mask=custom_mask, ) y_prefill = wrapper_prefill.run( From cf5661ea80ce9c741ec79ca95995551c245f036f Mon Sep 17 00:00:00 2001 From: Lucas Liebenwein <11156568+lucaslie@users.noreply.github.com> Date: Fri, 6 Feb 2026 07:40:52 -0800 Subject: [PATCH 03/17] follow new unit test folder structure Signed-off-by: Lucas Liebenwein <11156568+lucaslie@users.noreply.github.com> --- .../unit/singlegpu/custom_ops/{ => mla}/test_flashinfer_mla_op.py | 0 .../unit/singlegpu/custom_ops/{ => mla}/test_torch_mla_op.py | 0 2 files changed, 0 insertions(+), 0 deletions(-) rename tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/{ => mla}/test_flashinfer_mla_op.py (100%) rename tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/{ => mla}/test_torch_mla_op.py (100%) diff --git a/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_flashinfer_mla_op.py b/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/mla/test_flashinfer_mla_op.py similarity index 100% rename from tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_flashinfer_mla_op.py rename to tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/mla/test_flashinfer_mla_op.py diff --git a/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_torch_mla_op.py b/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/mla/test_torch_mla_op.py similarity index 100% rename from tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_torch_mla_op.py rename to tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/mla/test_torch_mla_op.py From 9d236ec8e59eef4496e617adb6d76027dd25687d Mon Sep 17 00:00:00 2001 From: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com> Date: Thu, 5 Feb 2026 23:44:30 -0800 Subject: [PATCH 04/17] lint Signed-off-by: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com> --- .../unit/singlegpu/custom_ops/test_update_kv_cache.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_update_kv_cache.py b/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_update_kv_cache.py index b9a635f609e7..47681b559602 100644 --- a/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_update_kv_cache.py +++ b/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_update_kv_cache.py @@ -1,6 +1,8 @@ import torch -from tensorrt_llm._torch.auto_deploy.custom_ops.attention.torch_backend_attention import _update_kv_cache +from tensorrt_llm._torch.auto_deploy.custom_ops.attention.torch_backend_attention import ( + _update_kv_cache, +) def test_update_kv_cache(): From a5dfc413ad3f7149806273d40c292e46a48e342d Mon Sep 17 00:00:00 2001 From: Balamurugan Marimuthu <246387390+bmarimuthu-nv@users.noreply.github.com> Date: Thu, 5 Feb 2026 16:02:39 -0800 Subject: [PATCH 05/17] added triton_rope_on_interleaved_qk_inputs fused kernel Signed-off-by: Balamurugan Marimuthu <246387390+bmarimuthu-nv@users.noreply.github.com> --- .../_torch/auto_deploy/custom_ops/README.md | 1 + .../custom_ops/rope/triton_rope.py | 106 +++++++++- .../custom_ops/rope/triton_rope_kernel.py | 185 ++++++++++++++++++ .../custom_ops/rope/test_triton_rope.py | 120 +++++++++++- 4 files changed, 410 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/README.md b/tensorrt_llm/_torch/auto_deploy/custom_ops/README.md index 27f18bed1b92..8e625955a731 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/README.md +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/README.md @@ -35,6 +35,7 @@ The table below lists the operators ordered by their backend. | `torch.ops.auto_deploy.triton_attention_flattened_mha_with_cache` | Triton flattened MHA with cache | | `torch.ops.auto_deploy.triton_attention_fused_flattened_mla_with_cache` | Triton fused flattened Multi-head Latent Attention with cache support | | `torch.ops.auto_deploy.triton_rope_on_flattened_inputs` | Triton RoPE on flattened inputs | +| `torch.ops.auto_deploy.triton_rope_on_interleaved_qk_inputs` | Triton fused RoPE on interleaved QK inputs (position lookup + de-interleave + RoPE) | | `torch.ops.auto_deploy.triton_rope_with_input_pos` | Triton RoPE with input positions | | `torch.ops.auto_deploy.trtllm_moe_fused` | TensorRT LLM fused MoE implementation | | `torch.ops.auto_deploy.trtllm_dist_all_gather` | Distributed all-gather operation (TRT-LLM backend, MPI mode) | diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/rope/triton_rope.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/rope/triton_rope.py index 3c5d79c0f0fb..5dff54543533 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/rope/triton_rope.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/rope/triton_rope.py @@ -12,11 +12,16 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +from typing import Tuple import torch import triton -from .triton_rope_kernel import rope_fwd_flattened_kernel, rope_fwd_kernel +from .triton_kernels.rope import ( + rope_fwd_flattened_kernel, + rope_fwd_interleaved_kernel, + rope_fwd_kernel, +) @torch.library.custom_op("auto_deploy::triton_rope_with_input_pos", mutates_args=()) @@ -147,3 +152,102 @@ def apply_rope_on_flattened_inputs( @apply_rope_on_flattened_inputs.register_fake def apply_rope_on_flattened_inputs_fake(x, freqs_cis, input_pos, seq_lens, seq_start_indices): return torch.empty_like(x) + + +@torch.library.custom_op("auto_deploy::triton_rope_on_interleaved_qk_inputs", mutates_args=()) +def apply_rope_on_interleaved_qk_inputs( + q: torch.Tensor, + k: torch.Tensor, + cos_cache: torch.Tensor, + sin_cache: torch.Tensor, + position_ids: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Fused RoPE for DeepSeek-style interleaved Q/K inputs. + + This kernel fuses: + 1. Position ID lookup for cos/sin from cache + 2. De-interleaving of Q/K + 3. RoPE application + + Args: + q: Query tensor with interleaved layout [B, S, H_Q, D] + Input: [q0_r, q0_i, q1_r, q1_i, ...] + k: Key tensor with interleaved layout [B, S, H_K, D] + Typically H_K=1 for MQA + cos_cache: Cosine cache [max_seq_len, D] + sin_cache: Sine cache [max_seq_len, D] + position_ids: Position indices [B, S] + + Returns: + Tuple of (q_rotated, k_rotated) with layout [B, S, H, D] + Output: [y0, y1, ..., y_{D/2-1}, y_{D/2}, ..., y_{D-1}] + where y_first = a*cos - b*sin, y_second = b*cos + a*sin + """ + assert q.dim() == 4, f"Q must be 4D [B, S, H, D], got {q.dim()}D" + assert k.dim() == 4, f"K must be 4D [B, S, H, D], got {k.dim()}D" + assert q.shape[-1] % 2 == 0, "Head dimension must be even" + assert q.shape[-1] == k.shape[-1], "Q and K must have same head dimension" + assert cos_cache.shape == sin_cache.shape, "cos and sin cache must have same shape" + + B, S, H_Q, D = q.shape + _, _, H_K, _ = k.shape + + # Allocate contiguous outputs + # The kernel computes contiguous strides internally for output writes + q_out = torch.empty_like(q) + k_out = torch.empty_like(k) + + # Block sizes + BLOCK_SIZE_H = 32 + BLOCK_SIZE_S = min(triton.next_power_of_2(S), 32) + + # Grid: (B, cdiv(H_Q, BLOCK_SIZE_H), cdiv(S, BLOCK_SIZE_S)) + # Note: We use H_Q for grid since H_Q >= H_K typically + grid = ( + B, + triton.cdiv(H_Q, BLOCK_SIZE_H), + triton.cdiv(S, BLOCK_SIZE_S), + ) + + rope_fwd_interleaved_kernel[grid]( + q, + k, + cos_cache, + sin_cache, + position_ids, + q_out, + k_out, + B, + S, + H_Q, + H_K, + D, + q.stride(0), + q.stride(1), + q.stride(2), + q.stride(3), + k.stride(0), + k.stride(1), + k.stride(2), + k.stride(3), + position_ids.stride(0), + position_ids.stride(1), + cos_cache.stride(0), + cos_cache.stride(1), + BLOCK_SIZE_H=BLOCK_SIZE_H, + BLOCK_SIZE_S=BLOCK_SIZE_S, + ) + + return q_out, k_out + + +@apply_rope_on_interleaved_qk_inputs.register_fake +def apply_rope_on_interleaved_qk_inputs_fake( + q: torch.Tensor, + k: torch.Tensor, + cos_cache: torch.Tensor, + sin_cache: torch.Tensor, + position_ids: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + return torch.empty_like(q), torch.empty_like(k) diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/rope/triton_rope_kernel.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/rope/triton_rope_kernel.py index 4139d3a9d9c7..29526a12e91c 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/rope/triton_rope_kernel.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/rope/triton_rope_kernel.py @@ -17,6 +17,191 @@ import triton.language as tl +@triton.jit +def rope_fwd_interleaved_kernel( + q_ptr, # [B, S, H_Q, D] input Q (interleaved layout) + k_ptr, # [B, S, H_K, D] input K (interleaved layout) + cos_cache_ptr, # [max_seq_len, D] cos cache + sin_cache_ptr, # [max_seq_len, D] sin cache + position_ids_ptr, # [B, S] position IDs + q_out_ptr, # [B, S, H_Q, D] output Q + k_out_ptr, # [B, S, H_K, D] output K + B, # batch size + S, # sequence length + H_Q, # number of Q heads + H_K, # number of K heads (typically 1 for MQA) + D: tl.constexpr, # head dimension + stride_qb, # Q batch stride + stride_qs, # Q seq stride + stride_qh, # Q head stride + stride_qd, # Q dim stride + stride_kb, # K batch stride + stride_ks, # K seq stride + stride_kh, # K head stride + stride_kd, # K dim stride + stride_pos_b, # position_ids batch stride + stride_pos_s, # position_ids seq stride + stride_cache_s, # cache seq stride + stride_cache_d, # cache dim stride + BLOCK_SIZE_H: tl.constexpr, + BLOCK_SIZE_S: tl.constexpr, +): + """ + Fused RoPE kernel for DeepSeek-style interleaved Q/K inputs. + + Input layout (interleaved): [q0_r, q0_i, q1_r, q1_i, ...] + After de-interleave: [q0_r, q1_r, ..., q0_i, q1_i, ...] + + This kernel does: + 1. Position ID lookup for cos/sin + 2. Reads even and odd indices directly from the interleaved input using strided loads + 3. RoPE application directly on the separated a/b values: y_first = a*cos - b*sin, y_second = b*cos + a*sin + 4. Writes the output in contiguous "half-split" layout — first half and second half stored separately + + Grid: (B, cdiv(H_Q, BLOCK_SIZE_H), cdiv(S, BLOCK_SIZE_S)) + """ + D2: tl.constexpr = D // 2 + D2_PADDED: tl.constexpr = triton.next_power_of_2(D2) + + # Program IDs + batch_id = tl.program_id(0) + head_block_id = tl.program_id(1) + seq_block_id = tl.program_id(2) + + # Head offsets and mask + head_offsets = head_block_id * BLOCK_SIZE_H + tl.arange(0, BLOCK_SIZE_H) + head_mask = head_offsets < H_Q + + # Sequence offsets and mask + seq_offsets = seq_block_id * BLOCK_SIZE_S + tl.arange(0, BLOCK_SIZE_S) + seq_mask = seq_offsets < S + + # Dimension offsets for half the head dim (we read pairs) + dim_offsets = tl.arange(0, D2_PADDED) + dim_mask = dim_offsets < D2 + + # ========== LOAD POSITION IDS ========== + # position_ids: [B, S] + pos_ptr = position_ids_ptr + batch_id * stride_pos_b + seq_offsets * stride_pos_s + pos_ids = tl.load(pos_ptr, mask=seq_mask, other=0) # [BLOCK_SIZE_S] + + # ========== LOAD COS/SIN FROM CACHE ========== + # cos_cache, sin_cache: [max_seq_len, D] + # For each position, load the corresponding cos/sin values + # We need cos/sin for both halves of head_dim + cache_offsets = ( + pos_ids[:, None] * stride_cache_s + dim_offsets[None, :] * stride_cache_d + ) # [BLOCK_SIZE_S, D2_PADDED] + + cache_mask = seq_mask[:, None] & dim_mask[None, :] # [BLOCK_SIZE_S, D2_PADDED] + + cos_first = tl.load(cos_cache_ptr + cache_offsets, mask=cache_mask) # [BLOCK_SIZE_S, D2_PADDED] + sin_first = tl.load(sin_cache_ptr + cache_offsets, mask=cache_mask) # [BLOCK_SIZE_S, D2_PADDED] + + # Second half of cos/sin (offset by D2) + cache_offsets_second = ( + pos_ids[:, None] * stride_cache_s + (dim_offsets[None, :] + D2) * stride_cache_d + ) + cos_second = tl.load(cos_cache_ptr + cache_offsets_second, mask=cache_mask) + sin_second = tl.load(sin_cache_ptr + cache_offsets_second, mask=cache_mask) + + # ========== PROCESS Q ========== + # Q layout: [B, S, H, D] with interleaved D + # Read even indices (a) and odd indices (b) for de-interleaving + # Input: [q0_r, q0_i, q1_r, q1_i, ...] -> a=[q0_r, q1_r, ...], b=[q0_i, q1_i, ...] + + q_base = batch_id * stride_qb + + # Compute offsets for reading interleaved data + # even_offsets: positions 0, 2, 4, ... (stride 2) + # odd_offsets: positions 1, 3, 5, ... (stride 2) + q_offsets_base = ( + seq_offsets[:, None, None] * stride_qs + + head_offsets[None, :, None] * stride_qh + + dim_offsets[None, None, :] * 2 * stride_qd # stride 2 for interleaved + ) # [BLOCK_SIZE_S, BLOCK_SIZE_H, D2_PADDED] + + even_offsets = q_base + q_offsets_base + odd_offsets = q_base + q_offsets_base + stride_qd + + # Combined mask + load_mask = seq_mask[:, None, None] & head_mask[None, :, None] & dim_mask[None, None, :] + + # Load Q values (even = a, odd = b) + q_a = tl.load(q_ptr + even_offsets, mask=load_mask) # [BLOCK_SIZE_S, BLOCK_SIZE_H, D2_PADDED] + q_b = tl.load(q_ptr + odd_offsets, mask=load_mask) # [BLOCK_SIZE_S, BLOCK_SIZE_H, D2_PADDED] + + # Broadcast cos/sin for heads: [BLOCK_SIZE_S, D2_PADDED] -> [BLOCK_SIZE_S, 1, D2_PADDED] + cos_first_bc = cos_first[:, None, :] + sin_first_bc = sin_first[:, None, :] + cos_second_bc = cos_second[:, None, :] + sin_second_bc = sin_second[:, None, :] + + # Apply RoPE formula + # y_first_half = a * cos - b * sin + # y_second_half = b * cos + a * sin + q_y1 = q_a * cos_first_bc - q_b * sin_first_bc + q_y2 = q_b * cos_second_bc + q_a * sin_second_bc + + # Store Q output (CONTIGUOUS layout) + # Output layout: [B, S, H_Q, D] with first half = y1, second half = y2 + # Compute contiguous strides: stride_b=S*H_Q*D, stride_s=H_Q*D, stride_h=D, stride_d=1 + q_out_stride_b = S * H_Q * D + q_out_stride_s = H_Q * D + q_out_stride_h = D + q_out_offsets_first = ( + batch_id * q_out_stride_b + + seq_offsets[:, None, None] * q_out_stride_s + + head_offsets[None, :, None] * q_out_stride_h + + dim_offsets[None, None, :] # stride_d = 1 + ) + q_out_offsets_second = q_out_offsets_first + D2 # D2 * 1 + + tl.store(q_out_ptr + q_out_offsets_first, q_y1, mask=load_mask) + tl.store(q_out_ptr + q_out_offsets_second, q_y2, mask=load_mask) + + # ========== PROCESS K ========== + # K typically has H_K=1 for MQA, but we handle general case + # Use head_offsets < H_K for K's mask + head_mask_k = head_offsets < H_K + + k_base = batch_id * stride_kb + + k_offsets_base = ( + seq_offsets[:, None, None] * stride_ks + + head_offsets[None, :, None] * stride_kh + + dim_offsets[None, None, :] * 2 * stride_kd + ) + + k_even_offsets = k_base + k_offsets_base + k_odd_offsets = k_base + k_offsets_base + stride_kd + + load_mask_k = seq_mask[:, None, None] & head_mask_k[None, :, None] & dim_mask[None, None, :] + + k_a = tl.load(k_ptr + k_even_offsets, mask=load_mask_k) + k_b = tl.load(k_ptr + k_odd_offsets, mask=load_mask_k) + + k_y1 = k_a * cos_first_bc - k_b * sin_first_bc + k_y2 = k_b * cos_second_bc + k_a * sin_second_bc + + # Store K output (CONTIGUOUS layout) + # Output layout: [B, S, H_K, D] with first half = y1, second half = y2 + # Compute contiguous strides: stride_b=S*H_K*D, stride_s=H_K*D, stride_h=D, stride_d=1 + k_out_stride_b = S * H_K * D + k_out_stride_s = H_K * D + k_out_stride_h = D + k_out_offsets_first = ( + batch_id * k_out_stride_b + + seq_offsets[:, None, None] * k_out_stride_s + + head_offsets[None, :, None] * k_out_stride_h + + dim_offsets[None, None, :] # stride_d = 1 + ) + k_out_offsets_second = k_out_offsets_first + D2 # D2 * 1 + + tl.store(k_out_ptr + k_out_offsets_first, k_y1, mask=load_mask_k) + tl.store(k_out_ptr + k_out_offsets_second, k_y2, mask=load_mask_k) + + @triton.jit def rope_fwd_kernel( x_ptr, diff --git a/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/rope/test_triton_rope.py b/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/rope/test_triton_rope.py index f7a8a5972e36..fccda7d4286c 100644 --- a/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/rope/test_triton_rope.py +++ b/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/rope/test_triton_rope.py @@ -1,9 +1,25 @@ -from typing import Optional +from typing import Optional, Tuple import pytest import torch from _custom_op_utils import torch_rope_reference +# Import after we've imported torch (to ensure custom ops are registered) +from tensorrt_llm._torch.auto_deploy.custom_ops import triton_rope # noqa: F401 + + +def _precompute_cos_sin_cache( + max_seq_len: int, head_dim: int, rope_theta: float = 10000.0, dtype: torch.dtype = torch.float16 +) -> Tuple[torch.Tensor, torch.Tensor]: + """Precompute cos and sin cache for RoPE (DeepSeek-style).""" + inv_freq = 1.0 / (rope_theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)) + positions = torch.arange(max_seq_len, dtype=torch.float32) + freqs = positions.unsqueeze(1) * inv_freq.unsqueeze(0) # [max_seq_len, head_dim//2] + emb = torch.cat((freqs, freqs), dim=-1) # [max_seq_len, head_dim] + cos_cache = emb.cos().to(dtype) + sin_cache = emb.sin().to(dtype) + return cos_cache, sin_cache + def _precompute_freqs_cis( seq_len: int, head_dim: int, rope_theta: Optional[float] = None @@ -70,3 +86,105 @@ def test_rope_flattened(d_head): y_reshaped = y.unflatten(-1, (2, N_ELEM // 2)).transpose(-2, -1).flatten(-2).contiguous() assert torch.allclose(y_ref.cpu(), y_reshaped.cpu(), atol=1e-2, rtol=1e-2) + + +@pytest.mark.parametrize( + "batch_size,seq_len,num_q_heads,num_k_heads,head_dim", + [ + (1, 4, 8, 8, 64), # Standard case + (2, 16, 20, 1, 128), # GLM-4 style: non-power-of-2 heads, MQA + (1, 1, 8, 8, 64), # Single token (decode) + (4, 32, 16, 2, 96), # GQA with non-standard head_dim + ], +) +def test_triton_rope_on_interleaved_qk_inputs( + batch_size: int, seq_len: int, num_q_heads: int, num_k_heads: int, head_dim: int +): + """ + Test that triton_rope_on_interleaved_qk_inputs produces the same output as + the PyTorch reference (index + torch_rope_with_qk_interleaving). + """ + device = "cuda" + dtype = torch.bfloat16 + max_seq_len = 1024 + + # Create random inputs with interleaved layout [B, S, H, D] + q = torch.randn(batch_size, seq_len, num_q_heads, head_dim, device=device, dtype=dtype) + k = torch.randn(batch_size, seq_len, num_k_heads, head_dim, device=device, dtype=dtype) + + # Precompute cos/sin cache + cos_cache, sin_cache = _precompute_cos_sin_cache(max_seq_len, head_dim, dtype=dtype) + cos_cache = cos_cache.to(device) + sin_cache = sin_cache.to(device) + + # Random position IDs (not necessarily sequential) + position_ids = torch.randint(0, max_seq_len - seq_len, (batch_size,), device=device) + position_ids = position_ids.unsqueeze(1) + torch.arange(seq_len, device=device).unsqueeze(0) + # position_ids: [B, S] + + # ========== PyTorch Reference ========== + # Step 1: Index cos/sin with position_ids + cos_indexed = cos_cache[position_ids] # [B, S, D] + sin_indexed = sin_cache[position_ids] # [B, S, D] + + # Step 2: Apply PyTorch rope with qk interleaving + # unsqueeze_dim=2 for [B, S, H, D] layout + q_ref, k_ref = torch.ops.auto_deploy.torch_rope_with_qk_interleaving( + q, k, cos_indexed, sin_indexed, unsqueeze_dim=2 + ) + + # ========== Triton Implementation ========== + q_triton, k_triton = torch.ops.auto_deploy.triton_rope_on_interleaved_qk_inputs( + q, k, cos_cache, sin_cache, position_ids + ) + + # ========== Compare Outputs ========== + # Use relative tolerance for bf16 + atol = 1e-2 + rtol = 1e-2 + + assert torch.allclose(q_ref, q_triton, atol=atol, rtol=rtol), ( + f"Q mismatch: max diff = {(q_ref - q_triton).abs().max().item()}" + ) + assert torch.allclose(k_ref, k_triton, atol=atol, rtol=rtol), ( + f"K mismatch: max diff = {(k_ref - k_triton).abs().max().item()}" + ) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_triton_rope_interleaved_dtype_consistency(dtype): + """Test that the Triton kernel works correctly with different dtypes.""" + device = "cuda" + batch_size, seq_len, num_heads, head_dim = 2, 8, 8, 64 + max_seq_len = 1024 + + q = torch.randn(batch_size, seq_len, num_heads, head_dim, device=device, dtype=dtype) + k = torch.randn(batch_size, seq_len, num_heads, head_dim, device=device, dtype=dtype) + + cos_cache, sin_cache = _precompute_cos_sin_cache(max_seq_len, head_dim, dtype=dtype) + cos_cache = cos_cache.to(device) + sin_cache = sin_cache.to(device) + + position_ids = torch.arange(seq_len, device=device).unsqueeze(0).expand(batch_size, -1) + + # PyTorch reference + cos_indexed = cos_cache[position_ids] + sin_indexed = sin_cache[position_ids] + q_ref, k_ref = torch.ops.auto_deploy.torch_rope_with_qk_interleaving( + q, k, cos_indexed, sin_indexed, unsqueeze_dim=2 + ) + + # Triton + q_triton, k_triton = torch.ops.auto_deploy.triton_rope_on_interleaved_qk_inputs( + q, k, cos_cache, sin_cache, position_ids + ) + + # Verify outputs match + atol = 1e-2 + rtol = 1e-2 + assert torch.allclose(q_ref, q_triton, atol=atol, rtol=rtol) + assert torch.allclose(k_ref, k_triton, atol=atol, rtol=rtol) + + # Verify output dtype is preserved + assert q_triton.dtype == dtype + assert k_triton.dtype == dtype From 097043a9509d7da85fa7acc33f499e43e3202044 Mon Sep 17 00:00:00 2001 From: Balamurugan Marimuthu <246387390+bmarimuthu-nv@users.noreply.github.com> Date: Thu, 5 Feb 2026 18:04:01 -0800 Subject: [PATCH 06/17] added transform to replace torch op with triton kernel Signed-off-by: Balamurugan Marimuthu <246387390+bmarimuthu-nv@users.noreply.github.com> --- .../auto_deploy/transform/library/rope.py | 120 +++++++++++++++++- .../library/test_rope_transformation.py | 61 +++++++++ 2 files changed, 180 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/rope.py b/tensorrt_llm/_torch/auto_deploy/transform/library/rope.py index 811dd34af52f..4e9318f0c101 100644 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/rope.py +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/rope.py @@ -363,7 +363,11 @@ def _apply( class OptimizeRope(BaseTransform): """ Scan the FX graph and replace calls to the torch-reference RoPE ops with - the optimized `rope::flashinfer` kernel. + optimized kernels: + - ``torch_rope_with_explicit_cos_sin`` → ``flashinfer_rope`` + - ``torch_rope_with_complex_freqs`` → ``flashinfer_rope`` + - ``torch_rope_with_qk_interleaving`` → ``triton_rope_on_interleaved_qk_inputs`` + Precomputes positional IDs and the fused cosine-sine cache as explicit nodes, and reuses those nodes when possible. """ @@ -385,6 +389,9 @@ def _apply( _optimize_explicit(graph, node, rope_flash_cache, rope_position_ids_cache) elif is_op(node, torch.ops.auto_deploy.torch_rope_with_complex_freqs): _optimize_complex(graph, node, rope_flash_cache, rope_position_ids_cache) + elif is_op(node, torch.ops.auto_deploy.torch_rope_with_qk_interleaving): + if not _optimize_interleaved(graph, node): + continue else: continue num_rope_optimizations += 1 @@ -546,6 +553,117 @@ def _optimize_complex( graph.erase_node(node) +def _trace_back_index(node: Node) -> Optional[Tuple[Node, Node]]: + """Trace back from a node to find an aten.index.Tensor producer. + + If ``node`` was produced by ``aten.index.Tensor(cache, [position_ids])``, + return ``(cache_node, position_ids_node)``. Otherwise return ``None``. + """ + if not is_op(node, torch.ops.aten.index.Tensor): + return None + cache_node = node.args[0] + indices = node.args[1] + if not isinstance(indices, (list, tuple)) or len(indices) != 1: + return None + position_ids_node = indices[0] + if not isinstance(cache_node, Node) or not isinstance(position_ids_node, Node): + return None + return cache_node, position_ids_node + + +def _validate_interleaved_rope_inputs(q_node: Node, k_node: Node) -> bool: + """Validate q/k inputs for the interleaved triton RoPE kernel. + + Relaxed compared to ``_validate_rope_inputs``: requires even head_dim + (not head_dim % 64 == 0) since the triton kernel pads to next-power-of-2. + """ + for node in (q_node, k_node): + fake_val = node.meta.get("val", None) + if fake_val is None: + return False + + # dtype must be half-precision + if fake_val.dtype not in (torch.float16, torch.bfloat16): + return False + + # Must be at least 4-D + if len(fake_val.shape) < 4: + return False + + # head_dim must be even + head_dim = fake_val.shape[-1] + if isinstance(head_dim, int) and head_dim % 2 != 0: + return False + + # BSND layout: dim 1 (S) should be symbolic, dim 2 (N) should be static + if not isinstance(fake_val.shape[1], torch.SymInt): + return False + if not isinstance(fake_val.shape[2], int): + return False + + return True + + +def _optimize_interleaved(graph: torch.fx.Graph, node: Node) -> bool: + """Replace ``torch_rope_with_qk_interleaving`` with ``triton_rope_on_interleaved_qk_inputs``. + + Traces back from cos/sin nodes to find the original ``aten.index.Tensor`` + producer and extracts the cached cos/sin tensors and position_ids, passing + them directly to the triton kernel (which fuses the position lookup). + + Returns True if the replacement was made, False if skipped. + """ + # --- extract arguments --------------------------------------------------- + q_node, k_node, cos_node, sin_node, *rest = node.args + q_rope_old, k_rope_old = extract_output_tuple(node, 2) + if q_rope_old is None or k_rope_old is None: + return False + + # --- validate inputs ----------------------------------------------------- + if not _validate_interleaved_rope_inputs(q_node, k_node): + return False + + # --- trace back cos/sin to find cache + position_ids --------------------- + cos_traced = _trace_back_index(cos_node) + if cos_traced is None: + return False + cos_cache_node, cos_position_ids_node = cos_traced + + sin_traced = _trace_back_index(sin_node) + if sin_traced is None: + return False + sin_cache_node, sin_position_ids_node = sin_traced + + # Both cos and sin must use the same position_ids + if cos_position_ids_node is not sin_position_ids_node: + return False + + position_ids_node = cos_position_ids_node + + # --- insert the triton op ------------------------------------------------ + with graph.inserting_before(node): + triton_node = graph.call_function( + torch.ops.auto_deploy.triton_rope_on_interleaved_qk_inputs, + args=(q_node, k_node, cos_cache_node, sin_cache_node, position_ids_node), + ) + + with graph.inserting_after(triton_node): + q_rope_new = graph.call_function(operator.getitem, args=(triton_node, 0)) + k_rope_new = graph.call_function(operator.getitem, args=(triton_node, 1)) + + # --- rewire outputs ------------------------------------------------------ + q_rope_new.meta["val"] = q_rope_old.meta.get("val", None) + k_rope_new.meta["val"] = k_rope_old.meta.get("val", None) + + q_rope_old.replace_all_uses_with(q_rope_new) + k_rope_old.replace_all_uses_with(k_rope_new) + + graph.erase_node(q_rope_old) + graph.erase_node(k_rope_old) + + return True + + def _match_input_interleave_pattern(node: Node) -> Optional[Dict[str, Node]]: """ Detect DeepSeek-style interleave on Q/K: diff --git a/tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_rope_transformation.py b/tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_rope_transformation.py index 291cd377bd99..e636803fe30a 100644 --- a/tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_rope_transformation.py +++ b/tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_rope_transformation.py @@ -490,3 +490,64 @@ def checker(gm): None, # check_num_matches False, # skip_output_assert ) + + +@pytest.mark.parametrize( + "num_heads,num_kv_heads", + [ + (8, 8), # Standard MHA + (8, 1), # MQA (DeepSeek-style) + ], +) +@torch.inference_mode() +def test_optimize_interleaved_rope(num_heads, num_kv_heads): + """Test that optimize_rope replaces torch_rope_with_qk_interleaving + with triton_rope_on_interleaved_qk_inputs by tracing back through + aten.index.Tensor to find the cached cos/sin and position_ids.""" + batch, seq, hid = 4, 12, 512 + model = DSModel(hid, 16, num_heads, num_kv_heads, layout="BSND", mode="optimize") + model = model.to("cuda", torch.float16) + + x = torch.randn(batch, seq, hid, device="cuda", dtype=torch.float16) + dynamic_shapes = model.get_dynamic_shapes() + gm = torch_export_to_gm(model, args=(x,), dynamic_shapes=(dynamic_shapes,), clone=True) + + # Verify the graph contains torch_rope_with_qk_interleaving before optimization + assert any( + is_op(n, torch.ops.auto_deploy.torch_rope_with_qk_interleaving) for n in gm.graph.nodes + ), "Expected torch_rope_with_qk_interleaving in graph before optimization" + + gm_transformed = InferenceOptimizer( + None, + { + "optimize_rope": { + "stage": "pattern_matcher", + }, + }, + )(None, gm) + gm_transformed.to("cuda") + + def checker(gm): + has_triton = any( + is_op(n, torch.ops.auto_deploy.triton_rope_on_interleaved_qk_inputs) + for n in gm.graph.nodes + ) + no_torch_rope = not any( + is_op(n, torch.ops.auto_deploy.torch_rope_with_qk_interleaving) for n in gm.graph.nodes + ) + return has_triton and no_torch_rope + + run_test_transformed_gm( + model, + x, + gm_transformed, + checker, + lambda num_p: num_p, + 1e-2, # atol + 1e-2, # rtol + True, # test_load_hook + True, # strict_loading + dynamic_shapes, # dynamic_shapes + None, # check_num_matches + False, # skip_output_assert + ) From 379916156fdf91864a0760d148289b2896450e74 Mon Sep 17 00:00:00 2001 From: Balamurugan Marimuthu <246387390+bmarimuthu-nv@users.noreply.github.com> Date: Thu, 5 Feb 2026 18:13:56 -0800 Subject: [PATCH 07/17] added more input validation Signed-off-by: Balamurugan Marimuthu <246387390+bmarimuthu-nv@users.noreply.github.com> --- .../_torch/auto_deploy/custom_ops/rope/triton_rope.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/rope/triton_rope.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/rope/triton_rope.py index 5dff54543533..30f5bf08d066 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/rope/triton_rope.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/rope/triton_rope.py @@ -192,6 +192,9 @@ def apply_rope_on_interleaved_qk_inputs( B, S, H_Q, D = q.shape _, _, H_K, _ = k.shape + assert k.shape[0] == B and k.shape[1] == S, "Q and K must have same batch and seq dims" + assert position_ids.shape == (B, S), f"position_ids must be [B, S], got {position_ids.shape}" + assert H_Q >= H_K, f"H_Q ({H_Q}) must be >= H_K ({H_K}) for grid sizing" # Allocate contiguous outputs # The kernel computes contiguous strides internally for output writes @@ -203,7 +206,7 @@ def apply_rope_on_interleaved_qk_inputs( BLOCK_SIZE_S = min(triton.next_power_of_2(S), 32) # Grid: (B, cdiv(H_Q, BLOCK_SIZE_H), cdiv(S, BLOCK_SIZE_S)) - # Note: We use H_Q for grid since H_Q >= H_K typically + # H_Q >= H_K is enforced above; K heads are masked within each block grid = ( B, triton.cdiv(H_Q, BLOCK_SIZE_H), From 331924218e62139553498cf984204092e8c3335f Mon Sep 17 00:00:00 2001 From: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com> Date: Thu, 5 Feb 2026 21:12:58 -0800 Subject: [PATCH 08/17] small cleanup in plan_decode to avoid redundant computations Signed-off-by: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com> Signed-off-by: Balamurugan Marimuthu <246387390+bmarimuthu-nv@users.noreply.github.com> --- .../custom_ops/mla/flashinfer_mla.py | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashinfer_mla.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashinfer_mla.py index 0171713aaee4..961135ff9c6e 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashinfer_mla.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashinfer_mla.py @@ -243,12 +243,6 @@ def plan_decode( plan_params: Parameters for planning. """ # Decode qo_indptr: [0, 1, 2, ..., batch_size] (1 token per sequence) - batch_size = kv_page_indptr.shape[0] - 1 - qo_indptr = torch.arange(batch_size + 1, device=kv_page_indptr.device, dtype=torch.int32) - - # Compute kv_len_arr for CUDA graph wrapper initialization - num_pages_per_seq = kv_page_indptr[1:] - kv_page_indptr[:-1] - kv_len_arr = (num_pages_per_seq - 1) * plan_params.page_size + kv_last_page_len # we want to plan during warm-up of cuda graph capture to ensure we have the plan cached if ( @@ -257,6 +251,14 @@ def plan_decode( ): # During CUDA graph capture, the metadata tensors provided by auto-deploy are stable. # Pass the buffer tensors to the wrapper for use_cuda_graph=True + batch_size = kv_page_indptr.shape[0] - 1 + qo_indptr = torch.arange( + batch_size + 1, device=kv_page_indptr.device, dtype=torch.int32 + ) + + # Compute kv_len_arr for CUDA graph wrapper initialization + num_pages_per_seq = kv_page_indptr[1:] - kv_page_indptr[:-1] + kv_len_arr = (num_pages_per_seq - 1) * plan_params.page_size + kv_last_page_len wrapper = self._init_decode_wrapper( use_cuda_graph=True, qo_indptr=qo_indptr, @@ -276,6 +278,11 @@ def plan_decode( # Re-plan if plan_params changed if plan_params != self.plan_params_decode: + batch_size = kv_page_indptr.shape[0] - 1 + qo_indptr = torch.arange( + batch_size + 1, device=kv_page_indptr.device, dtype=torch.int32 + ) + self._plan_mla_wrapper( self.decode_wrapper, qo_indptr, From 174b8ae0373c1588de2433c7464692f2d537c95d Mon Sep 17 00:00:00 2001 From: Balamurugan Marimuthu <246387390+bmarimuthu-nv@users.noreply.github.com> Date: Fri, 6 Feb 2026 13:10:48 -0800 Subject: [PATCH 09/17] fixed conflicts Signed-off-by: Balamurugan Marimuthu <246387390+bmarimuthu-nv@users.noreply.github.com> --- tensorrt_llm/_torch/auto_deploy/custom_ops/rope/triton_rope.py | 2 +- tests/integration/defs/accuracy/test_llm_api_autodeploy.py | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/rope/triton_rope.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/rope/triton_rope.py index 30f5bf08d066..59dd7f899ecd 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/rope/triton_rope.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/rope/triton_rope.py @@ -17,7 +17,7 @@ import torch import triton -from .triton_kernels.rope import ( +from .triton_rope_kernel import ( rope_fwd_flattened_kernel, rope_fwd_interleaved_kernel, rope_fwd_kernel, diff --git a/tests/integration/defs/accuracy/test_llm_api_autodeploy.py b/tests/integration/defs/accuracy/test_llm_api_autodeploy.py index ee9934cda3ca..630719919b3a 100644 --- a/tests/integration/defs/accuracy/test_llm_api_autodeploy.py +++ b/tests/integration/defs/accuracy/test_llm_api_autodeploy.py @@ -426,7 +426,8 @@ class TestGLM4Flash(LlmapiAccuracyTestHarness): ``` """ - MODEL_NAME = "zai-org/GLM-4.7-Flash" + # MODEL_NAME = "zai-org/GLM-4.7-Flash" + MODEL_NAME = "DeepInfra/GLM-4.7-Flash-NVFP4" MODEL_PATH = MODEL_NAME # Model is in HF_CACHE # Set minimum possible seq len + small buffer, for test speed & memory usage MAX_SEQ_LEN = max(MMLU.MAX_INPUT_LEN + MMLU.MAX_OUTPUT_LEN, From 5ead4b4cde8a063c49dc27dd127c81062a14d740 Mon Sep 17 00:00:00 2001 From: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com> Date: Fri, 6 Feb 2026 21:15:19 -0800 Subject: [PATCH 10/17] test multi-stream Signed-off-by: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com> --- examples/auto_deploy/glm_flash.yaml | 20 +++++++++ .../custom_ops/mla/flashinfer_mla.py | 41 ++++++++++++------- .../models/custom/modeling_glm4_moe_lite.py | 8 ++-- .../transform/library/multi_stream_moe.py | 5 ++- .../custom_ops/mla/test_flashinfer_mla_op.py | 1 + 5 files changed, 54 insertions(+), 21 deletions(-) create mode 100644 examples/auto_deploy/glm_flash.yaml diff --git a/examples/auto_deploy/glm_flash.yaml b/examples/auto_deploy/glm_flash.yaml new file mode 100644 index 000000000000..e20e7cbf1c86 --- /dev/null +++ b/examples/auto_deploy/glm_flash.yaml @@ -0,0 +1,20 @@ +runtime: trtllm +compile_backend: torch-cudagraph +max_seq_len: 4096 +max_num_tokens: 4096 +max_batch_size: 64 +enable_chunked_prefill: true +model_factory: AutoModelForCausalLM +disable_overlap_scheduler: false +cuda_graph_batch_sizes: [1, 2, 4, 8, 16, 32, 64] +kv_cache_config: + enable_block_reuse: false + free_gpu_memory_fraction: 0.7 + tokens_per_block: 64 +model_kwargs: + torch_dtype: bfloat16 + +transforms: + multi_stream_moe: + stage: compile + enabled: true diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashinfer_mla.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashinfer_mla.py index 961135ff9c6e..a3afd4abc8ae 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashinfer_mla.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/mla/flashinfer_mla.py @@ -226,6 +226,7 @@ def _plan_mla_wrapper( def plan_decode( self, + qo_indptr: torch.Tensor, kv_page_indptr: torch.Tensor, kv_page_indices: torch.Tensor, kv_last_page_len: torch.Tensor, @@ -237,13 +238,13 @@ def plan_decode( FlashInfer's optimized MLA kernels. Each sequence generates 1 token. Args: + qo_indptr: Cumulative query/output lengths [batch_size + 1]. + For decode, this is [0, 1, 2, ..., batch_size] (1 token per sequence). kv_page_indptr: Cumulative page counts [batch_size + 1]. kv_page_indices: Page indices for the KV cache. kv_last_page_len: Length of the last page per sequence. plan_params: Parameters for planning. """ - # Decode qo_indptr: [0, 1, 2, ..., batch_size] (1 token per sequence) - # we want to plan during warm-up of cuda graph capture to ensure we have the plan cached if ( cuda_graph_state.in_warm_up() @@ -251,10 +252,6 @@ def plan_decode( ): # During CUDA graph capture, the metadata tensors provided by auto-deploy are stable. # Pass the buffer tensors to the wrapper for use_cuda_graph=True - batch_size = kv_page_indptr.shape[0] - 1 - qo_indptr = torch.arange( - batch_size + 1, device=kv_page_indptr.device, dtype=torch.int32 - ) # Compute kv_len_arr for CUDA graph wrapper initialization num_pages_per_seq = kv_page_indptr[1:] - kv_page_indptr[:-1] @@ -278,11 +275,6 @@ def plan_decode( # Re-plan if plan_params changed if plan_params != self.plan_params_decode: - batch_size = kv_page_indptr.shape[0] - 1 - qo_indptr = torch.arange( - batch_size + 1, device=kv_page_indptr.device, dtype=torch.int32 - ) - self._plan_mla_wrapper( self.decode_wrapper, qo_indptr, @@ -332,6 +324,7 @@ def plan_chunked_prefill( def plan_generate_only( self, num_seq: int, + cu_seqlen: torch.Tensor, cu_num_pages: torch.Tensor, cache_loc: torch.Tensor, last_page_len: torch.Tensor, @@ -344,6 +337,8 @@ def plan_generate_only( Args: num_seq: Number of sequences in the decode batch. + cu_seqlen: Cumulative sequence lengths [num_seq + 1]. For decode-only batches, + this is [0, 1, 2, ..., num_seq] (1 token per sequence), serving as qo_indptr. cu_num_pages: Cumulative page counts, already sliced to [: num_seq + 1]. cache_loc: Page indices for the KV cache. last_page_len: Length of the last page per sequence, already sliced to [:num_seq]. @@ -352,16 +347,14 @@ def plan_generate_only( if plan_params.num_seq == num_seq: wrapper = self.cached_cuda_graph_decode_wrappers[plan_params] - # For a pure decode batch, qo_indptr is just [0, 1, 2, ..., batch_size] - qo_indptr = torch.arange(num_seq + 1, device=cu_num_pages.device, dtype=torch.int32) - # Compute actual KV lengths from paging metadata: # kv_len = (num_pages - 1) * page_size + last_page_len num_pages_per_seq = cu_num_pages[1:] - cu_num_pages[:-1] kv_len_arr = (num_pages_per_seq - 1) * plan_params.page_size + last_page_len + # For decode-only batches, cu_seqlen = [0, 1, 2, ..., num_seq] = qo_indptr wrapper.plan( - qo_indptr, + cu_seqlen, # qo_indptr cu_num_pages, # kv_page_indptr cache_loc, # kv_page_indices kv_len_arr, @@ -424,6 +417,7 @@ def prepare_flashinfer_mla_metadata_fake( def prepare_flashinfer_mla_metadata_host( batch_info_host: torch.Tensor, + cu_seqlen_host: torch.Tensor, cu_num_pages_host: torch.Tensor, cache_loc_host: torch.Tensor, last_page_len_host: torch.Tensor, @@ -432,12 +426,23 @@ def prepare_flashinfer_mla_metadata_host( For decode-only batches, this function pre-plans the cached CUDA graph wrappers to avoid planning during graph capture/replay. + + Args: + batch_info_host: Batch info tensor [num_prefill, num_prefill_tokens, num_decode]. + cu_seqlen_host: Cumulative sequence lengths on host. For decode-only batches, + this is [0, 1, 2, ..., num_decode] (1 token per sequence), + which serves as qo_indptr. + cu_num_pages_host: Cumulative page counts on host. + cache_loc_host: Page indices for the KV cache on host. + last_page_len_host: Length of the last page per sequence on host. """ num_prefill, num_prefill_tokens, num_decode = batch_info_host.tolist() if num_prefill == 0: + # For decode-only batches, cu_seqlen_host = [0, 1, 2, ..., num_decode] = qo_indptr _GlobalFlashInferMLAPlanner.plan_generate_only( num_decode, + cu_seqlen_host[: num_decode + 1], cu_num_pages_host[: num_decode + 1], cache_loc_host, last_page_len_host[:num_decode], @@ -729,7 +734,13 @@ def flashinfer_mla_with_cache( sm_scale=scale, ) + # Decode qo_indptr: [0, 1, 2, ..., num_decode] (1 token per sequence) + qo_indptr_decode = torch.arange( + num_decode + 1, device=cu_num_pages.device, dtype=torch.int32 + ) + wrapper_decode = _GlobalFlashInferMLAPlanner.plan_decode( + qo_indptr=qo_indptr_decode, kv_page_indptr=cu_num_pages[num_prefill : num_seq + 1], kv_page_indices=cache_loc, kv_last_page_len=last_page_len[num_prefill:num_seq], diff --git a/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_glm4_moe_lite.py b/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_glm4_moe_lite.py index 17aec049406a..9d672371af7d 100644 --- a/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_glm4_moe_lite.py +++ b/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_glm4_moe_lite.py @@ -425,7 +425,8 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: orig_shape = hidden_states.shape selected_experts, routing_weights = self.gate(hidden_states) - + if self.shared_experts is not None: + shared_output = self.shared_experts(identity) # Use torch_moe custom op for routed experts final_hidden_states = torch.ops.auto_deploy.torch_moe( hidden_states.view(-1, hidden_states.shape[-1]), @@ -439,10 +440,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: ) final_hidden_states = final_hidden_states.view(*orig_shape) - - # Add shared experts output if present - if self.shared_experts is not None: - final_hidden_states = final_hidden_states + self.shared_experts(identity) + final_hidden_states = final_hidden_states + shared_output return final_hidden_states.to(hidden_states.dtype) diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.py b/tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.py index 5cd02553d958..8c217d0d41d7 100644 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.py +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.py @@ -311,6 +311,8 @@ def _apply( torch.ops.auto_deploy.triton_moe_fused: torch.ops.auto_deploy.triton_moe_fused_aux, torch.ops.auto_deploy.trtllm_quant_fp8_moe_fused: torch.ops.auto_deploy.trtllm_quant_fp8_moe_fused_aux, } + with open("before_multi_stream.txt", "w") as f: + f.write(str(gm.graph)) # Ensure that aux stream and events for the current device are added to the CudaStreamManager. cuda_stream_manager.add_device(torch.cuda.current_device()) gm, num_matches = _execute_op_in_aux_stream(gm, op_dict) @@ -321,5 +323,6 @@ def _apply( is_clean=num_matches == 0, has_valid_shapes=num_matches == 0, ) - + with open("after_multi_stream.txt", "w") as f: + f.write(str(gm.graph)) return gm, info diff --git a/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/mla/test_flashinfer_mla_op.py b/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/mla/test_flashinfer_mla_op.py index c9020ffe49cd..23e088fca2c9 100644 --- a/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/mla/test_flashinfer_mla_op.py +++ b/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/mla/test_flashinfer_mla_op.py @@ -1940,6 +1940,7 @@ def test_flashinfer_mla_plan_generate_only( # to re-plan the cached wrappers before graph replay _GlobalFlashInferMLAPlanner.plan_generate_only( batch_size, + flashinfer_meta["cu_seqlen_host"][: batch_size + 1], # cu_seqlen (serves as qo_indptr) flashinfer_meta["cu_num_pages"][: batch_size + 1], flashinfer_meta["cache_loc"], flashinfer_meta["last_page_len"][:batch_size], From 4c0f4721803c94acf030e9315eadb723bba173a6 Mon Sep 17 00:00:00 2001 From: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com> Date: Sat, 7 Feb 2026 10:48:27 -0800 Subject: [PATCH 11/17] multi stream fix, fused siglu, fused add rmsnorm Signed-off-by: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com> --- .../_torch/auto_deploy/config/default.yaml | 5 + .../auto_deploy/custom_ops/linear/__init__.py | 2 + .../auto_deploy/custom_ops/linear/swiglu.py | 128 ++++++++ .../models/custom/modeling_glm4_moe_lite.py | 2 +- .../transform/library/fuse_swiglu.py | 284 +++++++++++++++++ .../transform/library/fused_add_rms_norm.py | 166 +++++++--- .../transform/library/multi_stream_moe.py | 90 ++++-- .../defs/accuracy/test_llm_api_autodeploy.py | 11 +- .../singlegpu/custom_ops/test_multi_stream.py | 53 ++-- .../library/test_fuse_swiglu.py | 188 +++++++++++ .../library/test_fused_add_rms_norm.py | 300 +++++++++++++++--- 11 files changed, 1085 insertions(+), 144 deletions(-) create mode 100644 tensorrt_llm/_torch/auto_deploy/custom_ops/linear/swiglu.py create mode 100644 tensorrt_llm/_torch/auto_deploy/transform/library/fuse_swiglu.py create mode 100644 tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_fuse_swiglu.py diff --git a/tensorrt_llm/_torch/auto_deploy/config/default.yaml b/tensorrt_llm/_torch/auto_deploy/config/default.yaml index 6dd3e617e1d7..f0b188ce55bb 100644 --- a/tensorrt_llm/_torch/auto_deploy/config/default.yaml +++ b/tensorrt_llm/_torch/auto_deploy/config/default.yaml @@ -48,6 +48,8 @@ transforms: match_rope_layout: stage: pattern_matcher expected_layout: bsnd + match_swiglu_pattern: + stage: pattern_matcher match_rmsnorm_pattern: stage: pattern_matcher match_l2norm_pattern: @@ -147,6 +149,9 @@ transforms: fuse_l2norm: stage: post_load_fusion backend: fla + fuse_swiglu: + stage: post_load_fusion + enabled: true fuse_add_rms_norm: stage: post_load_fusion enabled: true diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/linear/__init__.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/linear/__init__.py index b11ccc7ee61b..17a5d87480d7 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/linear/__init__.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/linear/__init__.py @@ -18,9 +18,11 @@ This module provides linear layer implementations: - linear: Linear layer operations - torch_router: MoE router operations +- swiglu: SwiGLU MLP custom operations """ __all__ = [ "linear", "torch_router", + "swiglu", ] diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/linear/swiglu.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/linear/swiglu.py new file mode 100644 index 000000000000..808291d3f588 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/linear/swiglu.py @@ -0,0 +1,128 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""SwiGLU MLP custom operations for graph transformation. + +This module provides custom operators for SwiGLU MLP fusion: +- torch_swiglu_mlp: Intermediate representation after pattern matching +- fused_swiglu_mlp: Fused implementation with concatenated gate+up weights +""" + +from typing import Optional + +import torch +import torch.nn.functional as F + + +@torch.library.custom_op("auto_deploy::torch_swiglu_mlp", mutates_args=()) +def torch_swiglu_mlp( + input: torch.Tensor, + gate_weight: torch.Tensor, + up_weight: torch.Tensor, + down_weight: torch.Tensor, + gate_bias: Optional[torch.Tensor], + up_bias: Optional[torch.Tensor], + down_bias: Optional[torch.Tensor], +) -> torch.Tensor: + """Standardized SwiGLU MLP operation. + + Computes: silu(x @ gate.T + gate_bias) * (x @ up.T + up_bias) @ down.T + down_bias + + This is the intermediate representation used after pattern matching, + before weight fusion is applied. + + Args: + input: Input tensor of shape [..., hidden_size]. + gate_weight: Gate projection weight of shape [intermediate_size, hidden_size]. + up_weight: Up projection weight of shape [intermediate_size, hidden_size]. + down_weight: Down projection weight of shape [hidden_size, intermediate_size]. + gate_bias: Optional gate projection bias of shape [intermediate_size]. + up_bias: Optional up projection bias of shape [intermediate_size]. + down_bias: Optional down projection bias of shape [hidden_size]. + + Returns: + Output tensor of shape [..., hidden_size]. + """ + gate_out = F.linear(input, gate_weight, gate_bias) + up_out = F.linear(input, up_weight, up_bias) + hidden = F.silu(gate_out) * up_out + return F.linear(hidden, down_weight, down_bias) + + +@torch_swiglu_mlp.register_fake +def _( + input: torch.Tensor, + gate_weight: torch.Tensor, + up_weight: torch.Tensor, + down_weight: torch.Tensor, + gate_bias: Optional[torch.Tensor], + up_bias: Optional[torch.Tensor], + down_bias: Optional[torch.Tensor], +) -> torch.Tensor: + """Fake implementation for tracing.""" + # Output shape is [..., hidden_size] where hidden_size = down_weight.shape[0] + output_shape = list(input.shape[:-1]) + [down_weight.shape[0]] + return input.new_empty(output_shape, dtype=input.dtype) + + +@torch.library.custom_op("auto_deploy::fused_swiglu_mlp", mutates_args=()) +def fused_swiglu_mlp( + input: torch.Tensor, + gate_up_weight: torch.Tensor, + down_weight: torch.Tensor, + gate_up_bias: Optional[torch.Tensor], + down_bias: Optional[torch.Tensor], +) -> torch.Tensor: + """Fused SwiGLU MLP with concatenated gate+up weights. + + Performs a single matmul for gate and up projections, then splits the result. + Computes: silu(gate_out) * up_out @ down.T + down_bias + where gate_out, up_out = split(x @ gate_up.T + gate_up_bias) + + Args: + input: Input tensor of shape [..., hidden_size]. + gate_up_weight: Concatenated gate+up weight of shape [2*intermediate_size, hidden_size]. + down_weight: Down projection weight of shape [hidden_size, intermediate_size]. + gate_up_bias: Optional concatenated gate+up bias of shape [2*intermediate_size]. + down_bias: Optional down projection bias of shape [hidden_size]. + + Returns: + Output tensor of shape [..., hidden_size]. + """ + # Single matmul for both gate and up projections + gate_up_out = F.linear(input, gate_up_weight, gate_up_bias) + + # Split into gate and up outputs + gate_out, up_out = gate_up_out.chunk(2, dim=-1) + + # Apply SwiGLU activation: silu(gate) * up + hidden = F.silu(gate_out) * up_out + + # Down projection + return F.linear(hidden, down_weight, down_bias) + + +@fused_swiglu_mlp.register_fake +def _( + input: torch.Tensor, + gate_up_weight: torch.Tensor, + down_weight: torch.Tensor, + gate_up_bias: Optional[torch.Tensor], + down_bias: Optional[torch.Tensor], +) -> torch.Tensor: + """Fake implementation for tracing.""" + # Output shape is [..., hidden_size] where hidden_size = down_weight.shape[0] + output_shape = list(input.shape[:-1]) + [down_weight.shape[0]] + return input.new_empty(output_shape, dtype=input.dtype) diff --git a/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_glm4_moe_lite.py b/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_glm4_moe_lite.py index 9d672371af7d..abc972f791d2 100644 --- a/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_glm4_moe_lite.py +++ b/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_glm4_moe_lite.py @@ -442,7 +442,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: final_hidden_states = final_hidden_states.view(*orig_shape) final_hidden_states = final_hidden_states + shared_output - return final_hidden_states.to(hidden_states.dtype) + return final_hidden_states class Glm4MoeLiteAttention(nn.Module): diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_swiglu.py b/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_swiglu.py new file mode 100644 index 000000000000..7c8ef59522fe --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_swiglu.py @@ -0,0 +1,284 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Graph transforms for SwiGLU MLP fusion. + +This module provides two-stage transformation for SwiGLU MLP: +1. MatchSwiGLUPattern: Detects SwiGLU patterns and replaces with torch_swiglu_mlp +2. FuseSwiGLU: Fuses gate+up weights into a single concatenated matmul + +The SwiGLU pattern is: silu(x @ gate.T) * (x @ up.T) @ down.T +""" + +from typing import Tuple, Type + +import torch +from pydantic import Field +from torch.fx import GraphModule, Node + +# Import the custom ops to ensure they are registered and for use in replacements +from ...custom_ops.linear.swiglu import torch_swiglu_mlp +from ...models.factory import ModelFactory +from ...shim.interface import CachedSequenceInterface +from ...utils._graph import delete_all_unused_submodules, eliminate_dead_code, get_attr_by_name +from ...utils.node_utils import is_op +from ...utils.pattern_matcher import ADPatternMatcherPass, register_ad_pattern +from ..interface import ( + BaseTransform, + SharedConfig, + TransformConfig, + TransformInfo, + TransformRegistry, +) + + +def _swiglu_pattern_no_bias(x, gate_weight, up_weight, down_weight): + """Pattern for SwiGLU MLP without biases. + + Matches: silu(linear(x, gate_weight, None)) * linear(x, up_weight, None) -> linear(down_weight, None) + """ + gate_out = torch.ops.auto_deploy.torch_linear_simple.default(x, gate_weight, None) + up_out = torch.ops.auto_deploy.torch_linear_simple.default(x, up_weight, None) + silu_out = torch.ops.aten.silu.default(gate_out) + mul_out = torch.ops.aten.mul.Tensor(silu_out, up_out) + down_out = torch.ops.auto_deploy.torch_linear_simple.default(mul_out, down_weight, None) + return down_out + + +def _swiglu_replacement_no_bias(x, gate_weight, up_weight, down_weight): + """Replacement for SwiGLU pattern without biases.""" + # Call the Python wrapper directly, not via torch.ops.auto_deploy + # This ensures proper FakeTensor mode handling during tracing + return torch_swiglu_mlp(x, gate_weight, up_weight, down_weight, None, None, None) + + +def _swiglu_pattern_with_bias( + x, gate_weight, up_weight, down_weight, gate_bias, up_bias, down_bias +): + """Pattern for SwiGLU MLP with biases. + + Matches: silu(linear(x, gate_weight, gate_bias)) * linear(x, up_weight, up_bias) -> linear(down_weight, down_bias) + """ + gate_out = torch.ops.auto_deploy.torch_linear_simple.default(x, gate_weight, gate_bias) + up_out = torch.ops.auto_deploy.torch_linear_simple.default(x, up_weight, up_bias) + silu_out = torch.ops.aten.silu.default(gate_out) + mul_out = torch.ops.aten.mul.Tensor(silu_out, up_out) + down_out = torch.ops.auto_deploy.torch_linear_simple.default(mul_out, down_weight, down_bias) + return down_out + + +def _swiglu_replacement_with_bias( + x, gate_weight, up_weight, down_weight, gate_bias, up_bias, down_bias +): + """Replacement for SwiGLU pattern with biases.""" + # Call the Python wrapper directly, not via torch.ops.auto_deploy + # This ensures proper FakeTensor mode handling during tracing + return torch_swiglu_mlp(x, gate_weight, up_weight, down_weight, gate_bias, up_bias, down_bias) + + +@TransformRegistry.register("match_swiglu_pattern") +class MatchSwiGLUPattern(BaseTransform): + """Matches SwiGLU MLP patterns and replaces with torch_swiglu_mlp op. + + This transform runs in the pattern_matcher stage and detects the following pattern: + silu(x @ gate.T) * (x @ up.T) @ down.T + + And replaces it with a single torch_swiglu_mlp op that can be fused later. + + Uses ADPatternMatcherPass for declarative pattern matching. + """ + + config: TransformConfig + + @classmethod + def get_config_class(cls) -> Type[TransformConfig]: + return TransformConfig + + def _apply( + self, + gm: GraphModule, + cm: CachedSequenceInterface, + factory: ModelFactory, + shared_config: SharedConfig, + ) -> Tuple[GraphModule, TransformInfo]: + patterns = ADPatternMatcherPass() + + # Dummy shapes for tracing - shapes don't matter for matching + hidden, intermediate = 128, 256 + + # Pattern 1: SwiGLU without biases (most common case) + dummy_args_no_bias = [ + torch.randn(2, hidden, device="meta", dtype=torch.float16), # x + torch.randn(intermediate, hidden, device="meta", dtype=torch.float16), # gate_weight + torch.randn(intermediate, hidden, device="meta", dtype=torch.float16), # up_weight + torch.randn(hidden, intermediate, device="meta", dtype=torch.float16), # down_weight + ] + register_ad_pattern( + search_fn=_swiglu_pattern_no_bias, + replace_fn=_swiglu_replacement_no_bias, + patterns=patterns, + dummy_args=dummy_args_no_bias, + ) + + # Pattern 2: SwiGLU with biases + dummy_args_with_bias = [ + torch.randn(2, hidden, device="meta", dtype=torch.float16), # x + torch.randn(intermediate, hidden, device="meta", dtype=torch.float16), # gate_weight + torch.randn(intermediate, hidden, device="meta", dtype=torch.float16), # up_weight + torch.randn(hidden, intermediate, device="meta", dtype=torch.float16), # down_weight + torch.randn(intermediate, device="meta", dtype=torch.float16), # gate_bias + torch.randn(intermediate, device="meta", dtype=torch.float16), # up_bias + torch.randn(hidden, device="meta", dtype=torch.float16), # down_bias + ] + register_ad_pattern( + search_fn=_swiglu_pattern_with_bias, + replace_fn=_swiglu_replacement_with_bias, + patterns=patterns, + dummy_args=dummy_args_with_bias, + ) + + num_matches = patterns.apply(gm.graph) + + if num_matches > 0: + gm.recompile() + + info = TransformInfo( + skipped=False, + num_matches=num_matches, + is_clean=num_matches == 0, + has_valid_shapes=num_matches == 0, + ) + + return gm, info + + +class FuseSwiGLUConfig(TransformConfig): + """Configuration for the SwiGLU fusion transform.""" + + enabled: bool = Field( + default=True, + description="Whether to enable SwiGLU fusion.", + ) + + +@TransformRegistry.register("fuse_swiglu") +class FuseSwiGLU(BaseTransform): + """Fuses torch_swiglu_mlp ops by concatenating gate and up weights. + + This transform runs in the post_load_fusion stage and replaces torch_swiglu_mlp ops + with fused_swiglu_mlp ops that use a single concatenated gate+up weight matrix. + + This reduces memory bandwidth by performing a single matmul instead of two + separate matmuls for gate and up projections. + """ + + config: FuseSwiGLUConfig + + @classmethod + def get_config_class(cls) -> Type[FuseSwiGLUConfig]: + return FuseSwiGLUConfig + + def _apply( + self, + gm: GraphModule, + cm: CachedSequenceInterface, + factory: ModelFactory, + shared_config: SharedConfig, + ) -> Tuple[GraphModule, TransformInfo]: + if not self.config.enabled: + return gm, TransformInfo(skipped=True, num_matches=0) + + graph = gm.graph + cnt = 0 + fused_weight_idx = 0 + + for node in list(graph.nodes): + if not is_op(node, torch.ops.auto_deploy.torch_swiglu_mlp.default): + continue + + # Extract args: (input, gate_weight, up_weight, down_weight, gate_bias, up_bias, down_bias) + input_node = node.args[0] + gate_weight_node = node.args[1] + up_weight_node = node.args[2] + down_weight_node = node.args[3] + gate_bias_node = node.args[4] if len(node.args) > 4 else None + up_bias_node = node.args[5] if len(node.args) > 5 else None + down_bias_node = node.args[6] if len(node.args) > 6 else None + + # Get the actual weight tensors + gate_weight = get_attr_by_name(gm, gate_weight_node.target) + up_weight = get_attr_by_name(gm, up_weight_node.target) + + # Concatenate gate and up weights: [intermediate, hidden] -> [2*intermediate, hidden] + gate_up_weight = torch.cat([gate_weight, up_weight], dim=0) + + # Create new attribute for the fused weight + fused_weight_name = f"fused_swiglu_gate_up_{fused_weight_idx}" + gm.register_buffer(fused_weight_name, gate_up_weight) + + # Handle biases + gate_up_bias_node = None + fused_bias_name = None + + if gate_bias_node is not None and gate_bias_node.op == "get_attr": + gate_bias = get_attr_by_name(gm, gate_bias_node.target) + up_bias = get_attr_by_name(gm, up_bias_node.target) if up_bias_node else None + + if up_bias is not None: + gate_up_bias = torch.cat([gate_bias, up_bias], dim=0) + fused_bias_name = f"fused_swiglu_gate_up_bias_{fused_weight_idx}" + gm.register_buffer(fused_bias_name, gate_up_bias) + + # Create get_attr node for the fused weight + with graph.inserting_before(node): + fused_weight_node = graph.get_attr(fused_weight_name) + + if fused_bias_name is not None: + gate_up_bias_node = graph.get_attr(fused_bias_name) + + # Create the fused_swiglu_mlp node + with graph.inserting_after(node): + fused_node: Node = graph.call_function( + torch.ops.auto_deploy.fused_swiglu_mlp.default, + args=( + input_node, + fused_weight_node, + down_weight_node, + gate_up_bias_node, + down_bias_node, + ), + ) + + # Replace uses and erase old node + node.replace_all_uses_with(fused_node) + graph.erase_node(node) + + fused_weight_idx += 1 + cnt += 1 + + if cnt > 0: + gm.recompile() + + # Clean up dead code and free memory from unfused weights + # The original gate_weight and up_weight tensors are no longer referenced + # after fusion, so we can delete them to save GPU memory. + eliminate_dead_code(gm) + delete_all_unused_submodules(gm) + + info = TransformInfo( + skipped=False, num_matches=cnt, is_clean=cnt == 0, has_valid_shapes=cnt == 0 + ) + + return gm, info diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/fused_add_rms_norm.py b/tensorrt_llm/_torch/auto_deploy/transform/library/fused_add_rms_norm.py index 103578ca01f0..dee221e742be 100644 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/fused_add_rms_norm.py +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/fused_add_rms_norm.py @@ -9,31 +9,43 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Transformation for fusing Add + Cast + RMSNorm.""" +"""Transformation for fusing Add + (optional Cast) + RMSNorm via direct FX graph manipulation.""" -from typing import Tuple +import operator +from typing import List, Optional, Tuple import torch -from torch.fx import GraphModule +from torch.fx import GraphModule, Node from ...custom_ops.normalization.flashinfer_fused_add_rms_norm import flashinfer_fused_add_rms_norm from ...models.factory import ModelFactory from ...shim.interface import CachedSequenceInterface -from ...utils.pattern_matcher import ADPatternMatcherPass, register_ad_pattern +from ...utils._graph import eliminate_dead_code +from ...utils.node_utils import is_op from ..interface import BaseTransform, SharedConfig, TransformInfo, TransformRegistry @TransformRegistry.register("fuse_add_rms_norm") class FuseAddRMSNorm(BaseTransform): - """Fuse (add + cast + RMSNorm) into one fused op. + """Fuse (add + optional cast + RMSNorm) into one fused op. - Matches: - x = add(input, residual) - y = x.to(dtype) - z = flashinfer_rms_norm(y, weight, eps) + Uses direct FX graph manipulation instead of the inductor pattern matcher + to correctly handle patterns where intermediate nodes (add, rms_norm) have + multiple users in the graph. - Replaces with: - z, x = flashinfer_fused_add_rms_norm(input, residual, weight, eps) + Pattern 1 (without cast): + %add = aten.add(%x, %residual) + %norm = flashinfer_rms_norm(%add, %weight, eps) + + Pattern 2 (with cast): + %add = aten.add(%x, %residual) + %cast = aten.to.dtype(%add, bfloat16) + %norm = flashinfer_rms_norm(%cast, %weight, eps) + + Both are replaced with: + %fused = flashinfer_fused_add_rms_norm(%x, %residual, %weight, eps) + %norm_out = getitem(%fused, 0) # norm result (replaces %norm) + %add_out = getitem(%fused, 1) # add result (replaces %add) """ def _apply( @@ -43,42 +55,102 @@ def _apply( factory: ModelFactory, shared_config: SharedConfig, ) -> Tuple[GraphModule, TransformInfo]: - patterns = ADPatternMatcherPass() - - # Dummy shapes for tracing - bsz, hidden = 2, 128 - dummy_args = [ - torch.randn(bsz, hidden, device="meta", dtype=torch.bfloat16), # x (bf16) - torch.randn(bsz, hidden, device="meta", dtype=torch.bfloat16), # residual (bf16) - torch.randn(hidden, device="meta", dtype=torch.bfloat16), # weight - 1e-5, # eps - ] - - op_ignore_types = {torch.ops.aten.to.dtype: (torch.dtype,)} - scalar_workaround = {"eps": 1e-5} - - def _fused_add_norm_pattern(x, residual, weight, eps): - added = torch.ops.aten.add.Tensor(x, residual) - cast = torch.ops.aten.to.dtype(added, torch.bfloat16) - # Note: we assume flashinfer_rms_norm is the target - norm = torch.ops.auto_deploy.flashinfer_rms_norm.default(cast, weight, eps) - return norm, added - - def _fused_add_norm_replacement(x, residual, weight, eps): - # Use the python wrapper directly, not via torch.ops.auto_deploy - return flashinfer_fused_add_rms_norm(x, residual, weight, eps) - - # Register pattern - register_ad_pattern( - search_fn=_fused_add_norm_pattern, - replace_fn=_fused_add_norm_replacement, - patterns=patterns, - dummy_args=dummy_args, - op_ignore_types=op_ignore_types, - scalar_workaround=scalar_workaround, - ) - - num_matches = patterns.apply(gm.graph) + graph = gm.graph + num_matches = 0 + + # --- Step 1: collect (add_node, optional cast_node, norm_node) triples --- + matches: List[Tuple[Node, Optional[Node], Node]] = [] + + for node in graph.nodes: + # Match flashinfer_rms_norm (handles both overload packet and .default) + if not is_op(node, torch.ops.auto_deploy.flashinfer_rms_norm): + continue + + input_to_norm = node.args[0] + cast_node: Optional[Node] = None + + # Check for an optional aten.to.dtype cast between add and norm + if isinstance(input_to_norm, Node) and is_op(input_to_norm, torch.ops.aten.to.dtype): + cast_node = input_to_norm + input_to_norm = cast_node.args[0] + + # The (possibly unwrapped) input must be an aten.add.Tensor + if not isinstance(input_to_norm, Node) or not is_op( + input_to_norm, torch.ops.aten.add.Tensor + ): + continue + + add_node = input_to_norm + matches.append((add_node, cast_node, node)) + + # --- Step 2: apply fusions --- + erased: set = set() # track erased node ids to skip stale matches + + for add_node, cast_node, norm_node in matches: + # Safety: skip if a node in this match was already consumed + if id(add_node) in erased or id(norm_node) in erased: + continue + + # Original operands + add_lhs = add_node.args[0] # e.g. previous residual + add_rhs = add_node.args[1] # e.g. attention/MoE output + weight = norm_node.args[1] + eps = norm_node.args[2] + + # Insert the fused call right before the norm node. Using + # inserting_before(norm_node) ensures correct topological order: + # fused_node → norm_out → add_out all appear before norm_node. + with graph.inserting_before(norm_node): + # flashinfer_fused_add_rms_norm(x, residual, weight, eps): + # residual += x → residual becomes add result + # x = rms_norm(residual) → x becomes norm result + # returns (x, residual) = (norm_result, add_result) + fused_node = graph.call_function( + flashinfer_fused_add_rms_norm, + args=(add_rhs, add_lhs, weight, eps), + ) + norm_out = graph.call_function(operator.getitem, args=(fused_node, 0)) + add_out = graph.call_function(operator.getitem, args=(fused_node, 1)) + + # Rewire all consumers of the original norm → norm_out + norm_node.replace_all_uses_with(norm_out) + + # Erase norm first so cast_node (if present) loses its only user + graph.erase_node(norm_node) + erased.add(id(norm_node)) + + # Erase cast_node *before* replacing add's uses, otherwise + # replace_all_uses_with would rewrite cast's input to add_out + # which sits after cast in the graph → topological violation. + if cast_node is not None: + if len(cast_node.users) == 0: + graph.erase_node(cast_node) + erased.add(id(cast_node)) + else: + # Rare: cast has users besides norm. Redirect them to a + # new cast placed after add_out so the ordering is valid. + with graph.inserting_before(list(cast_node.users)[0]): + new_cast = graph.call_function( + cast_node.target, + args=(add_out, *cast_node.args[1:]), + kwargs=cast_node.kwargs, + ) + cast_node.replace_all_uses_with(new_cast) + graph.erase_node(cast_node) + erased.add(id(cast_node)) + + # Rewire all consumers of the original add → add_out + # (includes the residual connection to the *next* layer's add) + add_node.replace_all_uses_with(add_out) + + graph.erase_node(add_node) + erased.add(id(add_node)) + + num_matches += 1 + + # Clean up any remaining dead code + if num_matches > 0: + eliminate_dead_code(gm) info = TransformInfo( skipped=False, diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.py b/tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.py index 8c217d0d41d7..ea1407dc681d 100644 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.py +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.py @@ -88,17 +88,23 @@ def wait_event(device: int, stream_name: str) -> None: event.wait() -# skip during compilation @torch._dynamo.disable -def record_event_wrapper( - fn: Callable, - *args: Tuple[Any, ...], - **kwargs: Dict[str, Any], +def record_event_passthrough( + x: torch.Tensor, + *, + device: int = -1, ) -> torch.Tensor: - device = kwargs.pop("device", torch.cuda.current_device()) - output = fn(*args, **kwargs) + """Record a CUDA event on the main stream and return the input unchanged. + + Inserted after the gating/routing computation to mark a synchronization + point. The aux stream waits for this event before starting the MoE + computation, enabling overlap between the shared expert (main stream) + and routed experts (aux stream). + """ + if device < 0: + device = torch.cuda.current_device() torch.ops.auto_deploy.record_event(device, cuda_stream_manager.MAIN_STREAM_NAME) - return output + return x @torch._dynamo.disable @@ -254,38 +260,58 @@ def trtllm_quant_fp8_moe_fused_aux_fake( def _execute_op_in_aux_stream( gm: GraphModule, op_dict: Dict[Callable, Callable] ) -> Tuple[GraphModule, int]: + """Replace MoE ops with aux-stream variants and insert CUDA event synchronisation. + + For each MoE op we: + 1. Identify its *routing* inputs (selected_experts, routing_weights) — + everything except the hidden-state tensor (``args[0]``) and static + weight parameters (``get_attr`` nodes). + 2. Insert ``record_event_passthrough`` right after the latest routing + input in graph order. This records a CUDA event on the main stream + *after* gating is done but *before* the shared expert starts, + enabling the shared expert (main stream) and routed experts (aux + stream) to execute concurrently. + 3. Replace the original MoE op with its ``_aux`` variant that runs on + the auxiliary CUDA stream. + """ graph = gm.graph num_replaced = 0 # Collect targets first to avoid mutating while iterating target_nodes = [n for n in graph.nodes if is_op(n, op_dict.keys())] + # Topological position map — used to find the latest routing input. + node_order = {node: i for i, node in enumerate(graph.nodes)} + for n in target_nodes: - target_input_node = None - for input_node in n.all_input_nodes: - if input_node.target == torch.ops.aten.view.default: - target_input_node = input_node - break - # Look through dtype cast nodes (aten.to) to find the view node - if input_node.target == torch.ops.aten.to: - for nested_input in input_node.all_input_nodes: - if nested_input.target == torch.ops.aten.view.default: - target_input_node = nested_input - break - if target_input_node is not None: - break - - assert target_input_node is not None, f"Target input node not found for node {n}" - with graph.inserting_before(target_input_node): - kwargs = target_input_node.kwargs.copy() - kwargs["device"] = torch.cuda.current_device() - new_node = graph.call_function( - record_event_wrapper, - args=(target_input_node.target, *target_input_node.args), - kwargs=kwargs, + # MoE ops follow the convention (x, selected_experts, routing_weights, ...weights...). + # args[0] is the hidden-state tensor; the remaining *compute* (non-get_attr) + # inputs are routing results produced by the gating computation. + hidden_state_node = n.args[0] + routing_inputs = [ + inp + for inp in n.all_input_nodes + if inp is not hidden_state_node and inp.op != "get_attr" + ] + assert routing_inputs, f"No routing inputs found for MoE node {n}" + + # Place the event record right after the last routing dependency so + # that the aux stream's wait_event sees all routing results as ready + # while the shared expert (later in graph order) can still overlap. + latest_routing = max(routing_inputs, key=lambda inp: node_order.get(inp, 0)) + + with graph.inserting_after(latest_routing): + rec_node = graph.call_function( + record_event_passthrough, + args=(latest_routing,), + kwargs={"device": torch.cuda.current_device()}, ) - target_input_node.replace_all_uses_with(new_node) - graph.erase_node(target_input_node) + + # Wire the event node into the MoE op so there is a true data + # dependency that prevents reordering. + n.args = tuple(rec_node if arg is latest_routing else arg for arg in n.args) + + # Replace MoE op with its aux-stream variant with graph.inserting_after(n): new_node = graph.call_function(op_dict[n.target], args=n.args, kwargs=n.kwargs) n.replace_all_uses_with(new_node) diff --git a/tests/integration/defs/accuracy/test_llm_api_autodeploy.py b/tests/integration/defs/accuracy/test_llm_api_autodeploy.py index 630719919b3a..638887059099 100644 --- a/tests/integration/defs/accuracy/test_llm_api_autodeploy.py +++ b/tests/integration/defs/accuracy/test_llm_api_autodeploy.py @@ -426,8 +426,8 @@ class TestGLM4Flash(LlmapiAccuracyTestHarness): ``` """ - # MODEL_NAME = "zai-org/GLM-4.7-Flash" - MODEL_NAME = "DeepInfra/GLM-4.7-Flash-NVFP4" + MODEL_NAME = "zai-org/GLM-4.7-Flash" + #MODEL_NAME = "DeepInfra/GLM-4.7-Flash-NVFP4" MODEL_PATH = MODEL_NAME # Model is in HF_CACHE # Set minimum possible seq len + small buffer, for test speed & memory usage MAX_SEQ_LEN = max(MMLU.MAX_INPUT_LEN + MMLU.MAX_OUTPUT_LEN, @@ -452,6 +452,13 @@ def get_default_kwargs(self, enable_chunked_prefill=False): "model_kwargs": { "torch_dtype": "bfloat16" }, + "transforms": { + "multi_stream_moe": { + "stage": "compile", + # multi-stream MOE currently does not work for world_size > 1 + "enabled": True, + }, + } } if enable_chunked_prefill: config["enable_chunked_prefill"] = True diff --git a/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_multi_stream.py b/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_multi_stream.py index 182c6a0a31bd..7f32cbbf6b0b 100644 --- a/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_multi_stream.py +++ b/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_multi_stream.py @@ -7,7 +7,7 @@ from tensorrt_llm._torch.auto_deploy.transform.library.multi_stream_moe import ( aux_stream_wrapper, cuda_stream_manager, - record_event_wrapper, + record_event_passthrough, ) from tensorrt_llm._torch.auto_deploy.utils._graph import canonicalize_graph from tensorrt_llm._torch.auto_deploy.utils.node_utils import is_op @@ -32,12 +32,14 @@ def multi_stream_linear_fake(input, weight0, weight1): def replace_multi_stream_linear_with_aux_stream_wrapper(gm: GraphModule) -> Tuple[GraphModule, int]: """Traverse ``gm`` and replace all ``auto_deploy::multi_stream_linear`` ops with ``aux_stream_wrapper``. - The replacement preserves the original args/kwargs of the node. - After rewriting, the graph is cleaned and recompiled. - - Args: - gm: The FX graph module to transform. - aux_stream_wrapper: A callable to replace the custom op with. + For each target op we: + 1. Find the shared input (a node with >1 users, e.g. relu output used + by both a main-stream consumer and the target op). + 2. Insert ``record_event_passthrough`` right after it to mark a CUDA + synchronisation point — the aux stream waits for this event before + starting. + 3. Wire the event node into the target op (other users of the shared + input are unaffected) and replace the op with ``aux_stream_wrapper``. Returns: A tuple of (gm, num_replaced) @@ -46,27 +48,34 @@ def replace_multi_stream_linear_with_aux_stream_wrapper(gm: GraphModule) -> Tupl num_replaced = 0 # Collect targets first to avoid mutating while iterating - target_nodes: list[Node] = [] - target_nodes = [n for n in graph.nodes if is_op(n, torch.ops.auto_deploy.multi_stream_linear)] + target_nodes: list[Node] = [ + n for n in graph.nodes if is_op(n, torch.ops.auto_deploy.multi_stream_linear) + ] for n in target_nodes: - target_input_node = None + # Find the shared input — a node consumed by multiple ops. + # Analogous to the routing input in the MoE transform. + shared_input = None for input_node in n.all_input_nodes: if len(input_node.users) > 1: - target_input_node = input_node + shared_input = input_node break - if target_input_node is None: - raise ValueError(f"Target input node not found for node {n}") - with graph.inserting_before(target_input_node): - kwargs = target_input_node.kwargs.copy() - kwargs["device"] = torch.cuda.current_device() - new_node = graph.call_function( - record_event_wrapper, - args=(target_input_node.target, *target_input_node.args), - kwargs=kwargs, + if shared_input is None: + raise ValueError(f"Shared input node not found for node {n}") + + # Record CUDA event right after the shared input so the aux stream + # can start as soon as this dependency is ready, while the main + # stream continues with other consumers (e.g. fc2). + with graph.inserting_after(shared_input): + rec_node = graph.call_function( + record_event_passthrough, + args=(shared_input,), + kwargs={"device": torch.cuda.current_device()}, ) - target_input_node.replace_all_uses_with(new_node) - graph.erase_node(target_input_node) + + # Wire the event node into the target op only. + n.args = tuple(rec_node if arg is shared_input else arg for arg in n.args) + with graph.inserting_after(n): new_node = graph.call_function( aux_stream_wrapper, args=(n.target, *n.args), kwargs=n.kwargs diff --git a/tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_fuse_swiglu.py b/tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_fuse_swiglu.py new file mode 100644 index 000000000000..6ea74d972123 --- /dev/null +++ b/tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_fuse_swiglu.py @@ -0,0 +1,188 @@ +import pytest +import torch +from torch.export import Dim + +from tensorrt_llm._torch.auto_deploy.custom_ops.linear.swiglu import * # noqa +from tensorrt_llm._torch.auto_deploy.export import torch_export_to_gm +from tensorrt_llm._torch.auto_deploy.transform.optimizer import InferenceOptimizer +from tensorrt_llm._torch.auto_deploy.utils.node_utils import is_op + + +class SwiGLUMLP(torch.nn.Module): + """SwiGLU MLP module: silu(x @ gate.T) * (x @ up.T) @ down.T""" + + def __init__(self, hidden_size: int, intermediate_size: int): + super().__init__() + self.gate_proj = torch.nn.Linear(hidden_size, intermediate_size, bias=False) + self.up_proj = torch.nn.Linear(hidden_size, intermediate_size, bias=False) + self.down_proj = torch.nn.Linear(intermediate_size, hidden_size, bias=False) + + def forward(self, x): + return self.down_proj(torch.nn.functional.silu(self.gate_proj(x)) * self.up_proj(x)) + + +class SwiGLUMLPWithBias(torch.nn.Module): + """SwiGLU MLP module with biases.""" + + def __init__(self, hidden_size: int, intermediate_size: int): + super().__init__() + self.gate_proj = torch.nn.Linear(hidden_size, intermediate_size, bias=True) + self.up_proj = torch.nn.Linear(hidden_size, intermediate_size, bias=True) + self.down_proj = torch.nn.Linear(intermediate_size, hidden_size, bias=True) + + def forward(self, x): + return self.down_proj(torch.nn.functional.silu(self.gate_proj(x)) * self.up_proj(x)) + + +class SwiGLUTestModel(torch.nn.Module): + """Test model with SwiGLU MLP sandwiched between linear layers.""" + + def __init__( + self, + hidden_size: int = 256, + intermediate_size: int = 512, + with_bias: bool = False, + ): + super().__init__() + self.linear1 = torch.nn.Linear(hidden_size, hidden_size, device="cuda", dtype=torch.float16) + if with_bias: + self.mlp = SwiGLUMLPWithBias(hidden_size, intermediate_size) + else: + self.mlp = SwiGLUMLP(hidden_size, intermediate_size) + self.mlp = self.mlp.to(device="cuda", dtype=torch.float16) + self.linear2 = torch.nn.Linear(hidden_size, hidden_size, device="cuda", dtype=torch.float16) + + def forward(self, x): + x = self.linear1(x) + x = self.mlp(x) + x = self.linear2(x) + return x + + +class SwiGLUTestModelMultipleMLP(torch.nn.Module): + """Test model with multiple SwiGLU MLPs to test multiple pattern matches.""" + + def __init__( + self, + hidden_size: int = 256, + intermediate_size: int = 512, + num_layers: int = 2, + ): + super().__init__() + self.layers = torch.nn.ModuleList() + for _ in range(num_layers): + self.layers.append( + torch.nn.ModuleDict( + { + "linear": torch.nn.Linear( + hidden_size, hidden_size, device="cuda", dtype=torch.float16 + ), + "mlp": SwiGLUMLP(hidden_size, intermediate_size).to( + device="cuda", dtype=torch.float16 + ), + } + ) + ) + + def forward(self, x): + for layer in self.layers: + x = layer["linear"](x) + x = layer["mlp"](x) + return x + + +def _run_fusion_test(model, expected_op, expected_num_matches=1): + """Run the SwiGLU fusion test. + + Args: + model: The test model to transform. + expected_op: The expected fused op to find in the transformed graph. + expected_num_matches: Expected number of fused ops. + """ + x = torch.randn(2, 256, device="cuda", dtype=torch.float16) + dynamic_shapes = {0: Dim.DYNAMIC} + gm = torch_export_to_gm(model, args=(x,), dynamic_shapes=(dynamic_shapes,), clone=True) + + # Apply transforms + gm_transformed = InferenceOptimizer( + None, + { + "match_swiglu_pattern": { + "stage": "pattern_matcher", + }, + "fuse_swiglu": { + "stage": "post_load_fusion", + "enabled": True, + }, + }, + )(None, gm) + + # Move to CUDA if needed + gm_transformed = gm_transformed.to("cuda") + + # Check that the expected op is present + count = sum(1 for n in gm_transformed.graph.nodes if is_op(n, expected_op)) + assert count == expected_num_matches, ( + f"Expected {expected_num_matches} {expected_op} ops, got {count}" + ) + + # Verify numerical correctness + y_transformed = gm_transformed(x) + y_model = model(x) + torch.testing.assert_close(y_transformed, y_model, atol=1e-2, rtol=1e-2) + + # Test with a different batch size + new_input = torch.randn(4, 256, device="cuda", dtype=torch.float16) + y_transformed_2 = gm_transformed(new_input) + y_model_2 = model(new_input) + torch.testing.assert_close(y_transformed_2, y_model_2, atol=1e-2, rtol=1e-2) + + +def test_swiglu_fusion_basic(): + """Test basic SwiGLU fusion without biases.""" + model = SwiGLUTestModel(with_bias=False) + _run_fusion_test(model, torch.ops.auto_deploy.fused_swiglu_mlp.default) + + +def test_swiglu_fusion_with_bias(): + """Test SwiGLU fusion with biases.""" + model = SwiGLUTestModel(with_bias=True) + _run_fusion_test(model, torch.ops.auto_deploy.fused_swiglu_mlp.default) + + +@pytest.mark.parametrize("num_layers", [2, 3]) +def test_swiglu_fusion_multiple_layers(num_layers): + """Test that multiple SwiGLU patterns are fused correctly.""" + model = SwiGLUTestModelMultipleMLP(num_layers=num_layers) + _run_fusion_test( + model, torch.ops.auto_deploy.fused_swiglu_mlp.default, expected_num_matches=num_layers + ) + + +def test_swiglu_pattern_match_only(): + """Test pattern matching stage only (without fusion).""" + model = SwiGLUTestModel() + x = torch.randn(2, 256, device="cuda", dtype=torch.float16) + dynamic_shapes = {0: Dim.DYNAMIC} + gm = torch_export_to_gm(model, args=(x,), dynamic_shapes=(dynamic_shapes,), clone=True) + + # Only run pattern matching, not fusion + gm_matched = InferenceOptimizer( + None, + { + "match_swiglu_pattern": { + "stage": "pattern_matcher", + }, + }, + )(None, gm) + + # Check that the intermediate op is present + has_swiglu_op = any( + is_op(n, torch.ops.auto_deploy.torch_swiglu_mlp.default) for n in gm_matched.graph.nodes + ) + assert has_swiglu_op, "Pattern matcher should produce torch_swiglu_mlp op" + + # Verify numerical correctness + y_matched = gm_matched(x) + y_model = model(x) + torch.testing.assert_close(y_matched, y_model, atol=1e-3, rtol=1e-3) diff --git a/tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_fused_add_rms_norm.py b/tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_fused_add_rms_norm.py index dc353e49936d..4dfd3f683ce5 100644 --- a/tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_fused_add_rms_norm.py +++ b/tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_fused_add_rms_norm.py @@ -1,14 +1,26 @@ +import operator + import torch from torch.export import Dim -from tensorrt_llm._torch.auto_deploy.custom_ops.normalization.flashinfer_fused_add_rms_norm import * # noqa +from tensorrt_llm._torch.auto_deploy.custom_ops.normalization.flashinfer_fused_add_rms_norm import ( # noqa + flashinfer_fused_add_rms_norm, +) from tensorrt_llm._torch.auto_deploy.custom_ops.normalization.rms_norm import * # noqa from tensorrt_llm._torch.auto_deploy.export import torch_export_to_gm +from tensorrt_llm._torch.auto_deploy.transform.interface import TransformConfig +from tensorrt_llm._torch.auto_deploy.transform.library.fused_add_rms_norm import FuseAddRMSNorm from tensorrt_llm._torch.auto_deploy.transform.optimizer import InferenceOptimizer from tensorrt_llm._torch.auto_deploy.utils.node_utils import is_op +# --------------------------------------------------------------------------- +# Models +# --------------------------------------------------------------------------- + + +class AddCastNormModel(torch.nn.Module): + """Pattern 1: add + cast(to.dtype) + rms_norm.""" -class TestModel(torch.nn.Module): def __init__(self, hidden_size=128, eps=1e-5): super().__init__() self.weight = torch.nn.Parameter( @@ -23,55 +35,263 @@ def forward(self, x, residual): return norm, added -def _run_test(model): - # The replacement uses flashinfer_fused_add_rms_norm python wrapper which calls the inplace op - # auto_deploy::flashinfer_fused_add_rms_norm_inplace - op = torch.ops.auto_deploy.flashinfer_fused_add_rms_norm_inplace +class AddNormModel(torch.nn.Module): + """Pattern 2: add + rms_norm (no intermediate cast).""" + + def __init__(self, hidden_size=128, eps=1e-5): + super().__init__() + self.weight = torch.nn.Parameter( + torch.ones(hidden_size, device="cuda", dtype=torch.bfloat16) + ) + self.eps = eps + + def forward(self, x, residual): + added = x + residual + norm = torch.ops.auto_deploy.flashinfer_rms_norm(added, self.weight, self.eps) + return norm, added + + +class MultiUserModel(torch.nn.Module): + """Both add and rms_norm outputs have multiple users (DeepSeek V3 MoE pattern). + + add_result has 2 users: rms_norm + next residual add + norm_result has 2 users: linear1 + linear2 + """ + + def __init__(self, hidden_size=128, eps=1e-5): + super().__init__() + self.weight = torch.nn.Parameter( + torch.ones(hidden_size, device="cuda", dtype=torch.bfloat16) + ) + self.linear1 = torch.nn.Linear( + hidden_size, hidden_size, bias=False, device="cuda", dtype=torch.bfloat16 + ) + self.linear2 = torch.nn.Linear( + hidden_size, hidden_size, bias=False, device="cuda", dtype=torch.bfloat16 + ) + self.eps = eps + + def forward(self, residual, attn_output, moe_output): + # add with 2 users (norm + next_add) + add_result = residual + attn_output + # rms_norm with 2 users (linear1, linear2) + norm_result = torch.ops.auto_deploy.flashinfer_rms_norm(add_result, self.weight, self.eps) + out1 = self.linear1(norm_result) + out2 = self.linear2(norm_result) + combined = out1 + out2 + # add_result also feeds into next residual add + next_residual = add_result + moe_output + return combined, next_residual + + +class ChainedModel(torch.nn.Module): + """Two consecutive add+norm pairs sharing residual (like transformer layers). + + Layer 1: add1 = embed + attn_out, norm1 = rms_norm(add1) -- add1 has 2 users + Layer 2: add2 = add1 + mlp_out, norm2 = rms_norm(add2) -- add2 has 2 users + """ + + def __init__(self, hidden_size=128, eps=1e-5): + super().__init__() + self.weight1 = torch.nn.Parameter( + torch.ones(hidden_size, device="cuda", dtype=torch.bfloat16) + ) + self.weight2 = torch.nn.Parameter( + torch.ones(hidden_size, device="cuda", dtype=torch.bfloat16) + ) + self.linear = torch.nn.Linear( + hidden_size, hidden_size, bias=False, device="cuda", dtype=torch.bfloat16 + ) + self.eps = eps + + def forward(self, embed, attn_out, mlp_out): + add1 = embed + attn_out + norm1 = torch.ops.auto_deploy.flashinfer_rms_norm(add1, self.weight1, self.eps) + branch1 = self.linear(norm1) + + add2 = add1 + mlp_out + norm2 = torch.ops.auto_deploy.flashinfer_rms_norm(add2, self.weight2, self.eps) + branch2 = self.linear(norm2) + + return branch1 + branch2, add2 - def checker(gm): - return any(is_op(n, op) for n in gm.graph.nodes) +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _count_fused_ops(gm): + """Count flashinfer_fused_add_rms_norm wrapper calls in the graph.""" + return sum( + 1 + for n in gm.graph.nodes + if n.op == "call_function" and n.target is flashinfer_fused_add_rms_norm + ) + + +def _count_rms_norm_ops(gm): + """Count flashinfer_rms_norm calls in the graph.""" + return sum(1 for n in gm.graph.nodes if is_op(n, torch.ops.auto_deploy.flashinfer_rms_norm)) + + +def _count_add_ops(gm): + """Count aten.add.Tensor calls in the graph.""" + return sum(1 for n in gm.graph.nodes if is_op(n, torch.ops.aten.add.Tensor)) + + +def _export_model(model, *inputs, dynamic_dim0=True): + """Export a model to a GraphModule, optionally with a dynamic batch dimension.""" + if dynamic_dim0: + dyn = Dim.DYNAMIC + ds = tuple({0: dyn} for _ in inputs) + else: + ds = None + return torch_export_to_gm(model, args=inputs, dynamic_shapes=ds, clone=True) + + +def _apply_transform(gm): + """Apply fuse_add_rms_norm via InferenceOptimizer (integration-style).""" + return InferenceOptimizer( + None, + {"fuse_add_rms_norm": {"stage": "post_load_fusion"}}, + )(None, gm) + + +def _apply_transform_direct(gm): + """Apply the transform directly (unit-test style).""" + config = TransformConfig(stage="post_load_fusion") + transform = FuseAddRMSNorm(config=config) + gm, info = transform._apply(gm, None, None, None) + return gm, info + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +def test_fuse_add_cast_rms_norm(): + """Original test: add + cast(bf16) + rms_norm → fused op.""" + model = AddCastNormModel() bsz, seq_len, hidden = 2, 8, 128 - # Inputs should be bfloat16 x = torch.randn(bsz, seq_len, hidden, device="cuda", dtype=torch.bfloat16) + residual = torch.randn_like(x) + + gm = _export_model(model, x, residual) + gm_t = _apply_transform(gm) + + # Structure check + assert _count_fused_ops(gm_t) >= 1, "fused op not found in graph" + assert _count_rms_norm_ops(gm_t) == 0, "unfused rms_norm still in graph" + + # Numerical check + y_fused = gm_t(x.clone(), residual.clone()) + y_ref = model(x.clone(), residual.clone()) + torch.testing.assert_close(y_fused[0], y_ref[0], atol=1e-2, rtol=1e-2) + torch.testing.assert_close(y_fused[1], y_ref[1], atol=1e-2, rtol=1e-2) + + +def test_fuse_add_rms_norm_no_cast(): + """Pattern 2: add + rms_norm (no cast) → fused op.""" + model = AddNormModel() + bsz, seq_len, hidden = 2, 8, 128 + x = torch.randn(bsz, seq_len, hidden, device="cuda", dtype=torch.bfloat16) + residual = torch.randn_like(x) + + gm = _export_model(model, x, residual) + gm_t = _apply_transform(gm) + + # Structure check + assert _count_fused_ops(gm_t) >= 1, "fused op not found in graph" + assert _count_rms_norm_ops(gm_t) == 0, "unfused rms_norm still in graph" + + # Numerical check + y_fused = gm_t(x.clone(), residual.clone()) + y_ref = model(x.clone(), residual.clone()) + torch.testing.assert_close(y_fused[0], y_ref[0], atol=1e-2, rtol=1e-2) + torch.testing.assert_close(y_fused[1], y_ref[1], atol=1e-2, rtol=1e-2) + + +def test_fuse_add_rms_norm_multi_user(): + """Multi-user: both add (2 users) and rms_norm (2 users) → fused op. + + This is the key pattern from the DeepSeek V3 / GLM4-MoE graph that failed + with the old inductor-based pattern matcher due to num_users constraints. + """ + model = MultiUserModel() + bsz, seq_len, hidden = 2, 8, 128 residual = torch.randn(bsz, seq_len, hidden, device="cuda", dtype=torch.bfloat16) + attn_out = torch.randn_like(residual) + moe_out = torch.randn_like(residual) - # Dynamic shapes - dyn_batch_size = Dim.DYNAMIC - ds_x = {0: dyn_batch_size} - ds_res = {0: dyn_batch_size} + gm = _export_model(model, residual, attn_out, moe_out) - gm = torch_export_to_gm(model, args=(x, residual), dynamic_shapes=(ds_x, ds_res), clone=True) + # Before: 1 add+norm fusible pair, add has 2 users, norm has 2 users + assert _count_rms_norm_ops(gm) == 1 - gm_transformed = InferenceOptimizer( - None, - { - "fuse_add_rms_norm": { - "stage": "post_load_fusion", - }, - }, - )(None, gm) + gm_t, info = _apply_transform_direct(gm) - # Check if transform happened - if not checker(gm_transformed): - raise AssertionError( - "flashinfer_fused_add_rms_norm_inplace op not found in transformed graph" - ) + # Structure check + assert info.num_matches == 1, f"Expected 1 match, got {info.num_matches}" + assert _count_fused_ops(gm_t) == 1, "fused op not found in graph" + assert _count_rms_norm_ops(gm_t) == 0, "unfused rms_norm still in graph" + + # Verify getitem nodes for both outputs + getitems = [ + n + for n in gm_t.graph.nodes + if n.op == "call_function" + and n.target is operator.getitem + and isinstance(n.args[0], torch.fx.Node) + and n.args[0].target is flashinfer_fused_add_rms_norm + ] + assert len(getitems) == 2, f"Expected 2 getitem nodes, got {len(getitems)}" + + # Numerical check + y_fused = gm_t(residual.clone(), attn_out.clone(), moe_out.clone()) + y_ref = model(residual.clone(), attn_out.clone(), moe_out.clone()) + torch.testing.assert_close(y_fused[0], y_ref[0], atol=1e-2, rtol=1e-2) + torch.testing.assert_close(y_fused[1], y_ref[1], atol=1e-2, rtol=1e-2) + + +def test_fuse_add_rms_norm_chained(): + """Chained: two consecutive add+norm pairs across transformer layers.""" + model = ChainedModel() + bsz, seq_len, hidden = 2, 8, 128 + embed = torch.randn(bsz, seq_len, hidden, device="cuda", dtype=torch.bfloat16) + attn_out = torch.randn_like(embed) + mlp_out = torch.randn_like(embed) + + gm = _export_model(model, embed, attn_out, mlp_out) - # Validation - # Clone inputs because the fused op is inplace - x_in = x.clone() - res_in = residual.clone() + # Before: 2 add+norm fusible pairs + assert _count_rms_norm_ops(gm) == 2 - # The fused op is inplace, so inputs x_in and res_in will be modified. - # gm_transformed returns (x_in, res_in) which are the modified tensors. - y_transformed = gm_transformed(x_in, res_in) + gm_t, info = _apply_transform_direct(gm) - y_model = model(x.clone(), residual.clone()) - torch.testing.assert_close(y_transformed[0], y_model[0], atol=1e-2, rtol=1e-2) - torch.testing.assert_close(y_transformed[1], y_model[1], atol=1e-2, rtol=1e-2) + # Structure check + assert info.num_matches == 2, f"Expected 2 matches, got {info.num_matches}" + assert _count_fused_ops(gm_t) == 2, "Expected 2 fused ops" + assert _count_rms_norm_ops(gm_t) == 0, "unfused rms_norm still in graph" + # Verify second fused op receives add_out from first fused op (residual chain) + fused_nodes = [ + n + for n in gm_t.graph.nodes + if n.op == "call_function" and n.target is flashinfer_fused_add_rms_norm + ] + assert len(fused_nodes) == 2 + # The second fused op's residual arg should be a getitem from the first fused op + second_residual_arg = fused_nodes[1].args[1] + assert ( + second_residual_arg.op == "call_function" + and second_residual_arg.target is operator.getitem + and second_residual_arg.args[0] is fused_nodes[0] + ), "Second fused op's residual should come from first fused op's add_out" -def test_fuse_add_rms_norm(): - model = TestModel() - _run_test(model) + # Numerical check + y_fused = gm_t(embed.clone(), attn_out.clone(), mlp_out.clone()) + y_ref = model(embed.clone(), attn_out.clone(), mlp_out.clone()) + torch.testing.assert_close(y_fused[0], y_ref[0], atol=1e-2, rtol=1e-2) + torch.testing.assert_close(y_fused[1], y_ref[1], atol=1e-2, rtol=1e-2) From c360cd228acdfa335c472f91d12b0b6c0884f5a3 Mon Sep 17 00:00:00 2001 From: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com> Date: Sat, 7 Feb 2026 14:15:44 -0800 Subject: [PATCH 12/17] clean up multi-stream Signed-off-by: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com> --- .../transform/library/multi_stream_moe.py | 180 ++++-------------- .../_torch/auto_deploy/utils/_graph.py | 94 +++++++++ .../utils/test_create_derived_custom_op.py | 163 ++++++++++++++++ 3 files changed, 293 insertions(+), 144 deletions(-) create mode 100644 tests/unittest/_torch/auto_deploy/unit/singlegpu/utils/test_create_derived_custom_op.py diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.py b/tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.py index ea1407dc681d..3a5540fa059c 100644 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.py +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.py @@ -6,10 +6,9 @@ import torch from torch.fx import GraphModule -from tensorrt_llm._torch.utils import ActivationType - from ...models.factory import ModelFactory from ...shim.interface import CachedSequenceInterface +from ...utils._graph import create_derived_custom_op from ...utils.logger import ad_logger from ...utils.node_utils import is_op from ..interface import BaseTransform, SharedConfig, TransformInfo, TransformRegistry @@ -123,143 +122,29 @@ def aux_stream_wrapper( return output -# trtllm bf16 -@torch.library.custom_op("auto_deploy::trtllm_moe_fused_aux", mutates_args=()) -def trtllm_moe_fused_aux( - x: torch.Tensor, - selected_experts: torch.Tensor, - routing_weights: torch.Tensor, - w3_w1_stacked_weight: torch.Tensor, - w2_stacked_weight: torch.Tensor, - is_gated_mlp: bool = True, - act_fn: int = int(ActivationType.Silu), -) -> torch.Tensor: - device = torch.cuda.current_device() - with torch.cuda.stream( - cuda_stream_manager.get_stream(device, cuda_stream_manager.AUX_STREAM_NAME) - ): - torch.ops.auto_deploy.wait_event(device, cuda_stream_manager.MAIN_STREAM_NAME) - output = torch.ops.auto_deploy.trtllm_moe_fused( - x, - selected_experts, - routing_weights, - w3_w1_stacked_weight, - w2_stacked_weight, - is_gated_mlp, - act_fn, - ) - torch.ops.auto_deploy.record_event(device, cuda_stream_manager.AUX_STREAM_NAME) - torch.ops.auto_deploy.wait_event(device, cuda_stream_manager.AUX_STREAM_NAME) - return output - - -@trtllm_moe_fused_aux.register_fake -def trtllm_moe_fused_aux_fake( - x: torch.Tensor, - selected_experts: torch.Tensor, - routing_weights: torch.Tensor, - w3_w1_stacked_weight: torch.Tensor, - w2_stacked_weight: torch.Tensor, - is_gated_mlp: bool = True, - act_fn: int = int(ActivationType.Silu), -) -> torch.Tensor: - return torch.empty_like(x) - +def _make_aux_stream_impl(base_overload: Callable) -> Callable: + """Build an implementation that runs *base_overload* on the auxiliary CUDA stream.""" -# triton bf16 -@torch.library.custom_op("auto_deploy::triton_moe_fused_aux", mutates_args=()) -def triton_moe_fused_aux( - x: torch.Tensor, - selected_experts: torch.Tensor, - routing_weights: torch.Tensor, - w1_stacked_weight: torch.Tensor, - w2_stacked_weight: torch.Tensor, -) -> torch.Tensor: - device = torch.cuda.current_device() - with torch.cuda.stream( - cuda_stream_manager.get_stream(device, cuda_stream_manager.AUX_STREAM_NAME) - ): - torch.ops.auto_deploy.wait_event(device, cuda_stream_manager.MAIN_STREAM_NAME) - output = torch.ops.auto_deploy.triton_moe_fused( - x, - selected_experts, - routing_weights, - w1_stacked_weight, - w2_stacked_weight, - ) - torch.ops.auto_deploy.record_event(device, cuda_stream_manager.AUX_STREAM_NAME) - torch.ops.auto_deploy.wait_event(device, cuda_stream_manager.AUX_STREAM_NAME) - return output + def _impl(*args, **kwargs): + device = torch.cuda.current_device() + with torch.cuda.stream( + cuda_stream_manager.get_stream(device, cuda_stream_manager.AUX_STREAM_NAME) + ): + torch.ops.auto_deploy.wait_event(device, cuda_stream_manager.MAIN_STREAM_NAME) + output = base_overload(*args, **kwargs) + torch.ops.auto_deploy.record_event(device, cuda_stream_manager.AUX_STREAM_NAME) + torch.ops.auto_deploy.wait_event(device, cuda_stream_manager.AUX_STREAM_NAME) + return output + return _impl -@triton_moe_fused_aux.register_fake -def triton_moe_fused_aux_fake( - x: torch.Tensor, - selected_experts: torch.Tensor, - routing_weights: torch.Tensor, - w1_stacked_weight: torch.Tensor, - w2_stacked_weight: torch.Tensor, -) -> torch.Tensor: - return torch.empty_like(x) +def _create_aux_op(base_op: Callable) -> Callable: + """Create an ``_aux`` variant of a MoE custom op that runs on the auxiliary CUDA stream.""" + return create_derived_custom_op(base_op, "_aux", _make_aux_stream_impl) -@torch.library.custom_op("auto_deploy::trtllm_quant_fp8_moe_fused_aux", mutates_args=()) -def trtllm_quant_fp8_moe_fused_aux( - x: torch.Tensor, - selected_experts: torch.Tensor, - routing_weights: torch.Tensor, - fc1_expert_weights: torch.Tensor, - fc2_expert_weights: torch.Tensor, - fc1_act_scale: torch.Tensor, - fc1_dequant_scale: torch.Tensor, - fc2_act_scale_reciprocal: torch.Tensor, - fc2_dequant_scale: torch.Tensor, - is_gated_mlp: bool = True, - act_fn: int = int(ActivationType.Silu), -) -> torch.Tensor: - device = torch.cuda.current_device() - with torch.cuda.stream( - cuda_stream_manager.get_stream(device, cuda_stream_manager.AUX_STREAM_NAME) - ): - torch.ops.auto_deploy.wait_event(device, cuda_stream_manager.MAIN_STREAM_NAME) - output = torch.ops.auto_deploy.trtllm_quant_fp8_moe_fused( - x, - selected_experts, - routing_weights, - fc1_expert_weights, - fc2_expert_weights, - fc1_act_scale, - fc1_dequant_scale, - fc2_act_scale_reciprocal, - fc2_dequant_scale, - is_gated_mlp, - act_fn, - ) - torch.ops.auto_deploy.record_event(device, cuda_stream_manager.AUX_STREAM_NAME) - torch.ops.auto_deploy.wait_event(device, cuda_stream_manager.AUX_STREAM_NAME) - return output - -@trtllm_quant_fp8_moe_fused_aux.register_fake -def trtllm_quant_fp8_moe_fused_aux_fake( - x: torch.Tensor, - selected_experts: torch.Tensor, - routing_weights: torch.Tensor, - fc1_expert_weights: torch.Tensor, - fc2_expert_weights: torch.Tensor, - fc1_act_scale: torch.Tensor, - fc1_dequant_scale: torch.Tensor, - fc2_act_scale_reciprocal: torch.Tensor, - fc2_dequant_scale: torch.Tensor, - is_gated_mlp: bool = True, - act_fn: int = int(ActivationType.Silu), -) -> torch.Tensor: - return torch.empty_like(x) - - -def _execute_op_in_aux_stream( - gm: GraphModule, op_dict: Dict[Callable, Callable] -) -> Tuple[GraphModule, int]: +def _execute_op_in_aux_stream(gm: GraphModule, base_ops: List[Callable]) -> Tuple[GraphModule, int]: """Replace MoE ops with aux-stream variants and insert CUDA event synchronisation. For each MoE op we: @@ -273,12 +158,21 @@ def _execute_op_in_aux_stream( stream) to execute concurrently. 3. Replace the original MoE op with its ``_aux`` variant that runs on the auxiliary CUDA stream. + + Aux-stream variants are created lazily — only for base ops that actually + appear in the graph. """ graph = gm.graph num_replaced = 0 # Collect targets first to avoid mutating while iterating - target_nodes = [n for n in graph.nodes if is_op(n, op_dict.keys())] + target_nodes = [n for n in graph.nodes if is_op(n, base_ops)] + if not target_nodes: + return gm, 0 + + # Create aux ops only for ops that are present in the graph. + ops_in_graph = {n.target for n in target_nodes} + op_dict = {op: _create_aux_op(op) for op in ops_in_graph} # Topological position map — used to find the latest routing input. node_order = {node: i for i, node in enumerate(graph.nodes)} @@ -332,16 +226,16 @@ def _apply( factory: ModelFactory, shared_config: SharedConfig, ) -> Tuple[GraphModule, TransformInfo]: - op_dict = { - torch.ops.auto_deploy.trtllm_moe_fused: torch.ops.auto_deploy.trtllm_moe_fused_aux, - torch.ops.auto_deploy.triton_moe_fused: torch.ops.auto_deploy.triton_moe_fused_aux, - torch.ops.auto_deploy.trtllm_quant_fp8_moe_fused: torch.ops.auto_deploy.trtllm_quant_fp8_moe_fused_aux, - } - with open("before_multi_stream.txt", "w") as f: - f.write(str(gm.graph)) + base_ops = [ + torch.ops.auto_deploy.trtllm_moe_fused, + torch.ops.auto_deploy.triton_moe_fused, + torch.ops.auto_deploy.trtllm_quant_fp8_moe_fused, + torch.ops.auto_deploy.trtllm_quant_fp4_moe_fused, + ] + # Ensure that aux stream and events for the current device are added to the CudaStreamManager. cuda_stream_manager.add_device(torch.cuda.current_device()) - gm, num_matches = _execute_op_in_aux_stream(gm, op_dict) + gm, num_matches = _execute_op_in_aux_stream(gm, base_ops) info = TransformInfo( skipped=False, @@ -349,6 +243,4 @@ def _apply( is_clean=num_matches == 0, has_valid_shapes=num_matches == 0, ) - with open("after_multi_stream.txt", "w") as f: - f.write(str(gm.graph)) return gm, info diff --git a/tensorrt_llm/_torch/auto_deploy/utils/_graph.py b/tensorrt_llm/_torch/auto_deploy/utils/_graph.py index 749b2cd130af..d54c86d0c37d 100644 --- a/tensorrt_llm/_torch/auto_deploy/utils/_graph.py +++ b/tensorrt_llm/_torch/auto_deploy/utils/_graph.py @@ -20,6 +20,100 @@ from .logger import ad_logger from .node_utils import get_weight_tensor, is_op +# --------------------------------------------------------------------------- +# Dynamic custom-op derivation helpers +# --------------------------------------------------------------------------- +# These are used to create new custom ops that share the schema of an existing +# op but wrap it with additional logic (e.g. stream management). A single +# module-level dict of ``Library`` objects (keyed by namespace) is used so that +# the registrations persist and are visible via ``torch.ops..*``. +# --------------------------------------------------------------------------- + +_derived_op_libs: Dict[str, torch.library.Library] = {} +_derived_op_registry: Dict[str, Callable] = {} + + +def _get_lib(namespace: str) -> torch.library.Library: + """Return (and lazily create) a ``FRAGMENT`` Library for *namespace*.""" + if namespace not in _derived_op_libs: + _derived_op_libs[namespace] = torch.library.Library(namespace, "FRAGMENT") + return _derived_op_libs[namespace] + + +def create_derived_custom_op( + base_op: Callable, + suffix: str, + make_impl: Callable[[Callable], Callable], + make_fake: Optional[Callable[[Callable], Callable]] = None, +) -> Callable: + """Dynamically create a new custom op derived from an existing one. + + The new op has the **same** schema (arguments, default values, and return + type) as *base_op* but with a different name (````) and a + custom implementation produced by *make_impl*. + + Args: + base_op: The base custom op — either an ``OpOverloadPacket`` + (e.g. ``torch.ops.auto_deploy.trtllm_moe_fused``) or an + ``OpOverload`` (e.g. ``…trtllm_moe_fused.default``). + suffix: Suffix appended to the base op name to form the new op name + (e.g. ``"_aux"``). + make_impl: A factory ``(base_overload) -> impl_fn`` that receives the + resolved base ``OpOverload`` and returns the *implementation* + function for the new op. ``impl_fn`` will be called with the + same positional/keyword arguments as *base_op*. + make_fake: Optional factory ``(base_overload) -> fake_fn`` that returns + the *Meta / fake-tensor* implementation. When ``None`` the + default fake implementation ``torch.empty_like(args[0])`` is used. + + Returns: + The newly registered op as an ``OpOverloadPacket`` + (e.g. ``torch.ops.auto_deploy.``). Repeated calls with + the same *base_op* and *suffix* return the cached op. + """ + base_overload = base_op.default if hasattr(base_op, "default") else base_op + schema = base_overload._schema + + # e.g. "auto_deploy::trtllm_moe_fused" + qualified_name = schema.name + namespace, base_name = qualified_name.split("::") + new_name = f"{base_name}{suffix}" + new_qualified = f"{namespace}::{new_name}" + + # Return the cached op if it was already created. + if new_qualified in _derived_op_registry: + return _derived_op_registry[new_qualified] + + # Build the schema string for the derived op. ``str(schema)`` produces a + # fully-qualified string such as + # auto_deploy::trtllm_moe_fused(Tensor x, …) -> Tensor + # We replace the qualified name with the bare new name (the Library already + # knows its namespace). + new_schema_str = str(schema).replace(qualified_name, new_name, 1) + + lib = _get_lib(namespace) + lib.define(new_schema_str) + + # Register the real implementation for all devices. + # We use "CompositeExplicitAutograd" so that we can provide a separate + # Meta / fake kernel for shape inference. + lib.impl(new_name, make_impl(base_overload), "CompositeExplicitAutograd") + + # Register the Meta / fake implementation. + if make_fake is not None: + lib.impl(new_name, make_fake(base_overload), "Meta") + else: + + def _default_fake(*args, **kwargs): + return torch.empty_like(args[0]) + + lib.impl(new_name, _default_fake, "Meta") + + new_op = getattr(getattr(torch.ops, namespace), new_name) + _derived_op_registry[new_qualified] = new_op + return new_op + + _NoValType = type("_NoValType", (), {}) diff --git a/tests/unittest/_torch/auto_deploy/unit/singlegpu/utils/test_create_derived_custom_op.py b/tests/unittest/_torch/auto_deploy/unit/singlegpu/utils/test_create_derived_custom_op.py new file mode 100644 index 000000000000..14eed45a63f9 --- /dev/null +++ b/tests/unittest/_torch/auto_deploy/unit/singlegpu/utils/test_create_derived_custom_op.py @@ -0,0 +1,163 @@ +"""Tests for ``create_derived_custom_op`` in ``_graph.py``.""" + +import torch +from torch._subclasses import FakeTensorMode + +from tensorrt_llm._torch.auto_deploy.utils._graph import create_derived_custom_op + +# --------------------------------------------------------------------------- +# Helpers – tiny custom ops used as base ops for the tests +# --------------------------------------------------------------------------- + + +@torch.library.custom_op("ad_test_derived::double", mutates_args=()) +def _double(x: torch.Tensor) -> torch.Tensor: + return x * 2 + + +@_double.register_fake +def _double_fake(x: torch.Tensor) -> torch.Tensor: + return torch.empty_like(x) + + +@torch.library.custom_op("ad_test_derived::weighted_add", mutates_args=()) +def _weighted_add(x: torch.Tensor, y: torch.Tensor, alpha: float = 1.0) -> torch.Tensor: + return x + alpha * y + + +@_weighted_add.register_fake +def _weighted_add_fake(x: torch.Tensor, y: torch.Tensor, alpha: float = 1.0) -> torch.Tensor: + return torch.empty_like(x) + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +class TestCreateDerivedCustomOp: + """Tests for the ``create_derived_custom_op`` utility.""" + + def test_basic_derived_op(self): + """A derived op should be callable and produce correct results.""" + + def make_impl(base_overload): + # Wrapper that calls the base op then negates the result. + def impl(*args, **kwargs): + return -base_overload(*args, **kwargs) + + return impl + + base_op = torch.ops.ad_test_derived.double + derived = create_derived_custom_op(base_op, "_neg", make_impl) + + x = torch.tensor([1.0, 2.0, 3.0]) + result = derived(x) + expected = -(x * 2) + torch.testing.assert_close(result, expected) + + def test_derived_op_is_registered(self): + """The derived op must be accessible via ``torch.ops``.""" + + def make_impl(base_overload): + return lambda *a, **kw: base_overload(*a, **kw) + + create_derived_custom_op(torch.ops.ad_test_derived.double, "_registered", make_impl) + assert hasattr(torch.ops.ad_test_derived, "double_registered") + + def test_caching(self): + """Repeated calls with the same base_op + suffix must return the same object.""" + + def make_impl(base_overload): + return lambda *a, **kw: base_overload(*a, **kw) + + op1 = create_derived_custom_op(torch.ops.ad_test_derived.double, "_cached", make_impl) + op2 = create_derived_custom_op(torch.ops.ad_test_derived.double, "_cached", make_impl) + assert op1 is op2 + + def test_different_suffix_produces_different_op(self): + """Different suffixes must create distinct ops.""" + + def make_impl(base_overload): + return lambda *a, **kw: base_overload(*a, **kw) + + op_a = create_derived_custom_op(torch.ops.ad_test_derived.double, "_sfx_a", make_impl) + op_b = create_derived_custom_op(torch.ops.ad_test_derived.double, "_sfx_b", make_impl) + assert op_a is not op_b + + def test_default_fake_implementation(self): + """When *make_fake* is None the default (empty_like) must be used.""" + + def make_impl(base_overload): + return lambda *a, **kw: base_overload(*a, **kw) + + derived = create_derived_custom_op( + torch.ops.ad_test_derived.double, "_dflt_fake", make_impl + ) + # Calling the Meta implementation via FakeTensorMode + with FakeTensorMode(): + x = torch.empty(4) + out = derived(x) + assert out.shape == x.shape + + def test_custom_fake_implementation(self): + """A user-supplied *make_fake* must override the default.""" + + def make_impl(base_overload): + return lambda *a, **kw: base_overload(*a, **kw) + + # Fake that always returns shape (1,) regardless of input shape. + def make_fake(base_overload): + def fake(*args, **kwargs): + return args[0].new_empty(1) + + return fake + + derived = create_derived_custom_op( + torch.ops.ad_test_derived.double, + "_custom_fake", + make_impl, + make_fake=make_fake, + ) + + with FakeTensorMode(): + x = torch.empty(10) + out = derived(x) + assert out.shape == (1,) + + def test_preserves_schema_with_defaults(self): + """Derived op must preserve the base op's argument defaults.""" + + def make_impl(base_overload): + def impl(*args, **kwargs): + return base_overload(*args, **kwargs) * 10 + + return impl + + base_op = torch.ops.ad_test_derived.weighted_add + derived = create_derived_custom_op(base_op, "_x10", make_impl) + + x = torch.ones(3) + y = torch.ones(3) * 2.0 + + # With default alpha=1.0 → (x + 1.0*y) * 10 = 30 + result_default = derived(x, y) + torch.testing.assert_close(result_default, torch.full((3,), 30.0)) + + # With explicit alpha=0.5 → (x + 0.5*y) * 10 = 20 + result_alpha = derived(x, y, alpha=0.5) + torch.testing.assert_close(result_alpha, torch.full((3,), 20.0)) + + def test_accepts_op_overload(self): + """The function should accept an OpOverload (e.g. ``.default``) as well.""" + + def make_impl(base_overload): + return lambda *a, **kw: base_overload(*a, **kw) + 1 + + derived = create_derived_custom_op( + torch.ops.ad_test_derived.double.default, "_from_overload", make_impl + ) + + x = torch.tensor([5.0]) + # double → 10, +1 → 11 + torch.testing.assert_close(derived(x), torch.tensor([11.0])) From bb38432b68b0e90f840d38cfac6c75310094e3fd Mon Sep 17 00:00:00 2001 From: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com> Date: Sat, 7 Feb 2026 14:26:30 -0800 Subject: [PATCH 13/17] fix fp4 typo Signed-off-by: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com> --- .../_torch/auto_deploy/transform/library/multi_stream_moe.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.py b/tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.py index 3a5540fa059c..db4b9c09b8e4 100644 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.py +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_moe.py @@ -230,7 +230,7 @@ def _apply( torch.ops.auto_deploy.trtllm_moe_fused, torch.ops.auto_deploy.triton_moe_fused, torch.ops.auto_deploy.trtllm_quant_fp8_moe_fused, - torch.ops.auto_deploy.trtllm_quant_fp4_moe_fused, + torch.ops.auto_deploy.trtllm_quant_nvfp4_moe_fused, ] # Ensure that aux stream and events for the current device are added to the CudaStreamManager. @@ -243,4 +243,5 @@ def _apply( is_clean=num_matches == 0, has_valid_shapes=num_matches == 0, ) + return gm, info From 5b5c78414dda5b7eb18b4e855af9f2c17be7e902 Mon Sep 17 00:00:00 2001 From: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com> Date: Sat, 7 Feb 2026 15:53:59 -0800 Subject: [PATCH 14/17] fuse nvf4 swiglu, multi-stream mla attn Signed-off-by: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com> --- examples/auto_deploy/glm_flash.yaml | 3 + .../_torch/auto_deploy/config/default.yaml | 17 +- .../auto_deploy/custom_ops/linear/swiglu.py | 173 ++++++++ .../transform/library/fuse_swiglu.py | 298 ++++++++++++++ .../transform/library/multi_stream_attn.py | 219 ++++++++++ .../custom_ops/test_multi_stream_attn.py | 230 +++++++++++ .../library/test_nvfp4_swiglu.py | 377 ++++++++++++++++++ 7 files changed, 1315 insertions(+), 2 deletions(-) create mode 100644 tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_attn.py create mode 100644 tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_multi_stream_attn.py create mode 100644 tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_nvfp4_swiglu.py diff --git a/examples/auto_deploy/glm_flash.yaml b/examples/auto_deploy/glm_flash.yaml index e20e7cbf1c86..87482e47423e 100644 --- a/examples/auto_deploy/glm_flash.yaml +++ b/examples/auto_deploy/glm_flash.yaml @@ -18,3 +18,6 @@ transforms: multi_stream_moe: stage: compile enabled: true + multi_stream_mla_attn: + stage: compile + enabled: true diff --git a/tensorrt_llm/_torch/auto_deploy/config/default.yaml b/tensorrt_llm/_torch/auto_deploy/config/default.yaml index f0b188ce55bb..c4e92fecd18c 100644 --- a/tensorrt_llm/_torch/auto_deploy/config/default.yaml +++ b/tensorrt_llm/_torch/auto_deploy/config/default.yaml @@ -48,10 +48,9 @@ transforms: match_rope_layout: stage: pattern_matcher expected_layout: bsnd - match_swiglu_pattern: - stage: pattern_matcher match_rmsnorm_pattern: stage: pattern_matcher + run_shape_prop: true match_l2norm_pattern: stage: pattern_matcher ############################################################################################ @@ -77,6 +76,15 @@ transforms: stage: pattern_matcher quantize_nvfp4_from_graph: stage: pattern_matcher + # SwiGLU pattern matching must run AFTER quantization transforms. For pre-quantized + # checkpoints (e.g., NVFP4), quantization converts torch_linear_simple ops to quantized + # ops first, and then match_nvfp4_swiglu_pattern captures the NVFP4 SwiGLU pattern. + # For non-quantized models, quantization transforms are no-ops, so match_swiglu_pattern + # proceeds normally. + match_swiglu_pattern: + stage: pattern_matcher + match_nvfp4_swiglu_pattern: + stage: pattern_matcher quantize_fp8_moe: stage: pattern_matcher quantize_nvfp4_moe: @@ -128,6 +136,8 @@ transforms: fuse_nvfp4_linear: stage: post_load_fusion backend: trtllm + fuse_nvfp4_swiglu: + stage: post_load_fusion fuse_moe: stage: post_load_fusion expect_mem_change: true @@ -203,6 +213,9 @@ transforms: multi_stream_moe: stage: compile enabled: false + multi_stream_mla_attn: + stage: compile + enabled: false compile_model: stage: compile expect_mem_change: true diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/linear/swiglu.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/linear/swiglu.py index 808291d3f588..306744d588ec 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/linear/swiglu.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/linear/swiglu.py @@ -126,3 +126,176 @@ def _( # Output shape is [..., hidden_size] where hidden_size = down_weight.shape[0] output_shape = list(input.shape[:-1]) + [down_weight.shape[0]] return input.new_empty(output_shape, dtype=input.dtype) + + +# ── NVFP4 quantized SwiGLU ops ────────────────────────────────────────────── + + +@torch.library.custom_op("auto_deploy::torch_nvfp4_swiglu_mlp", mutates_args=()) +def torch_nvfp4_swiglu_mlp( + input: torch.Tensor, + gate_weight: torch.Tensor, + up_weight: torch.Tensor, + down_weight: torch.Tensor, + gate_input_scale: torch.Tensor, + gate_weight_scale: torch.Tensor, + gate_alpha: torch.Tensor, + up_input_scale: torch.Tensor, + up_weight_scale: torch.Tensor, + up_alpha: torch.Tensor, + down_input_scale: torch.Tensor, + down_weight_scale: torch.Tensor, + down_alpha: torch.Tensor, +) -> torch.Tensor: + """NVFP4 quantized SwiGLU MLP operation (intermediate representation). + + Computes: silu(nvfp4_linear(x, gate)) * nvfp4_linear(x, up) -> nvfp4_linear(down) + + This is the intermediate representation used after pattern matching for NVFP4 + quantized checkpoints, before gate+up weight fusion is applied. + + Args: + input: Input tensor of shape [..., hidden_size]. + gate_weight: FP4 packed gate weight [intermediate_size, hidden_size/2] uint8. + up_weight: FP4 packed up weight [intermediate_size, hidden_size/2] uint8. + down_weight: FP4 packed down weight [hidden_size, intermediate_size/2] uint8. + gate_input_scale: Input scale for gate projection. + gate_weight_scale: Per-block weight scale for gate projection. + gate_alpha: Alpha (combined scale) for gate projection. + up_input_scale: Input scale for up projection. + up_weight_scale: Per-block weight scale for up projection. + up_alpha: Alpha (combined scale) for up projection. + down_input_scale: Input scale for down projection. + down_weight_scale: Per-block weight scale for down projection. + down_alpha: Alpha (combined scale) for down projection. + + Returns: + Output tensor of shape [..., hidden_size]. + """ + gate_out = torch.ops.auto_deploy.torch_fake_quant_nvfp4_linear( + input, + gate_weight, + None, + input_scale=[gate_input_scale], + weight_scale=[gate_weight_scale, gate_alpha], + input_zp=[], + weight_zp=[], + ) + up_out = torch.ops.auto_deploy.torch_fake_quant_nvfp4_linear( + input, + up_weight, + None, + input_scale=[up_input_scale], + weight_scale=[up_weight_scale, up_alpha], + input_zp=[], + weight_zp=[], + ) + hidden = F.silu(gate_out) * up_out + return torch.ops.auto_deploy.torch_fake_quant_nvfp4_linear( + hidden, + down_weight, + None, + input_scale=[down_input_scale], + weight_scale=[down_weight_scale, down_alpha], + input_zp=[], + weight_zp=[], + ) + + +@torch_nvfp4_swiglu_mlp.register_fake +def _( + input: torch.Tensor, + gate_weight: torch.Tensor, + up_weight: torch.Tensor, + down_weight: torch.Tensor, + gate_input_scale: torch.Tensor, + gate_weight_scale: torch.Tensor, + gate_alpha: torch.Tensor, + up_input_scale: torch.Tensor, + up_weight_scale: torch.Tensor, + up_alpha: torch.Tensor, + down_input_scale: torch.Tensor, + down_weight_scale: torch.Tensor, + down_alpha: torch.Tensor, +) -> torch.Tensor: + """Fake implementation for tracing.""" + # Output shape: [..., hidden_size] where hidden_size = down_weight.shape[0] + output_shape = list(input.shape[:-1]) + [down_weight.shape[0]] + return input.new_empty(output_shape, dtype=input.dtype) + + +@torch.library.custom_op("auto_deploy::fused_nvfp4_swiglu_mlp", mutates_args=()) +def fused_nvfp4_swiglu_mlp( + input: torch.Tensor, + gate_up_weight: torch.Tensor, + down_weight: torch.Tensor, + gate_up_input_scale: torch.Tensor, + gate_up_weight_scale: torch.Tensor, + gate_up_alpha: torch.Tensor, + down_input_scale: torch.Tensor, + down_weight_scale: torch.Tensor, + down_alpha: torch.Tensor, +) -> torch.Tensor: + """Fused NVFP4 SwiGLU MLP with concatenated gate+up weights. + + Performs a single NVFP4 matmul for gate and up projections, then splits, + applies SwiGLU activation, and does the down NVFP4 matmul. + + Args: + input: Input tensor of shape [..., hidden_size]. + gate_up_weight: Concatenated FP4 packed gate+up weight + [2*intermediate_size, hidden_size/2] uint8. + down_weight: FP4 packed down weight [hidden_size, intermediate_size/2] uint8. + gate_up_input_scale: Shared input scale for gate+up projection. + gate_up_weight_scale: Concatenated per-block weight scale for gate+up. + gate_up_alpha: Shared alpha for gate+up projection. + down_input_scale: Input scale for down projection. + down_weight_scale: Per-block weight scale for down projection. + down_alpha: Alpha for down projection. + + Returns: + Output tensor of shape [..., hidden_size]. + """ + # Single NVFP4 linear for both gate and up projections + gate_up_out = torch.ops.auto_deploy.torch_quant_nvfp4_linear( + input, + gate_up_weight, + bias=None, + input_scale=gate_up_input_scale, + weight_scale=gate_up_weight_scale, + alpha=gate_up_alpha, + ) + + # Split into gate and up outputs + gate_out, up_out = gate_up_out.chunk(2, dim=-1) + + # Apply SwiGLU activation: silu(gate) * up + hidden = F.silu(gate_out) * up_out + + # Down projection + return torch.ops.auto_deploy.torch_quant_nvfp4_linear( + hidden, + down_weight, + bias=None, + input_scale=down_input_scale, + weight_scale=down_weight_scale, + alpha=down_alpha, + ) + + +@fused_nvfp4_swiglu_mlp.register_fake +def _( + input: torch.Tensor, + gate_up_weight: torch.Tensor, + down_weight: torch.Tensor, + gate_up_input_scale: torch.Tensor, + gate_up_weight_scale: torch.Tensor, + gate_up_alpha: torch.Tensor, + down_input_scale: torch.Tensor, + down_weight_scale: torch.Tensor, + down_alpha: torch.Tensor, +) -> torch.Tensor: + """Fake implementation for tracing.""" + # Output shape: [..., hidden_size] where hidden_size = down_weight.shape[0] + output_shape = list(input.shape[:-1]) + [down_weight.shape[0]] + return input.new_empty(output_shape, dtype=input.dtype) diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_swiglu.py b/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_swiglu.py index 7c8ef59522fe..4d97475cfc80 100644 --- a/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_swiglu.py +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/fuse_swiglu.py @@ -282,3 +282,301 @@ def _apply( ) return gm, info + + +# ── NVFP4 quantized SwiGLU pattern matching and fusion ────────────────────── + +from ...custom_ops.linear.swiglu import torch_nvfp4_swiglu_mlp # noqa: E402 + + +def _nvfp4_swiglu_pattern_no_bias( + x, + gate_weight, + gate_input_scale, + gate_weight_scale, + gate_alpha, + up_weight, + up_input_scale, + up_weight_scale, + up_alpha, + down_weight, + down_input_scale, + down_weight_scale, + down_alpha, +): + """Pattern for NVFP4 quantized SwiGLU MLP without biases. + + Matches: silu(nvfp4_linear(x, gate)) * nvfp4_linear(x, up) -> nvfp4_linear(down) + """ + gate_out = torch.ops.auto_deploy.torch_fake_quant_nvfp4_linear.default( + x, + gate_weight, + None, + input_scale=[gate_input_scale], + weight_scale=[gate_weight_scale, gate_alpha], + input_zp=[], + weight_zp=[], + ) + up_out = torch.ops.auto_deploy.torch_fake_quant_nvfp4_linear.default( + x, + up_weight, + None, + input_scale=[up_input_scale], + weight_scale=[up_weight_scale, up_alpha], + input_zp=[], + weight_zp=[], + ) + silu_out = torch.ops.aten.silu.default(gate_out) + mul_out = torch.ops.aten.mul.Tensor(silu_out, up_out) + down_out = torch.ops.auto_deploy.torch_fake_quant_nvfp4_linear.default( + mul_out, + down_weight, + None, + input_scale=[down_input_scale], + weight_scale=[down_weight_scale, down_alpha], + input_zp=[], + weight_zp=[], + ) + return down_out + + +def _nvfp4_swiglu_replacement_no_bias( + x, + gate_weight, + gate_input_scale, + gate_weight_scale, + gate_alpha, + up_weight, + up_input_scale, + up_weight_scale, + up_alpha, + down_weight, + down_input_scale, + down_weight_scale, + down_alpha, +): + """Replacement for NVFP4 quantized SwiGLU pattern without biases.""" + return torch_nvfp4_swiglu_mlp( + x, + gate_weight, + up_weight, + down_weight, + gate_input_scale, + gate_weight_scale, + gate_alpha, + up_input_scale, + up_weight_scale, + up_alpha, + down_input_scale, + down_weight_scale, + down_alpha, + ) + + +@TransformRegistry.register("match_nvfp4_swiglu_pattern") +class MatchNVFP4SwiGLUPattern(BaseTransform): + """Matches NVFP4 quantized SwiGLU MLP patterns and replaces with torch_nvfp4_swiglu_mlp. + + This transform runs in the pattern_matcher stage AFTER quantize_nvfp4_linear_from_config + has converted torch_linear_simple ops to torch_fake_quant_nvfp4_linear ops. + + It detects the following NVFP4 pattern: + silu(nvfp4_linear(x, gate)) * nvfp4_linear(x, up) -> nvfp4_linear(down) + + And replaces it with a single torch_nvfp4_swiglu_mlp op that can be fused later. + """ + + config: TransformConfig + + @classmethod + def get_config_class(cls) -> Type[TransformConfig]: + return TransformConfig + + def _apply( + self, + gm: GraphModule, + cm: CachedSequenceInterface, + factory: ModelFactory, + shared_config: SharedConfig, + ) -> Tuple[GraphModule, TransformInfo]: + patterns = ADPatternMatcherPass() + + # FP4 shape params for dummy args (shapes don't matter for matching) + N = 32 # intermediate_size + K_packed = 32 # hidden_size / 2 (FP4 packing) + K_eff = 2 * K_packed # actual hidden_size + N_down = K_eff # hidden_size (output of down proj) + K_down_packed = N // 2 # intermediate_size / 2 (down proj input) + + # Weight scale sizes (per-block scale: N * K / 16) + gate_cutlass_len = N * (K_eff // 16) + down_cutlass_len = N_down * (N // 16) + + x = torch.randn(2, K_eff, device="meta", dtype=torch.float16) + + # Gate args + gate_w = torch.randint(0, 255, (N, K_packed), device="meta", dtype=torch.uint8) + gate_is = torch.tensor(0.01, device="meta", dtype=torch.float32) + gate_ws = torch.randint(0, 255, (gate_cutlass_len,), device="meta", dtype=torch.uint8) + gate_a = torch.tensor(1.2345, device="meta", dtype=torch.float32) + + # Up args (same shapes as gate) + up_w = torch.randint(0, 255, (N, K_packed), device="meta", dtype=torch.uint8) + up_is = torch.tensor(0.02, device="meta", dtype=torch.float32) + up_ws = torch.randint(0, 255, (gate_cutlass_len,), device="meta", dtype=torch.uint8) + up_a = torch.tensor(2.3456, device="meta", dtype=torch.float32) + + # Down args + down_w = torch.randint(0, 255, (N_down, K_down_packed), device="meta", dtype=torch.uint8) + down_is = torch.tensor(0.03, device="meta", dtype=torch.float32) + down_ws = torch.randint(0, 255, (down_cutlass_len,), device="meta", dtype=torch.uint8) + down_a = torch.tensor(3.4567, device="meta", dtype=torch.float32) + + dummy_args = [ + x, + gate_w, + gate_is, + gate_ws, + gate_a, + up_w, + up_is, + up_ws, + up_a, + down_w, + down_is, + down_ws, + down_a, + ] + + register_ad_pattern( + search_fn=_nvfp4_swiglu_pattern_no_bias, + replace_fn=_nvfp4_swiglu_replacement_no_bias, + patterns=patterns, + dummy_args=dummy_args, + ) + + num_matches = patterns.apply(gm.graph) + + if num_matches > 0: + gm.recompile() + + info = TransformInfo( + skipped=False, + num_matches=num_matches, + is_clean=num_matches == 0, + has_valid_shapes=num_matches == 0, + ) + + return gm, info + + +@TransformRegistry.register("fuse_nvfp4_swiglu") +class FuseNVFP4SwiGLU(BaseTransform): + """Fuses torch_nvfp4_swiglu_mlp ops by concatenating gate and up FP4 weights. + + This transform runs in the post_load_fusion stage and replaces torch_nvfp4_swiglu_mlp + ops with fused_nvfp4_swiglu_mlp ops that use a single concatenated gate+up weight matrix. + + FP4 weight fusion: + - gate+up packed weights are concatenated along dim=0 + - gate+up per-block weight scales are concatenated along dim=0 + - gate+up input_scale and alpha must match (shared input) + """ + + config: TransformConfig + + @classmethod + def get_config_class(cls) -> Type[TransformConfig]: + return TransformConfig + + def _apply( + self, + gm: GraphModule, + cm: CachedSequenceInterface, + factory: ModelFactory, + shared_config: SharedConfig, + ) -> Tuple[GraphModule, TransformInfo]: + graph = gm.graph + cnt = 0 + fused_weight_idx = 0 + + for node in list(graph.nodes): + if not is_op(node, torch.ops.auto_deploy.torch_nvfp4_swiglu_mlp.default): + continue + + # Extract args: + # (input, gate_weight, up_weight, down_weight, + # gate_input_scale, gate_weight_scale, gate_alpha, + # up_input_scale, up_weight_scale, up_alpha, + # down_input_scale, down_weight_scale, down_alpha) + input_node = node.args[0] + gate_weight_node = node.args[1] + up_weight_node = node.args[2] + down_weight_node = node.args[3] + gate_input_scale_node = node.args[4] + gate_weight_scale_node = node.args[5] + gate_alpha_node = node.args[6] + up_input_scale_node = node.args[7] # noqa: F841 + up_weight_scale_node = node.args[8] + up_alpha_node = node.args[9] # noqa: F841 + down_input_scale_node = node.args[10] + down_weight_scale_node = node.args[11] + down_alpha_node = node.args[12] + + # Get the actual weight tensors + gate_weight = get_attr_by_name(gm, gate_weight_node.target) + up_weight = get_attr_by_name(gm, up_weight_node.target) + + # Concatenate gate and up FP4 packed weights along dim=0 + gate_up_weight = torch.cat([gate_weight, up_weight], dim=0) + + # Get and concatenate weight scales + gate_weight_scale = get_attr_by_name(gm, gate_weight_scale_node.target) + up_weight_scale = get_attr_by_name(gm, up_weight_scale_node.target) + gate_up_weight_scale = torch.cat([gate_weight_scale, up_weight_scale], dim=0) + + # Register fused buffers + prefix = f"fused_nvfp4_swiglu_{fused_weight_idx}" + gm.register_buffer(f"{prefix}_gate_up_weight", gate_up_weight) + gm.register_buffer(f"{prefix}_gate_up_weight_scale", gate_up_weight_scale) + + # Create get_attr nodes for fused weights/scales + with graph.inserting_before(node): + fused_gate_up_weight_node = graph.get_attr(f"{prefix}_gate_up_weight") + fused_gate_up_weight_scale_node = graph.get_attr(f"{prefix}_gate_up_weight_scale") + + # Create the fused_nvfp4_swiglu_mlp node + # Use gate's input_scale and alpha (same as up's since they share input) + with graph.inserting_after(node): + fused_node: Node = graph.call_function( + torch.ops.auto_deploy.fused_nvfp4_swiglu_mlp.default, + args=( + input_node, + fused_gate_up_weight_node, + down_weight_node, + gate_input_scale_node, # shared input_scale for gate+up + fused_gate_up_weight_scale_node, + gate_alpha_node, # shared alpha for gate+up + down_input_scale_node, + down_weight_scale_node, + down_alpha_node, + ), + ) + + # Replace uses and erase old node + node.replace_all_uses_with(fused_node) + graph.erase_node(node) + + fused_weight_idx += 1 + cnt += 1 + + if cnt > 0: + gm.recompile() + eliminate_dead_code(gm) + delete_all_unused_submodules(gm) + + info = TransformInfo( + skipped=False, num_matches=cnt, is_clean=cnt == 0, has_valid_shapes=cnt == 0 + ) + + return gm, info diff --git a/tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_attn.py b/tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_attn.py new file mode 100644 index 000000000000..1c531036840a --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/transform/library/multi_stream_attn.py @@ -0,0 +1,219 @@ +"""Transform for multi-stream execution of Q and KV projection chains in MLA attention. + +In DeepSeek-style MLA (Multi-head Latent Attention), the input layernorm output +forks into two independent projection chains that merge at the RoPE + attention op: + + - **Q chain** (heavier): q_a_proj -> rms_norm -> q_b_proj -> view -> split + - **KV chain** (lighter): kv_a_proj_with_mqa -> split -> rms_norm + view + +The Q chain is ~9x heavier than the KV chain. This transform moves the KV +projection linear onto the auxiliary CUDA stream so it executes concurrently +with the Q chain on the main stream. +""" + +from collections import deque +from typing import Callable, List, Tuple + +import torch +from torch.fx import GraphModule, Node + +# Reuse CachedSequenceInterface for the _apply signature. +from ...shim.interface import CachedSequenceInterface +from ...utils._graph import create_derived_custom_op +from ...utils.node_utils import is_op +from ..interface import BaseTransform, SharedConfig, TransformInfo, TransformRegistry + +# Reuse the stream infrastructure from the MoE multi-stream transform. +from .multi_stream_moe import ( + ModelFactory, + _make_aux_stream_impl, + cuda_stream_manager, + record_event_passthrough, +) + +# --------------------------------------------------------------------------- +# Supported linear op targets. Extend this list to cover quantised variants. +# --------------------------------------------------------------------------- +_LINEAR_OPS: List[Callable] = [ + torch.ops.auto_deploy.torch_linear_simple, + torch.ops.aten.linear, +] + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _is_linear(node: Node) -> bool: + """Return ``True`` if *node* is a call to one of the supported linear ops.""" + return is_op(node, _LINEAR_OPS) + + +def _has_downstream_linear(start: Node, max_depth: int = 3) -> bool: + """BFS from *start* through its users and return ``True`` if a linear op is reachable. + + The search only follows *user* edges (downstream in the data-flow graph) + and stops after *max_depth* hops. ``start`` itself is **not** checked. + """ + visited: set[Node] = {start} + queue: deque[Tuple[Node, int]] = deque() + + for user in start.users: + queue.append((user, 1)) + + while queue: + node, depth = queue.popleft() + if node in visited: + continue + visited.add(node) + + if _is_linear(node): + return True + + if depth < max_depth: + for user in node.users: + queue.append((user, depth + 1)) + + return False + + +def _find_kv_proj_linears(gm: GraphModule) -> List[Tuple[Node, Node]]: + """Find (fork_point, kv_linear) pairs suitable for aux-stream execution. + + A *fork point* is a node that directly feeds two or more supported linear + ops. Among these linears the one that does **not** lead to another linear + within a small BFS depth is the KV projection candidate (the lighter + branch). + + Returns a list of ``(fork_point, kv_linear_node)`` tuples. + """ + results: List[Tuple[Node, Node]] = [] + + for node in gm.graph.nodes: + # Collect direct linear users of this node. + linear_users = [u for u in node.users if _is_linear(u)] + if len(linear_users) < 2: + continue + + # Separate into "has downstream linear" (Q-like) and "does not" (KV-like). + kv_candidates = [ln for ln in linear_users if not _has_downstream_linear(ln)] + q_candidates = [ln for ln in linear_users if _has_downstream_linear(ln)] + + if not kv_candidates or not q_candidates: + continue + + # Pick the KV candidate(s). In MLA there is exactly one per fork point. + for kv_linear in kv_candidates: + results.append((node, kv_linear)) + + return results + + +def _create_aux_op(base_op: Callable) -> Callable: + """Create an ``_aux`` variant of a linear op that runs on the auxiliary CUDA stream. + + Uses a custom ``make_fake`` that delegates to the base op's registered fake + so that output shapes are computed correctly (linear output shape != input shape). + """ + return create_derived_custom_op( + base_op, + "_aux", + _make_aux_stream_impl, + make_fake=lambda base: lambda *a, **kw: base(*a, **kw), + ) + + +def _execute_kv_proj_in_aux_stream(gm: GraphModule) -> Tuple[GraphModule, int]: + """Replace KV projection linears with aux-stream variants. + + For each matched ``(fork_point, kv_linear)`` the rewriter: + + 1. Inserts ``record_event_passthrough(fork_point)`` so the main-stream + event is recorded *before* the Q-chain kernels are submitted. + 2. Replaces the KV linear's target with its ``_aux`` variant and wires the + ``record_event_passthrough`` output as the hidden-state input + (creating a true data dependency). + + The remaining KV-chain ops (split, rms_norm, view) stay on the main + stream — they are lightweight and run after the aux wait that is built + into the derived op. + + Aux-stream variants are created lazily — only for base ops that actually + appear in the matched KV positions. + """ + pairs = _find_kv_proj_linears(gm) + if not pairs: + return gm, 0 + + graph = gm.graph + node_order = {n: i for i, n in enumerate(graph.nodes)} + + # Create aux ops lazily for whatever linear op types are found. + ops_in_graph = {kv_linear.target for _, kv_linear in pairs} + op_dict = {op: _create_aux_op(op) for op in ops_in_graph} + + num_replaced = 0 + + for fork_point, kv_linear in pairs: + # Find the Q-chain linear(s) so we can insert the event record + # *before* the earliest Q-chain op in graph order. + q_linears = [u for u in fork_point.users if _is_linear(u) and u is not kv_linear] + earliest_q = min(q_linears, key=lambda n: node_order.get(n, 0)) + + # Insert record_event_passthrough right before the first Q-chain + # linear so the event is recorded before Q kernels hit the GPU. + with graph.inserting_before(earliest_q): + rec_node = graph.call_function( + record_event_passthrough, + args=(fork_point,), + kwargs={"device": torch.cuda.current_device()}, + ) + + # Replace KV linear with its aux-stream variant. The hidden-state + # input (args[0]) is rewired to ``rec_node`` to create a data + # dependency that ensures the event is recorded first. + new_args = tuple(rec_node if arg is fork_point else arg for arg in kv_linear.args) + + with graph.inserting_after(kv_linear): + new_node = graph.call_function( + op_dict[kv_linear.target], args=new_args, kwargs=kv_linear.kwargs + ) + + kv_linear.replace_all_uses_with(new_node) + graph.erase_node(kv_linear) + num_replaced += 1 + + return gm, num_replaced + + +# --------------------------------------------------------------------------- +# Transform class +# --------------------------------------------------------------------------- + + +@TransformRegistry.register("multi_stream_mla_attn") +class MultiStreamMLAAttn(BaseTransform): + """Multi-stream Q/KV projection parallelism for MLA attention blocks.""" + + def _apply( + self, + gm: GraphModule, + cm: CachedSequenceInterface, + factory: ModelFactory, + shared_config: SharedConfig, + ) -> Tuple[GraphModule, TransformInfo]: + # Ensure aux stream and events are set up for the current device. + cuda_stream_manager.add_device(torch.cuda.current_device()) + + gm, num_matches = _execute_kv_proj_in_aux_stream(gm) + + info = TransformInfo( + skipped=False, + num_matches=num_matches, + is_clean=num_matches == 0, + has_valid_shapes=num_matches == 0, + ) + with open("/tmp/after_multi_stream_mla_attn.txt", "w") as f: + f.write(str(gm.graph)) + return gm, info diff --git a/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_multi_stream_attn.py b/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_multi_stream_attn.py new file mode 100644 index 000000000000..474e9e6d5d3c --- /dev/null +++ b/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/test_multi_stream_attn.py @@ -0,0 +1,230 @@ +"""Tests for multi-stream Q/KV projection parallelism in MLA attention. + +The test builds a minimal mock model that mirrors the MLA fork pattern: +a shared input feeds two parallel linear chains (one heavier "Q-like", +one lighter "KV-like") whose outputs are combined with an add. + +The transform should: + 1. Detect the fork point (shared input with 2+ linear users). + 2. Identify the lighter KV-like linear (no downstream linear within + a few hops) vs. the heavier Q-like chain (has a downstream linear). + 3. Move the KV linear onto the auxiliary CUDA stream. + 4. Preserve numerical correctness. + 5. Be compatible with CUDA graph capture & replay. +""" + +import torch +import torch.nn as nn + +from tensorrt_llm._torch.auto_deploy.transform.library.multi_stream_attn import ( + _execute_kv_proj_in_aux_stream, + _find_kv_proj_linears, +) +from tensorrt_llm._torch.auto_deploy.transform.library.multi_stream_moe import cuda_stream_manager + +# --------------------------------------------------------------------------- +# Helpers -- mock MLA-like module +# --------------------------------------------------------------------------- + + +class MockMLABlock(nn.Module): + """Simplified MLA-like attention block with Q and KV projection chains. + + Q chain (heavier): q_a_proj -> relu (stand-in for rms_norm) -> q_b_proj + KV chain (lighter): kv_a_proj + Merge: add(q_b_proj_output, kv_a_proj_output) + + The layernorm at the output simulates the inter-layer distance in a real + transformer (output projection, residual add, layernorm) so that the + next layer's fork point is beyond the BFS max_depth from this layer's + KV linear. + """ + + def __init__(self, hidden_dim: int, q_inner_dim: int, kv_out_dim: int): + super().__init__() + # Q chain: two linears with a non-linearity in between + self.q_a_proj = nn.Linear(hidden_dim, q_inner_dim, bias=False) + self.q_b_proj = nn.Linear(q_inner_dim, kv_out_dim, bias=False) + # KV chain: single linear + self.kv_a_proj = nn.Linear(hidden_dim, kv_out_dim, bias=False) + # Inter-layer distance (layernorm + relu simulate residual + norm) + self.layernorm = nn.LayerNorm(kv_out_dim) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # Q chain: q_a_proj -> relu -> q_b_proj + q = self.q_a_proj(x) + q = torch.nn.functional.relu(q) + q = self.q_b_proj(q) + # KV chain: kv_a_proj + kv = self.kv_a_proj(x) + out = q + kv + # Inter-layer distance to push next layer's linears beyond BFS depth + return self.layernorm(torch.nn.functional.relu(out)) + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +def _build_gm(model, example_input): + """Export *model* to an FX GraphModule.""" + egm = torch.export.export(model, (example_input,)) + return egm.module() + + +def test_pattern_matching_single_block(): + """The pattern matcher should find exactly one pair for a single MLA block.""" + model = MockMLABlock(128, 64, 128).eval().to("cuda") + example_input = torch.randn(4, 128, device="cuda") + gm = _build_gm(model, example_input) + + pairs = _find_kv_proj_linears(gm) + assert len(pairs) == 1, f"Expected 1 fork-point pair, got {len(pairs)}" + + +def test_pattern_matching_multi_block(): + """Multiple layers with sufficient inter-layer distance should all be matched.""" + hidden_dim, q_inner_dim, kv_out_dim = 128, 64, 128 + model = ( + nn.Sequential( + MockMLABlock(hidden_dim, q_inner_dim, kv_out_dim), + MockMLABlock(kv_out_dim, q_inner_dim, kv_out_dim), + ) + .eval() + .to("cuda") + ) + example_input = torch.randn(4, hidden_dim, device="cuda") + gm = _build_gm(model, example_input) + + pairs = _find_kv_proj_linears(gm) + assert len(pairs) == 2, f"Expected 2 fork-point pairs, got {len(pairs)}" + + +def test_numerical_correctness(): + """After the transform the GraphModule must produce the same output as the original model.""" + hidden_dim, q_inner_dim, kv_out_dim = 128, 64, 128 + cuda_stream_manager.add_device(torch.cuda.current_device()) + + model = MockMLABlock(hidden_dim, q_inner_dim, kv_out_dim).eval().to("cuda") + example_input = torch.randn(4, hidden_dim, device="cuda") + gm = _build_gm(model, example_input) + + test_x = torch.randn(4, hidden_dim, device="cuda") + ref_output = model(test_x) + + gm, num_replaced = _execute_kv_proj_in_aux_stream(gm) + + assert num_replaced == 1, f"Expected 1 replacement, got {num_replaced}" + + y = gm(test_x) + assert torch.allclose(y, ref_output, atol=1e-5), ( + f"Output mismatch: max diff = {(y - ref_output).abs().max().item()}" + ) + + +def test_numerical_correctness_multi_block(): + """Multi-block correctness test.""" + hidden_dim, q_inner_dim, kv_out_dim = 128, 64, 128 + cuda_stream_manager.add_device(torch.cuda.current_device()) + + model = ( + nn.Sequential( + MockMLABlock(hidden_dim, q_inner_dim, kv_out_dim), + MockMLABlock(kv_out_dim, q_inner_dim, kv_out_dim), + ) + .eval() + .to("cuda") + ) + example_input = torch.randn(4, hidden_dim, device="cuda") + gm = _build_gm(model, example_input) + + test_x = torch.randn(4, hidden_dim, device="cuda") + ref_output = model(test_x) + + gm, num_replaced = _execute_kv_proj_in_aux_stream(gm) + + assert num_replaced == 2, f"Expected 2 replacements, got {num_replaced}" + + y = gm(test_x) + assert torch.allclose(y, ref_output, atol=1e-5), ( + f"Output mismatch: max diff = {(y - ref_output).abs().max().item()}" + ) + + +def test_cuda_graph_compatibility(): + """The transformed GraphModule must work under CUDA graph capture and replay.""" + hidden_dim, q_inner_dim, kv_out_dim = 128, 64, 128 + cuda_stream_manager.add_device(torch.cuda.current_device()) + + model = MockMLABlock(hidden_dim, q_inner_dim, kv_out_dim).eval().to("cuda") + example_input = torch.randn(4, hidden_dim, device="cuda") + gm = _build_gm(model, example_input) + + test_x = torch.randn(4, hidden_dim, device="cuda") + ref_output = model(test_x) + + gm, num_replaced = _execute_kv_proj_in_aux_stream(gm) + assert num_replaced == 1 + + # Allocate static buffers for CUDA graph capture. + static_x = torch.randn(4, hidden_dim, device="cuda") + static_output = torch.randn(4, kv_out_dim, device="cuda") + + # Warm up (required before capture). + for _ in range(3): + static_output.copy_(gm(static_x)) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + static_output.copy_(gm(static_x)) + + static_x.copy_(test_x) + graph.replay() + + assert torch.allclose(static_output, ref_output, atol=1e-5), ( + f"CUDA graph output mismatch: max diff = {(static_output - ref_output).abs().max().item()}" + ) + + +def test_no_match_on_single_linear(): + """A node with only one linear user should not be matched.""" + + class SingleLinear(nn.Module): + def __init__(self, dim): + super().__init__() + self.fc = nn.Linear(dim, dim, bias=False) + + def forward(self, x): + return self.fc(x) + + model = SingleLinear(64).eval().to("cuda") + example_input = torch.randn(4, 64, device="cuda") + gm = _build_gm(model, example_input) + + pairs = _find_kv_proj_linears(gm) + assert len(pairs) == 0, f"Expected 0 matches, got {len(pairs)}" + + +def test_no_match_when_both_have_downstream_linear(): + """When *both* branches have downstream linears the pattern should not match.""" + + class BothHeavy(nn.Module): + def __init__(self, dim, inner): + super().__init__() + self.fc_a1 = nn.Linear(dim, inner, bias=False) + self.fc_a2 = nn.Linear(inner, dim, bias=False) + self.fc_b1 = nn.Linear(dim, inner, bias=False) + self.fc_b2 = nn.Linear(inner, dim, bias=False) + + def forward(self, x): + a = self.fc_a2(torch.relu(self.fc_a1(x))) + b = self.fc_b2(torch.relu(self.fc_b1(x))) + return a + b + + model = BothHeavy(64, 32).eval().to("cuda") + example_input = torch.randn(4, 64, device="cuda") + gm = _build_gm(model, example_input) + + pairs = _find_kv_proj_linears(gm) + assert len(pairs) == 0, f"Expected 0 matches, got {len(pairs)}" diff --git a/tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_nvfp4_swiglu.py b/tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_nvfp4_swiglu.py new file mode 100644 index 000000000000..8e4e3f329163 --- /dev/null +++ b/tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_nvfp4_swiglu.py @@ -0,0 +1,377 @@ +"""Tests for NVFP4 quantized SwiGLU pattern matching and fusion transforms. + +Tests the parallel NVFP4 SwiGLU path: +1. match_nvfp4_swiglu_pattern: Matches torch_fake_quant_nvfp4_linear SwiGLU -> torch_nvfp4_swiglu_mlp +2. fuse_nvfp4_swiglu: Fuses gate+up FP4 weights -> fused_nvfp4_swiglu_mlp +""" + +import pytest +import torch +import torch.nn as nn +from _torch_test_utils import fp4_compatible, trtllm_ops_available + +import tensorrt_llm._torch.auto_deploy.custom_ops # noqa: F401 +from tensorrt_llm._torch.auto_deploy.export import torch_export_to_gm +from tensorrt_llm._torch.auto_deploy.transform.optimizer import InferenceOptimizer +from tensorrt_llm._torch.auto_deploy.utils.node_utils import is_op +from tensorrt_llm._torch.auto_deploy.utils.quantization_utils import fp4_global_scale + +_skip_reason = "Requires NVFP4 (Blackwell+) and TRT-LLM ops" +_skip_condition = not (fp4_compatible() and trtllm_ops_available()) + + +class NVFP4SwiGLUMLP(nn.Module): + """SwiGLU MLP using NVFP4 quantized linear ops. + + Mimics the graph structure produced by quantize_nvfp4_linear_from_config + applied to a standard SwiGLU MLP: silu(gate(x)) * up(x) -> down(hidden). + """ + + def __init__(self, hidden_size: int = 128, intermediate_size: int = 128): + super().__init__() + device = torch.device("cuda") + scaling_vector_size = 16 + + # Create random weights and quantize them to FP4 + gate_weight = ( + torch.randn(intermediate_size, hidden_size, dtype=torch.half, device=device) * 0.05 + ) + up_weight = ( + torch.randn(intermediate_size, hidden_size, dtype=torch.half, device=device) * 0.05 + ) + down_weight = ( + torch.randn(hidden_size, intermediate_size, dtype=torch.half, device=device) * 0.05 + ) + + # Quantize gate projection + s_w_gate = fp4_global_scale(gate_weight) + gate_fp4, gate_cutlass = torch.ops.trtllm.fp4_quantize( + gate_weight, s_w_gate, scaling_vector_size, False + ) + # Use a shared input scale for gate and up (same input x) + s_in = fp4_global_scale(torch.randn(1, hidden_size, dtype=torch.half, device=device)) + gate_alpha = (1.0 / (s_in * s_w_gate)).to(torch.float32) + + self.register_buffer("gate_weight", gate_fp4) + self.register_buffer("gate_input_scale", s_in.to(torch.float32)) + self.register_buffer("gate_weight_scale", gate_cutlass) + self.register_buffer("gate_alpha", gate_alpha) + + # Quantize up projection (same input scale as gate) + s_w_up = fp4_global_scale(up_weight) + up_fp4, up_cutlass = torch.ops.trtllm.fp4_quantize( + up_weight, s_w_up, scaling_vector_size, False + ) + up_alpha = (1.0 / (s_in * s_w_up)).to(torch.float32) + + self.register_buffer("up_weight", up_fp4) + self.register_buffer("up_input_scale", s_in.to(torch.float32)) + self.register_buffer("up_weight_scale", up_cutlass) + self.register_buffer("up_alpha", up_alpha) + + # Quantize down projection (different input: the hidden state) + s_in_down = fp4_global_scale( + torch.randn(1, intermediate_size, dtype=torch.half, device=device) + ) + s_w_down = fp4_global_scale(down_weight) + down_fp4, down_cutlass = torch.ops.trtllm.fp4_quantize( + down_weight, s_w_down, scaling_vector_size, False + ) + down_alpha = (1.0 / (s_in_down * s_w_down)).to(torch.float32) + + self.register_buffer("down_weight", down_fp4) + self.register_buffer("down_input_scale", s_in_down.to(torch.float32)) + self.register_buffer("down_weight_scale", down_cutlass) + self.register_buffer("down_alpha", down_alpha) + + def forward(self, x): + gate_out = torch.ops.auto_deploy.torch_fake_quant_nvfp4_linear( + x, + self.gate_weight, + None, + [self.gate_input_scale], + [self.gate_weight_scale, self.gate_alpha], + [], + [], + ) + up_out = torch.ops.auto_deploy.torch_fake_quant_nvfp4_linear( + x, + self.up_weight, + None, + [self.up_input_scale], + [self.up_weight_scale, self.up_alpha], + [], + [], + ) + hidden = torch.nn.functional.silu(gate_out) * up_out + return torch.ops.auto_deploy.torch_fake_quant_nvfp4_linear( + hidden, + self.down_weight, + None, + [self.down_input_scale], + [self.down_weight_scale, self.down_alpha], + [], + [], + ) + + +class NVFP4SwiGLUTestModel(nn.Module): + """Test model wrapping NVFP4 SwiGLU MLP between linear layers.""" + + def __init__(self, hidden_size: int = 128, intermediate_size: int = 128): + super().__init__() + device = torch.device("cuda") + self.linear_in = nn.Linear(hidden_size, hidden_size, device=device, dtype=torch.float16) + self.mlp = NVFP4SwiGLUMLP(hidden_size, intermediate_size) + self.linear_out = nn.Linear(hidden_size, hidden_size, device=device, dtype=torch.float16) + + def forward(self, x): + x = self.linear_in(x) + x = self.mlp(x) + x = self.linear_out(x) + return x + + +class NVFP4SwiGLUMultiLayerModel(nn.Module): + """Test model with multiple NVFP4 SwiGLU MLP layers.""" + + def __init__( + self, + hidden_size: int = 128, + intermediate_size: int = 128, + num_layers: int = 2, + ): + super().__init__() + device = torch.device("cuda") + self.layers = nn.ModuleList() + for _ in range(num_layers): + self.layers.append( + nn.ModuleDict( + { + "linear": nn.Linear( + hidden_size, + hidden_size, + device=device, + dtype=torch.float16, + ), + "mlp": NVFP4SwiGLUMLP(hidden_size, intermediate_size), + } + ) + ) + + def forward(self, x): + for layer in self.layers: + x = layer["linear"](x) + x = layer["mlp"](x) + return x + + +# -- Test helpers -------------------------------------------------------------- + + +def _count_ops(gm, op): + """Count how many nodes in the graph match the given op.""" + return sum(1 for n in gm.graph.nodes if is_op(n, op)) + + +def _has_no_fake_quant_nvfp4(gm): + """Verify no torch_fake_quant_nvfp4_linear ops remain.""" + return _count_ops(gm, torch.ops.auto_deploy.torch_fake_quant_nvfp4_linear) == 0 + + +# -- Tests --------------------------------------------------------------------- + + +@pytest.mark.skipif(_skip_condition, reason=_skip_reason) +def test_nvfp4_swiglu_pattern_match_only(): + """Test that match_nvfp4_swiglu_pattern produces torch_nvfp4_swiglu_mlp op.""" + torch.manual_seed(0) + model = NVFP4SwiGLUMLP().to("cuda") + x = torch.randn(2, 128, device="cuda", dtype=torch.float16) + + gm = torch_export_to_gm(model, args=(x,), clone=True) + + # Verify the graph has torch_fake_quant_nvfp4_linear ops before transform + assert _count_ops(gm, torch.ops.auto_deploy.torch_fake_quant_nvfp4_linear) == 3, ( + "Expected 3 torch_fake_quant_nvfp4_linear ops (gate, up, down) before transform" + ) + + # Apply only pattern matching + gm_matched = InferenceOptimizer( + None, + { + "match_nvfp4_swiglu_pattern": { + "stage": "pattern_matcher", + }, + }, + )(None, gm) + + # Check the intermediate op is present + nvfp4_swiglu_count = _count_ops( + gm_matched, torch.ops.auto_deploy.torch_nvfp4_swiglu_mlp.default + ) + assert nvfp4_swiglu_count == 1, ( + f"Expected 1 torch_nvfp4_swiglu_mlp op, got {nvfp4_swiglu_count}" + ) + + # All 3 fake_quant_nvfp4 ops should be consumed + assert _has_no_fake_quant_nvfp4(gm_matched), ( + "torch_fake_quant_nvfp4_linear ops should be consumed by pattern matcher" + ) + + # Verify numerical correctness + gm_matched = gm_matched.to("cuda") + y_matched = gm_matched(x) + y_model = model(x) + torch.testing.assert_close(y_matched, y_model, atol=1e-3, rtol=1e-3) + + +@pytest.mark.skipif(_skip_condition, reason=_skip_reason) +def test_nvfp4_swiglu_full_fusion(): + """Test full pipeline: pattern match -> fuse -> fused_nvfp4_swiglu_mlp.""" + torch.manual_seed(0) + model = NVFP4SwiGLUTestModel().to("cuda") + x = torch.randn(2, 128, device="cuda", dtype=torch.float16) + + gm = torch_export_to_gm(model, args=(x,), clone=True) + + # Apply pattern matching + fusion + gm_fused = InferenceOptimizer( + None, + { + "match_nvfp4_swiglu_pattern": { + "stage": "pattern_matcher", + }, + "fuse_nvfp4_swiglu": { + "stage": "post_load_fusion", + }, + }, + )(None, gm) + + gm_fused = gm_fused.to("cuda") + + # Check the fused op is present + fused_count = _count_ops(gm_fused, torch.ops.auto_deploy.fused_nvfp4_swiglu_mlp.default) + assert fused_count == 1, f"Expected 1 fused_nvfp4_swiglu_mlp op, got {fused_count}" + + # No intermediate or unfused ops should remain + assert _count_ops(gm_fused, torch.ops.auto_deploy.torch_nvfp4_swiglu_mlp.default) == 0, ( + "Intermediate torch_nvfp4_swiglu_mlp should be replaced by fused version" + ) + assert _has_no_fake_quant_nvfp4(gm_fused), ( + "No torch_fake_quant_nvfp4_linear ops should remain after fusion" + ) + + # Verify numerical correctness (fused uses TRT-LLM kernel, allow wider tolerance) + y_fused = gm_fused(x) + y_model = model(x) + torch.testing.assert_close(y_fused, y_model, atol=0.15, rtol=0.05) + + # Test with a different batch size to verify dynamic shapes work + x2 = torch.randn(4, 128, device="cuda", dtype=torch.float16) + y_fused_2 = gm_fused(x2) + y_model_2 = model(x2) + torch.testing.assert_close(y_fused_2, y_model_2, atol=0.15, rtol=0.05) + + +@pytest.mark.skipif(_skip_condition, reason=_skip_reason) +@pytest.mark.parametrize("num_layers", [2, 3]) +def test_nvfp4_swiglu_fusion_multiple_layers(num_layers): + """Test that multiple NVFP4 SwiGLU patterns are fused correctly.""" + torch.manual_seed(0) + model = NVFP4SwiGLUMultiLayerModel(num_layers=num_layers).to("cuda") + x = torch.randn(2, 128, device="cuda", dtype=torch.float16) + + gm = torch_export_to_gm(model, args=(x,), clone=True) + + # Apply pattern matching + fusion + gm_fused = InferenceOptimizer( + None, + { + "match_nvfp4_swiglu_pattern": { + "stage": "pattern_matcher", + }, + "fuse_nvfp4_swiglu": { + "stage": "post_load_fusion", + }, + }, + )(None, gm) + + gm_fused = gm_fused.to("cuda") + + # Check that all layers are fused + fused_count = _count_ops(gm_fused, torch.ops.auto_deploy.fused_nvfp4_swiglu_mlp.default) + assert fused_count == num_layers, ( + f"Expected {num_layers} fused_nvfp4_swiglu_mlp ops, got {fused_count}" + ) + + # Verify numerical correctness + y_fused = gm_fused(x) + y_model = model(x) + torch.testing.assert_close(y_fused, y_model, atol=0.2, rtol=0.1) + + +@pytest.mark.skipif(_skip_condition, reason=_skip_reason) +def test_nvfp4_swiglu_does_not_match_non_swiglu(): + """Test that the NVFP4 SwiGLU matcher does not match non-SwiGLU NVFP4 linears.""" + torch.manual_seed(0) + device = torch.device("cuda") + hidden_size = 128 + + # Model with two sequential NVFP4 linears + relu (NOT a SwiGLU pattern) + class NonSwiGLUModel(nn.Module): + def __init__(self): + super().__init__() + w1 = torch.randn(hidden_size, hidden_size, dtype=torch.half, device=device) * 0.05 + w2 = torch.randn(hidden_size, hidden_size, dtype=torch.half, device=device) * 0.05 + + s_in = fp4_global_scale(torch.randn(1, hidden_size, dtype=torch.half, device=device)) + s_w1 = fp4_global_scale(w1) + s_w2 = fp4_global_scale(w2) + + w1_fp4, w1_cutlass = torch.ops.trtllm.fp4_quantize(w1, s_w1, 16, False) + w2_fp4, w2_cutlass = torch.ops.trtllm.fp4_quantize(w2, s_w2, 16, False) + + self.register_buffer("w1", w1_fp4) + self.register_buffer("w1_is", s_in.to(torch.float32)) + self.register_buffer("w1_ws", w1_cutlass) + self.register_buffer("w1_a", (1.0 / (s_in * s_w1)).to(torch.float32)) + + self.register_buffer("w2", w2_fp4) + self.register_buffer("w2_is", s_in.to(torch.float32)) + self.register_buffer("w2_ws", w2_cutlass) + self.register_buffer("w2_a", (1.0 / (s_in * s_w2)).to(torch.float32)) + + def forward(self, x): + # Sequential linears without SwiGLU pattern + y = torch.ops.auto_deploy.torch_fake_quant_nvfp4_linear( + x, self.w1, None, [self.w1_is], [self.w1_ws, self.w1_a], [], [] + ) + y = torch.nn.functional.relu(y) + return torch.ops.auto_deploy.torch_fake_quant_nvfp4_linear( + y, self.w2, None, [self.w2_is], [self.w2_ws, self.w2_a], [], [] + ) + + model = NonSwiGLUModel().to("cuda") + x = torch.randn(2, hidden_size, device="cuda", dtype=torch.float16) + + gm = torch_export_to_gm(model, args=(x,), clone=True) + + gm_result = InferenceOptimizer( + None, + { + "match_nvfp4_swiglu_pattern": { + "stage": "pattern_matcher", + }, + }, + )(None, gm) + + # No SwiGLU ops should be found + assert _count_ops(gm_result, torch.ops.auto_deploy.torch_nvfp4_swiglu_mlp.default) == 0, ( + "Non-SwiGLU NVFP4 pattern should not match" + ) + + # Original NVFP4 linear ops should still be present + assert _count_ops(gm_result, torch.ops.auto_deploy.torch_fake_quant_nvfp4_linear) == 2, ( + "Original NVFP4 linear ops should be unchanged" + ) From ab2e4f5d45e93265ef829239a4c7b5ad961dff59 Mon Sep 17 00:00:00 2001 From: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com> Date: Sat, 7 Feb 2026 19:26:44 -0800 Subject: [PATCH 15/17] fix swiglu for b200 Signed-off-by: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com> --- .../auto_deploy/custom_ops/linear/swiglu.py | 30 ++++++++++++------- .../defs/accuracy/test_llm_api_autodeploy.py | 5 +++- .../library/test_nvfp4_swiglu.py | 3 +- 3 files changed, 26 insertions(+), 12 deletions(-) diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/linear/swiglu.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/linear/swiglu.py index 306744d588ec..e13acb6f79fa 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/linear/swiglu.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/linear/swiglu.py @@ -25,6 +25,22 @@ import torch import torch.nn.functional as F +try: + from flashinfer.activation import silu_and_mul as _flashinfer_silu_and_mul +except ImportError: + _flashinfer_silu_and_mul = None + + +def _silu_and_mul(x: torch.Tensor) -> torch.Tensor: + """SwiGLU activation: split x in half, apply silu to first half, multiply with second half. + + Uses FlashInfer's fused kernel when available, falls back to manual implementation. + """ + if _flashinfer_silu_and_mul is not None: + return _flashinfer_silu_and_mul(x) + gate, up = x.chunk(2, dim=-1) + return F.silu(gate) * up + @torch.library.custom_op("auto_deploy::torch_swiglu_mlp", mutates_args=()) def torch_swiglu_mlp( @@ -104,11 +120,8 @@ def fused_swiglu_mlp( # Single matmul for both gate and up projections gate_up_out = F.linear(input, gate_up_weight, gate_up_bias) - # Split into gate and up outputs - gate_out, up_out = gate_up_out.chunk(2, dim=-1) - - # Apply SwiGLU activation: silu(gate) * up - hidden = F.silu(gate_out) * up_out + # Apply SwiGLU activation: split, silu(gate) * up (uses FlashInfer when available) + hidden = _silu_and_mul(gate_up_out) # Down projection return F.linear(hidden, down_weight, down_bias) @@ -266,11 +279,8 @@ def fused_nvfp4_swiglu_mlp( alpha=gate_up_alpha, ) - # Split into gate and up outputs - gate_out, up_out = gate_up_out.chunk(2, dim=-1) - - # Apply SwiGLU activation: silu(gate) * up - hidden = F.silu(gate_out) * up_out + # Apply SwiGLU activation: split, silu(gate) * up (uses FlashInfer when available) + hidden = _silu_and_mul(gate_up_out) # Down projection return torch.ops.auto_deploy.torch_quant_nvfp4_linear( diff --git a/tests/integration/defs/accuracy/test_llm_api_autodeploy.py b/tests/integration/defs/accuracy/test_llm_api_autodeploy.py index 638887059099..acfe0f2637a9 100644 --- a/tests/integration/defs/accuracy/test_llm_api_autodeploy.py +++ b/tests/integration/defs/accuracy/test_llm_api_autodeploy.py @@ -455,7 +455,10 @@ def get_default_kwargs(self, enable_chunked_prefill=False): "transforms": { "multi_stream_moe": { "stage": "compile", - # multi-stream MOE currently does not work for world_size > 1 + "enabled": True, + }, + "multi_stream_mla_attn": { + "stage": "compile", "enabled": True, }, } diff --git a/tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_nvfp4_swiglu.py b/tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_nvfp4_swiglu.py index 8e4e3f329163..0687b1fd53d6 100644 --- a/tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_nvfp4_swiglu.py +++ b/tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_nvfp4_swiglu.py @@ -9,6 +9,7 @@ import torch import torch.nn as nn from _torch_test_utils import fp4_compatible, trtllm_ops_available +from torch.export import Dim import tensorrt_llm._torch.auto_deploy.custom_ops # noqa: F401 from tensorrt_llm._torch.auto_deploy.export import torch_export_to_gm @@ -233,7 +234,7 @@ def test_nvfp4_swiglu_full_fusion(): model = NVFP4SwiGLUTestModel().to("cuda") x = torch.randn(2, 128, device="cuda", dtype=torch.float16) - gm = torch_export_to_gm(model, args=(x,), clone=True) + gm = torch_export_to_gm(model, args=(x,), clone=True, dynamic_shapes=({0: Dim.DYNAMIC},)) # Apply pattern matching + fusion gm_fused = InferenceOptimizer( From 37c6eb173d5e782ffc22471a06f16ec8c2027631 Mon Sep 17 00:00:00 2001 From: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com> Date: Sat, 7 Feb 2026 19:29:39 -0800 Subject: [PATCH 16/17] fix default.yaml Signed-off-by: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com> --- tensorrt_llm/_torch/auto_deploy/config/default.yaml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/auto_deploy/config/default.yaml b/tensorrt_llm/_torch/auto_deploy/config/default.yaml index c4e92fecd18c..3a7e60c926fa 100644 --- a/tensorrt_llm/_torch/auto_deploy/config/default.yaml +++ b/tensorrt_llm/_torch/auto_deploy/config/default.yaml @@ -85,6 +85,7 @@ transforms: stage: pattern_matcher match_nvfp4_swiglu_pattern: stage: pattern_matcher + requires_shape_prop: true quantize_fp8_moe: stage: pattern_matcher quantize_nvfp4_moe: @@ -168,7 +169,7 @@ transforms: gather_logits_before_lm_head: stage: post_load_fusion # TODO: fix https://github.com/NVIDIA/TensorRT-LLM/issues/9878 to enable by default - enabled: false + enabled: true ############################################################################################ # VISUALIZE GRAPH ############################################################################################ From 382456a7240c45084bbf2754798d64c48e42375a Mon Sep 17 00:00:00 2001 From: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com> Date: Sat, 7 Feb 2026 20:33:51 -0800 Subject: [PATCH 17/17] fix defaults, fix triton rope unit test Signed-off-by: Suyog Gupta <41447211+suyoggupta@users.noreply.github.com> --- examples/auto_deploy/glm_flash.yaml | 23 ------------------- .../model_registry/configs/glm-4.7-flash.yaml | 7 ++++++ .../_torch/auto_deploy/config/default.yaml | 2 +- .../custom_ops/rope/test_triton_rope.py | 2 +- 4 files changed, 9 insertions(+), 25 deletions(-) delete mode 100644 examples/auto_deploy/glm_flash.yaml diff --git a/examples/auto_deploy/glm_flash.yaml b/examples/auto_deploy/glm_flash.yaml deleted file mode 100644 index 87482e47423e..000000000000 --- a/examples/auto_deploy/glm_flash.yaml +++ /dev/null @@ -1,23 +0,0 @@ -runtime: trtllm -compile_backend: torch-cudagraph -max_seq_len: 4096 -max_num_tokens: 4096 -max_batch_size: 64 -enable_chunked_prefill: true -model_factory: AutoModelForCausalLM -disable_overlap_scheduler: false -cuda_graph_batch_sizes: [1, 2, 4, 8, 16, 32, 64] -kv_cache_config: - enable_block_reuse: false - free_gpu_memory_fraction: 0.7 - tokens_per_block: 64 -model_kwargs: - torch_dtype: bfloat16 - -transforms: - multi_stream_moe: - stage: compile - enabled: true - multi_stream_mla_attn: - stage: compile - enabled: true diff --git a/examples/auto_deploy/model_registry/configs/glm-4.7-flash.yaml b/examples/auto_deploy/model_registry/configs/glm-4.7-flash.yaml index fb434e8bb7c0..cc8892728a80 100644 --- a/examples/auto_deploy/model_registry/configs/glm-4.7-flash.yaml +++ b/examples/auto_deploy/model_registry/configs/glm-4.7-flash.yaml @@ -3,3 +3,10 @@ max_batch_size: 64 max_seq_len: 4096 enable_chunked_prefill: true cuda_graph_batch_sizes: [1, 2, 4, 8, 16, 32, 64] +transforms: + multi_stream_moe: + stage: compile + enabled: true + multi_stream_mla_attn: + stage: compile + enabled: true diff --git a/tensorrt_llm/_torch/auto_deploy/config/default.yaml b/tensorrt_llm/_torch/auto_deploy/config/default.yaml index 3a7e60c926fa..1a75ac010de3 100644 --- a/tensorrt_llm/_torch/auto_deploy/config/default.yaml +++ b/tensorrt_llm/_torch/auto_deploy/config/default.yaml @@ -169,7 +169,7 @@ transforms: gather_logits_before_lm_head: stage: post_load_fusion # TODO: fix https://github.com/NVIDIA/TensorRT-LLM/issues/9878 to enable by default - enabled: true + enabled: false ############################################################################################ # VISUALIZE GRAPH ############################################################################################ diff --git a/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/rope/test_triton_rope.py b/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/rope/test_triton_rope.py index fccda7d4286c..748394648b8e 100644 --- a/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/rope/test_triton_rope.py +++ b/tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/rope/test_triton_rope.py @@ -5,7 +5,7 @@ from _custom_op_utils import torch_rope_reference # Import after we've imported torch (to ensure custom ops are registered) -from tensorrt_llm._torch.auto_deploy.custom_ops import triton_rope # noqa: F401 +from tensorrt_llm._torch.auto_deploy.custom_ops.rope import triton_rope # noqa: F401 def _precompute_cos_sin_cache(