commit f37510c88f181eaf9b07da4d8c585b75eb9167a6 Author: Sun Ruoxi Date: Wed Oct 21 15:52:23 2026 +0800 first commit diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..8de7abc --- /dev/null +++ b/.dockerignore @@ -0,0 +1,6 @@ +.git +__pycache__ +*.pyc +upstream_archive/ +patches/ +README.md diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..21bd9fa --- /dev/null +++ b/.gitignore @@ -0,0 +1,3 @@ +__pycache__/ +*.pyc +upstream_archive/ diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..77016ae --- /dev/null +++ b/Dockerfile @@ -0,0 +1,16 @@ +ARG BASE_IMAGE +FROM ${BASE_IMAGE} + +COPY manifest.json /opt/ppu-deepseek-v4-pro-patch/manifest.json +COPY scripts/apply.py /opt/ppu-deepseek-v4-pro-patch/scripts/apply.py +COPY src/ /opt/ppu-deepseek-v4-pro-patch/src/ +COPY upstream/ /opt/ppu-deepseek-v4-pro-patch/upstream/ +COPY examples/ /opt/ppu-deepseek-v4-pro-patch/examples/ + +RUN python /opt/ppu-deepseek-v4-pro-patch/scripts/apply.py --apply + +ENV VLLM_DISABLED_KERNELS=PPUCutlassFp8BlockScaledMMKernel \ + VLLM_PPU_BF16_O_PROJ=1 \ + VLLM_PPU_SERIAL_ATTENTION=1 \ + VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD=0 \ + VLLM_ENGINE_READY_TIMEOUT_S=3600 diff --git a/README.md b/README.md new file mode 100644 index 0000000..7534e6a --- /dev/null +++ b/README.md @@ -0,0 +1,123 @@ +# DeepSeek V4 Pro / PPU-ZW810E vLLM 补丁 + +这是 2026-10-20 至 2026-10-21 在 `PPU-ZW810E` 上验证过的 DeepSeek V4 Pro 推理补丁。它以原始 vLLM 镜像为基础，修改四个 vLLM Python 源文件；**不修改 `/model` 的 checkpoint 或模型 `config.json`**。当前验证环境是 vLLM 0.26.0、PyTorch 2.11.0、Python 3.12、16 张 PPU，TP=8、DP=2、EP 开启。 + +已验证的基础镜像： + +```text +swr.cn-north-4.myhuaweicloud.com/inference-engines/llm:v2.1.1-pytorch2.11.0-ubuntu24.04-cuda13.0-vllm0.26.0-py312 +``` + +## 修复了什么 + +1. **FP8 block linear 无法通过 PPU CUTLASS sm80 fallback。**原服务加载完 64 个权重 shard 后，在 `cutlass_scaled_mm_sm80` / `scaled_mm_c2x.cu:110` 报错。运行时设置 `VLLM_DISABLED_KERNELS=PPUCutlassFp8BlockScaledMMKernel`，从候选列表中移除这个出错的 kernel；本次模型日志确认随后选择了 Marlin 的 FP8 weight-only 路径。PPU DeepGEMM 仍保留在候选列表中，其他形状或模型不保证总是选择 Marlin。Marlin 可能比原生 FP8 计算慢。 +2. **输出投影的 FP8 融合算子不支持该架构。**`_fused_inv_rope_fp8_quant_per_head` 使用 `fp8e4nv`，而 PPU 编译器只支持 `fp8e4b15`、`fp8e5`。`VLLM_PPU_BF16_O_PROJ=1` 时，在 Marlin 打包前将 `wo_a` 的原始 block FP8 权重与 scale 解码成 BF16；输出投影使用 FP32 中间值完成 inverse RoPE、BF16 分组 `bmm`，然后调用现有 `wo_b`。这是 FP8 **解码到 BF16**，不是重新编码成 uint8 FP8。FP8 channelwise 路径也会遇到 `fp8e4nv`，因此没有采用。 +3. **compressor GEMM 触发设备参数错误。**原 BF16 输入、FP32 输出的 GEMM 在 PPU 上报 `CUBLAS_STATUS_INVALID_VALUE`；失败后再在同一流上转换 FP32 会继续报 `invalid argument`。补丁直接把两侧输入转换为 FP32，再执行普通 `torch.mm`，避免先调用已知失败的 GEMM。 +4. **profile / 稀疏 MLA 预热仍有异步设备错误。**曾在 MoE 加法以及 indexer 的 `torch.empty` 处看到 `invalid argument`；这些位置不一定是实际出错算子。仅把 `VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD` 设为 `0` 不会关闭注意力的辅助流。新增 `VLLM_PPU_SERIAL_ATTENTION=1`，使注意力、indexer、compressor 使用单流。该配置完成了启动和实际推理，但**尚未单独证明多流是唯一根因**。 + +代码位置： + +| 文件 | 改动 | +| --- | --- | +| `src/vllm/model_executor/layers/quantization/fp8.py` | 在权重处理阶段为 `wo_a` 启用 BF16 解码 | +| `src/vllm/models/deepseek_v4/ppu/ops/o_proj.py` | BF16 输出投影与 FP32 compressor GEMM | +| `src/vllm/models/deepseek_v4/ppu/flashmla.py` | 将 PPU 输出投影切换到 BF16 路径 | +| `src/vllm/models/deepseek_v4/attention.py` | 两处 compressor 调用与注意力单流开关 | + +`upstream/` 是从原始文件备份取得的精确版本，`src/` 是成功运行的精确版本，`patches/vllm.patch` 是两者的统一 diff，`manifest.json` 记录两组 SHA-256。`scripts/apply.py` 在写入前核对全部四个目标文件和包内源码；遇到不同的 vLLM 版本会拒绝安装。首次安装会在目标文件旁生成 `.ppu-dsv4.orig` 备份。 + +该版本也接受 `vllm serve ... --linear-backend marlin`，它会要求所有适用的量化 linear 使用 Marlin；如果某种层没有可用的 Marlin kernel，启动会直接报错。它**不是** `VLLM_DISABLED_KERNELS=PPUCutlassFp8BlockScaledMMKernel` 的等价替代，且本项目的完整模型成功验证没有使用该选项。`--quantization marlin` 是另一类量化配置，不能把当前 FP8 checkpoint 当作 Marlin 格式 checkpoint 来指定。 + +## 构建包含补丁的新镜像（推荐） + +先将本目录放到有原始镜像的 PPU 主机上。例如，在本机 `/Users/4paradigm/projects` 下运行： + +```sh +scp -i llm_benchmark/bench-text_generation-tco/utils/id_rsa -r \ + ppu_deepseek_v4_pro_patch root@172.28.4.222:/tmp/ +``` + +在 PPU 主机上执行： + +```sh +cd /tmp/ppu_deepseek_v4_pro_patch +nerdctl build \ + --build-arg BASE_IMAGE=swr.cn-north-4.myhuaweicloud.com/inference-engines/llm:v2.1.1-pytorch2.11.0-ubuntu24.04-cuda13.0-vllm0.26.0-py312 \ + -t deepseek-v4-pro:ppu-bf16-serial . +``` + +若主机提示 BuildKit daemon 未运行，可启动仅用于本次构建的 BuildKit containerd worker，再指定其 socket；这样会复用 containerd 的镜像层缓存： + +```sh +nohup buildkitd --oci-worker=false --containerd-worker=true \ + --containerd-worker-namespace default \ + --addr unix:///tmp/ppu-buildkit.sock --root /tmp/ppu-buildkit \ + > /tmp/ppu-buildkit.log 2>&1 < /dev/null & +nerdctl build --buildkit-host unix:///tmp/ppu-buildkit.sock \ + --build-arg BASE_IMAGE=swr.cn-north-4.myhuaweicloud.com/inference-engines/llm:v2.1.1-pytorch2.11.0-ubuntu24.04-cuda13.0-vllm0.26.0-py312 \ + -t deepseek-v4-pro:ppu-bf16-serial . +``` + +构建时的 `RUN python .../scripts/apply.py --apply` 会先验证原始镜像内四个文件的 SHA-256，再替换源码。`nerdctl build` 成功即表示补丁已写入**新镜像**，原镜像标签不会被覆盖。构建失败若提示 `Installed source does not match this patch`，说明基础镜像中的 vLLM 源码与本项目的原始版本不同；应先比较源码，不要强行跳过校验。 + +创建服务容器时，沿用原容器的 PPU 设备映射、权限、共享内存和网络配置，改用 `deepseek-v4-pro:ppu-bf16-serial` 镜像，并将模型挂载到 `/model`。这台机器原容器的模型挂载源是 `/bmcp_lvm_fs/data/model/DeepSeek-V4-Pro`，容器目标是 `/model`，共享内存配置为 64 GiB。PPU 设备如何注入由机器的容器运行环境决定，应复用原容器的创建参数。新镜像默认设置 `VLLM_DISABLED_KERNELS`、`VLLM_PPU_BF16_O_PROJ`、`VLLM_PPU_SERIAL_ATTENTION`、`VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD` 和 `VLLM_ENGINE_READY_TIMEOUT_S=3600`。后者为权重加载与首次 JIT 编译留出时间；默认的 600 秒在本机使 API 进程提前退出。 + +在新容器里启动已验证的配置（把 `deepseek_v4_patched` 换成实际容器名）： + +```sh +nerdctl exec -d \ + -e CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15 \ + deepseek_v4_patched \ + python /opt/ppu-deepseek-v4-pro-patch/examples/launch.py serve.log +``` + +启动器读取 `examples/command-serial.json`，将日志写到 `/opt/ppu-deepseek-v4-pro-patch/examples/serve.log`。它运行的核心命令如下；如果按原有编排系统直接运行 `vllm serve`，需同时传入镜像默认的五个环境变量，并明确指定全部 16 张卡： + +```sh +vllm serve /model \ + --trust-remote-code --kv-cache-dtype fp8 --block-size 256 \ + --tensor-parallel-size 8 --data-parallel-size 2 --enable-expert-parallel \ + --no-enable-flashinfer-autotune \ + --tokenizer-mode deepseek_v4 --tool-call-parser deepseek_v4 \ + --enable-auto-tool-choice --reasoning-parser deepseek_v4 \ + --max-model-len 4096 --max-num-batched-tokens 1024 \ + --enforce-eager --port 18000 +``` + +该启动器将 `RLIMIT_MEMLOCK` 设为 unlimited；如果容器运行时不允许此设置，需在创建容器时提供相应 `memlock` 限额。`CUDA_VISIBLE_DEVICES` 不在镜像中固定，因为实际可见的 PPU 取决于部署环境。 + +## 检查、直接安装与回滚 + +新镜像内可检查四个文件的状态： + +```sh +nerdctl exec deepseek_v4_patched \ + python /opt/ppu-deepseek-v4-pro-patch/scripts/apply.py --check +``` + +如果要对一个**尚未修改**的原始容器做临时验证，可以将本目录复制进去，再检查、安装；这只修改该容器的可写层，不能替代上面的镜像构建： + +```sh +cd /workspace/pytorch +nerdctl cp ./ppu_deepseek_v4_pro_patch deepseek_v4_original:/opt/ppu-deepseek-v4-pro-patch +nerdctl exec deepseek_v4_original python /opt/ppu-deepseek-v4-pro-patch/scripts/apply.py --check +nerdctl exec deepseek_v4_original python /opt/ppu-deepseek-v4-pro-patch/scripts/apply.py --apply +``` + +`--apply` 可重复运行；已经是补丁版本时不再替换。容器内回滚四处源码： + +```sh +nerdctl exec deepseek_v4_original python /opt/ppu-deepseek-v4-pro-patch/scripts/apply.py --rollback +``` + +回滚后应停止服务并用原始环境变量重新启动；对于基于新镜像创建的容器，更直接的完整回滚方式是重新使用原始镜像。不要在正在执行推理的进程内热替换 Python 源码。 + +## 验证范围与限制 + +- 四个源码文件均通过语法检查；`patches/vllm.patch` 可对 `upstream/` 通过 `git apply --check`。安装脚本的幂等安装、回滚和错误版本拒绝有离线测试。 +- 合成尺寸的 BF16 `wo_a` 解码与输出投影相对参考计算的 L2 误差约 0.0016；`(1024,7168) @ (7168,1024)` 的 FP32 compressor GEMM 单卡测试通过。 +- 实际模型加载约 55 GiB/卡；TP8/DP2/EP 的服务通过 profile、KV cache 分配、稀疏 MLA 预热。一次短请求和四次并发短请求成功；两个 API 进程均处理了请求。 +- 长上下文测试中，1037、3085、3853、4045 输入 token 的请求各生成 32 token，均成功。一个 3619 输入 token 的聊天请求正确返回文中预埋的“蓝色”。测试后日志没有 `ERROR`。 +- 已验证配置使用 `--enforce-eager` 和 `--max-model-len 4096`。超过 4096 token、较高并发、CUDA graph 模式以及生产性能尚未验证。原始未禁用辅助流的失败不能仅凭当前结果断定就是多流造成的。 + +历史排查细节与原始日志保存在同一工作区的 `llm_benchmark/bench-text_generation-tco/DEEPSEEK_V4_PPU_2026-10-20.md` 和 `ppu_debug_snapshot_2026-10-20/`。 diff --git a/examples/command-serial.json b/examples/command-serial.json new file mode 100644 index 0000000..85448c9 --- /dev/null +++ b/examples/command-serial.json @@ -0,0 +1,40 @@ +{ + "args": [ + "vllm", + "serve", + "/model", + "--trust-remote-code", + "--kv-cache-dtype", + "fp8", + "--block-size", + "256", + "--tensor-parallel-size", + "8", + "--data-parallel-size", + "2", + "--enable-expert-parallel", + "--no-enable-flashinfer-autotune", + "--tokenizer-mode", + "deepseek_v4", + "--tool-call-parser", + "deepseek_v4", + "--enable-auto-tool-choice", + "--reasoning-parser", + "deepseek_v4", + "--max-model-len", + "4096", + "--max-num-batched-tokens", + "1024", + "--enforce-eager", + "--port", + "18000" + ], + "env": { + "VLLM_DISABLED_KERNELS": "PPUCutlassFp8BlockScaledMMKernel", + "VLLM_PPU_BF16_O_PROJ": "1", + "CUDA_VISIBLE_DEVICES": "0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15", + "VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD": "0", + "VLLM_PPU_SERIAL_ATTENTION": "1", + "VLLM_ENGINE_READY_TIMEOUT_S": "3600" + } +} diff --git a/examples/launch.py b/examples/launch.py new file mode 100644 index 0000000..34d8a34 --- /dev/null +++ b/examples/launch.py @@ -0,0 +1,20 @@ +import os,sys,json,resource,subprocess,signal,time +from pathlib import Path +base=Path(__file__).parent +resource.setrlimit(resource.RLIMIT_MEMLOCK,(resource.RLIM_INFINITY,resource.RLIM_INFINITY)) +cfg=json.loads((base/"command-serial.json").read_text()) +env=os.environ.copy();env.update(cfg["env"]) +log=base/(sys.argv[1] if len(sys.argv)>1 else "serve-fixed.log") +with log.open("w") as f: + p=subprocess.Popen(cfg["args"],env=env,stdout=f,stderr=subprocess.STDOUT,start_new_session=True,cwd="/workspace/pytorch") +(base/"serve.pid").write_text(str(p.pid)) +print("SERVER_PID",p.pid,"LOG",log,flush=True) +rc=p.wait() +print("SERVER_EXIT",rc,flush=True) +if rc: + try: os.killpg(p.pid,signal.SIGTERM) + except ProcessLookupError: pass + time.sleep(3) + try: os.killpg(p.pid,signal.SIGKILL) + except ProcessLookupError: pass +sys.exit(rc) diff --git a/manifest.json b/manifest.json new file mode 100644 index 0000000..5e620fb --- /dev/null +++ b/manifest.json @@ -0,0 +1,21 @@ +{ + "format": 1, + "files": { + "vllm/model_executor/layers/quantization/fp8.py": { + "original_sha256": "f412967fe6e764e91ed0f3e5cfd8fcd49372349af7a180360329fc6dc707dd69", + "patched_sha256": "e9864d722e0438b2657434a75ae47d87b6716dfd8deaf2e4224d71b999609828" + }, + "vllm/models/deepseek_v4/attention.py": { + "original_sha256": "8499ba6857f023622574fae7d61cff44d3f83c3e25d54e1371e7332e2ffb1221", + "patched_sha256": "fe2fe82e53994721957dc12bba53e498151f8bf3e827098340c7f815787aa563" + }, + "vllm/models/deepseek_v4/ppu/flashmla.py": { + "original_sha256": "f94eb897af22af3ac2f33b8cf0ebdc129d9686f45940d735a01658cabd9ab44f", + "patched_sha256": "c269fce0c6bd153ee01f4e8a71f1d8b5e0a0bbcbb3a007e8c24d2649f004b3bb" + }, + "vllm/models/deepseek_v4/ppu/ops/o_proj.py": { + "original_sha256": "4f90a12b18f7a7c17c9e1fcfdae9851f20f7792adecc76a775a09a037d74ab74", + "patched_sha256": "84953951f7582afd7c4eae4c0efe9eaa3823bf68284affc9806e2344ae566260" + } + } +} diff --git a/patches/vllm.patch b/patches/vllm.patch new file mode 100644 index 0000000..0ff4b18 --- /dev/null +++ b/patches/vllm.patch @@ -0,0 +1,153 @@ +diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py +index 56f43d5..4a76eca 100644 +--- a/vllm/model_executor/layers/quantization/fp8.py ++++ b/vllm/model_executor/layers/quantization/fp8.py +@@ -402,6 +402,15 @@ class Fp8LinearMethod(LinearMethodBase): + self.use_marlin = isinstance(self.fp8_linear, MarlinFP8ScaledMMLinearKernel) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: ++ # Opt-in fallback before Marlin repacks wo_a. ++ import os ++ if (os.environ.get("VLLM_PPU_BF16_O_PROJ") == "1" ++ and current_platform.is_ppu() ++ and getattr(layer, "is_bmm", False) ++ and self.block_quant): ++ from vllm.models.deepseek_v4.ppu.ops.o_proj import prepare_bf16_o_proj_weight ++ prepare_bf16_o_proj_weight(layer) ++ return + if self.use_marlin: + if not self.block_quant: + # Canonicalize to (K, N) for the kernel. +diff --git a/vllm/models/deepseek_v4/attention.py b/vllm/models/deepseek_v4/attention.py +index fedea3c..985654a 100644 +--- a/vllm/models/deepseek_v4/attention.py ++++ b/vllm/models/deepseek_v4/attention.py +@@ -169,6 +169,14 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): + cache_config = vllm_config.cache_config + tp_size = get_tensor_model_parallel_world_size() + layer_id = extract_layer_index(prefix) ++ # Diagnostic PPU path: run attention/indexer/compressor on one stream. ++ import os ++ from vllm.platforms import current_platform ++ if (current_platform.is_ppu() ++ and os.environ.get("VLLM_PPU_SERIAL_ATTENTION") == "1"): ++ aux_stream_list = None ++ if layer_id == 0: ++ logger.info("PPU attention auxiliary streams disabled") + + self.prefix = prefix # Alias for compatibility with compressor + self.hidden_size = config.hidden_size +@@ -395,6 +403,12 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): + compressor = self.compressor + + def compressor_kv_score() -> torch.Tensor: ++ import os ++ from vllm.platforms import current_platform ++ if (current_platform.is_ppu() ++ and os.environ.get("VLLM_PPU_BF16_O_PROJ") == "1"): ++ from vllm.models.deepseek_v4.ppu.ops.o_proj import ppu_compressor_mm ++ return ppu_compressor_mm(hidden_states, compressor.fused_wkv_wgate.weight.T) + return torch.mm( + hidden_states, + compressor.fused_wkv_wgate.weight.T, +@@ -412,6 +426,12 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): + return weights + + def indexer_compressor_kv_score() -> torch.Tensor: ++ import os ++ from vllm.platforms import current_platform ++ if (current_platform.is_ppu() ++ and os.environ.get("VLLM_PPU_BF16_O_PROJ") == "1"): ++ from vllm.models.deepseek_v4.ppu.ops.o_proj import ppu_compressor_mm ++ return ppu_compressor_mm(hidden_states, indexer.compressor.fused_wkv_wgate.weight.T) + return torch.mm( + hidden_states, + indexer.compressor.fused_wkv_wgate.weight.T, +diff --git a/vllm/models/deepseek_v4/ppu/flashmla.py b/vllm/models/deepseek_v4/ppu/flashmla.py +index 7e6b944..20cbfe8 100644 +--- a/vllm/models/deepseek_v4/ppu/flashmla.py ++++ b/vllm/models/deepseek_v4/ppu/flashmla.py +@@ -43,6 +43,16 @@ class DeepseekV4FlashMLAAttention(DeepseekV4Attention): + self._einsum_recipe, self._tma_aligned_scales = compute_fp8_einsum_recipe() + + def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: ++ import os ++ if os.environ.get("VLLM_PPU_BF16_O_PROJ") == "1": ++ from vllm.models.deepseek_v4.ppu.ops.o_proj import bf16_o_proj ++ return bf16_o_proj( ++ o, positions, self.rotary_emb.cos_sin_cache, self.wo_a, self.wo_b, ++ n_groups=self.n_local_groups, ++ heads_per_group=self.n_local_heads // self.n_local_groups, ++ nope_dim=self.nope_head_dim, rope_dim=self.rope_head_dim, ++ o_lora_rank=self.o_lora_rank, ++ ) + if self.wo_a.weight.dtype == torch.int8: + return deep_gemm_int8_o_proj( + o, +diff --git a/vllm/models/deepseek_v4/ppu/ops/o_proj.py b/vllm/models/deepseek_v4/ppu/ops/o_proj.py +index d6872f6..6941f10 100644 +--- a/vllm/models/deepseek_v4/ppu/ops/o_proj.py ++++ b/vllm/models/deepseek_v4/ppu/ops/o_proj.py +@@ -241,4 +241,60 @@ direct_register_custom_op( + op_func=deepseek_v4_fp8_channel_einsum, + mutates_args=["out"], + fake_impl=deepseek_v4_fp8_channel_einsum_fake, +-) +\ No newline at end of file ++) ++# Experimental ZW810E BF16 O-projection fallback (opt-in). ++def prepare_bf16_o_proj_weight(layer): ++ """Decode checkpoint block FP8 before any Marlin/DeepGEMM repacking.""" ++ from vllm.model_executor.utils import replace_parameter ++ w = layer.weight ++ s = layer.weight_scale_inv ++ if w.dtype != torch.float8_e4m3fn or w.ndim != 2: ++ raise ValueError(f"BF16 O-proj expects raw 2D FP8 weight, got {w.dtype} {w.shape}") ++ bn, bk = layer.weight_block_size ++ n, k = w.shape ++ expected = ((n + bn - 1) // bn, (k + bk - 1) // bk) ++ if tuple(s.shape) != expected: ++ raise ValueError(f"BF16 O-proj scale shape {s.shape}, expected {expected}") ++ # Convert row blocks to keep temporary FP32 memory bounded. ++ decoded = torch.empty((n, k), device=w.device, dtype=torch.bfloat16) ++ sf = s.float() ++ for row in range(0, n, bn): ++ expanded = sf[row // bn].repeat_interleave(bk)[:k] ++ decoded[row:row + bn] = ( ++ w[row:row + bn].float() * expanded[None, :] ++ ).to(torch.bfloat16) ++ replace_parameter(layer, "weight", decoded) ++ layer._ppu_bf16_o_proj_ready = True ++ ++ ++def bf16_o_proj( ++ o, positions, cos_sin_cache, wo_a, wo_b, *, ++ n_groups, heads_per_group, nope_dim, rope_dim, o_lora_rank, **kwargs ++): ++ """Inverse GPT-J RoPE followed by grouped BF16 GEMM and existing wo_b.""" ++ if not getattr(wo_a, "_ppu_bf16_o_proj_ready", False): ++ raise RuntimeError("BF16 O-proj weight was not prepared before repacking") ++ t, h, d = o.shape ++ if h != n_groups * heads_per_group or d != nope_dim + rope_dim: ++ raise ValueError(f"Unexpected O-proj input shape: {o.shape}") ++ # Match fused kernel: even = x*cos + partner*sin; odd = x*cos - partner*sin. ++ cs = cos_sin_cache.index_select(0, positions.reshape(-1).long()).float() ++ half = rope_dim // 2 ++ cos = cs[:, None, :half] ++ sin = cs[:, None, half:] ++ tail = o[..., nope_dim:].float().reshape(t, h, half, 2) ++ even, odd = tail[..., 0], tail[..., 1] ++ rotated = torch.stack((even * cos + odd * sin, ++ odd * cos - even * sin), dim=-1) ++ x = torch.cat((o[..., :nope_dim], ++ rotated.flatten(-2).to(o.dtype)), dim=-1) ++ r = heads_per_group * d ++ w = wo_a.weight.reshape(n_groups, o_lora_rank, r) ++ x = x.reshape(t, n_groups, r).transpose(0, 1).contiguous() ++ z = torch.bmm(x.to(torch.bfloat16), w.transpose(1, 2)) ++ return wo_b(z.transpose(0, 1).contiguous().flatten(1)) ++ ++# PPU cannot use the BF16-input / FP32-output compressor GEMM here. ++# Convert before GEMM so an invalid device operation cannot poison the stream. ++def ppu_compressor_mm(a, b): ++ return torch.mm(a.float().contiguous(), b.float().contiguous()) diff --git a/scripts/apply.py b/scripts/apply.py new file mode 100644 index 0000000..5651444 --- /dev/null +++ b/scripts/apply.py @@ -0,0 +1,114 @@ +#!/usr/bin/env python3 +"""Install, inspect, or roll back the verified DeepSeek V4 PPU source patch.""" + +from __future__ import annotations + +import argparse +import hashlib +import importlib.util +import json +import os +import stat +import sys +import tempfile +from pathlib import Path + + +ROOT = Path(__file__).resolve().parents[1] +BACKUP_SUFFIX = ".ppu-dsv4.orig" + + +def digest(data: bytes) -> str: + return hashlib.sha256(data).hexdigest() + + +def atomic_write(path: Path, data: bytes, mode: int) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + fd, temporary = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent) + try: + with os.fdopen(fd, "wb") as output: + output.write(data) + output.flush() + os.fsync(output.fileno()) + os.chmod(temporary, mode) + os.replace(temporary, path) + finally: + if os.path.exists(temporary): + os.unlink(temporary) + + +def site_packages_path(override: str | None) -> Path: + if override: + return Path(override).resolve() + spec = importlib.util.find_spec("vllm") + if spec is None or not spec.submodule_search_locations: + raise RuntimeError("Cannot locate installed vllm; pass --site-packages") + return Path(next(iter(spec.submodule_search_locations))).resolve().parent + + +def run() -> None: + parser = argparse.ArgumentParser(description=__doc__) + group = parser.add_mutually_exclusive_group(required=True) + group.add_argument("--check", action="store_true") + group.add_argument("--apply", action="store_true") + group.add_argument("--rollback", action="store_true") + parser.add_argument("--site-packages", help="Override site-packages root for testing") + args = parser.parse_args() + + site = site_packages_path(args.site_packages) + manifest = json.loads((ROOT / "manifest.json").read_text()) + if manifest.get("format") != 1: + raise RuntimeError("Unknown manifest format") + + planned = [] + for name, expected in manifest["files"].items(): + relative = Path(name) + if relative.is_absolute() or ".." in relative.parts or relative.parts[0] != "vllm": + raise RuntimeError(f"Unsafe manifest path: {name}") + target = site / relative + patched = (ROOT / "src" / relative).read_bytes() + original = (ROOT / "upstream" / relative).read_bytes() + if digest(patched) != expected["patched_sha256"]: + raise RuntimeError(f"Packaged patched source changed: {name}") + if digest(original) != expected["original_sha256"]: + raise RuntimeError(f"Packaged original source changed: {name}") + current = target.read_bytes() + current_hash = digest(current) + if current_hash == expected["original_sha256"]: + state = "original" + elif current_hash == expected["patched_sha256"]: + state = "patched" + else: + raise RuntimeError(f"Installed source does not match this patch: {target} ({current_hash})") + backup = target.with_name(target.name + BACKUP_SUFFIX) + if backup.exists() and digest(backup.read_bytes()) != expected["original_sha256"]: + raise RuntimeError(f"Existing backup does not match original source: {backup}") + planned.append((target, backup, current, original, patched, state)) + + if args.check: + for target, _, _, _, _, state in planned: + print(f"{state:8} {target}") + return + + # All source and backup hashes have been validated before the first write. + for target, backup, current, original, patched, state in planned: + mode = stat.S_IMODE(target.stat().st_mode) + if args.apply and state == "original": + if not backup.exists(): + atomic_write(backup, current, mode) + atomic_write(target, patched, mode) + print(f"PATCHED {target}") + elif args.rollback and state == "patched": + # The packaged original is identical to the verified backup. + atomic_write(target, original, mode) + print(f"RESTORED {target}") + else: + print(f"UNCHANGED {target}") + + +if __name__ == "__main__": + try: + run() + except (OSError, ValueError, RuntimeError) as error: + print(f"ERROR: {error}", file=sys.stderr) + sys.exit(1) diff --git a/src/vllm/model_executor/layers/quantization/fp8.py b/src/vllm/model_executor/layers/quantization/fp8.py new file mode 100644 index 0000000..4a76eca --- /dev/null +++ b/src/vllm/model_executor/layers/quantization/fp8.py @@ -0,0 +1,880 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import TYPE_CHECKING, Any + +import torch +from torch.utils._python_dispatch import TorchDispatchMode + +import vllm.envs as envs +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.config import get_current_vllm_config +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.logger import init_logger +from vllm.model_executor.kernels.linear import ( + init_fp8_linear_kernel, +) +from vllm.model_executor.kernels.linear.scaled_mm import ( + CutlassFP8ScaledMMLinearKernel, + MarlinFP8ScaledMMLinearKernel, +) +from vllm.model_executor.layers.attention import Attention +from vllm.model_executor.layers.fused_moe import ( + FusedMoEMethodBase, + FusedMoeWeightScaleSupported, + RoutedExperts, + SharedExperts, + UnquantizedFusedMoEMethod, +) +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEQuantConfig, +) +from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( + Fp8MoeBackend, + convert_to_fp8_moe_kernel_format, + make_fp8_moe_kernel, + make_fp8_moe_quant_config, + select_fp8_moe_backend, +) +from vllm.model_executor.layers.linear import ( + LinearBase, + LinearMethodBase, + UnquantizedLinearMethod, +) +from vllm.model_executor.layers.quantization import QuantizationMethods +from vllm.model_executor.layers.quantization.base_config import ( + QuantizationConfig, + QuantizeMethodBase, +) +from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod +from vllm.model_executor.layers.quantization.utils.fp8_utils import ( + create_fp8_input_scale, + create_fp8_scale_parameter, + create_fp8_weight_parameter, + process_fp8_input_tensor_strategy_moe, + process_fp8_weight_tensor_strategy, + process_fp8_weight_tensor_strategy_moe, + validate_fp8_block_shape, +) +from vllm.model_executor.layers.quantization.utils.marlin_utils import ( + get_marlin_input_dtype, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + GroupShape, + create_fp8_quant_key, + is_layer_skipped, + kFp8Dynamic128Sym, + kFp8DynamicTensorSym, + kFp8DynamicTokenSym, + kFp8Static128BlockSym, + kFp8StaticTensorSym, +) +from vllm.model_executor.layers.quantization.utils.w8a8_utils import ( + cutlass_block_fp8_supported, + cutlass_fp8_supported, + normalize_e4m3fn_to_e4m3fnuz, +) +from vllm.model_executor.parameter import ( + BlockQuantScaleParameter, + PerTensorScaleParameter, +) +from vllm.model_executor.utils import replace_parameter, set_weight_attrs +from vllm.platforms import current_platform + +if current_platform.is_ppu(): + from vllm.utils.ppu_deep_gemm import ( + is_deep_gemm_supported, + ) +else: + from vllm.utils.deep_gemm import ( + is_deep_gemm_supported, + ) + +if TYPE_CHECKING: + from vllm.model_executor.models.utils import WeightsMapper + +ACTIVATION_SCHEMES = ["static", "dynamic"] + +logger = init_logger(__name__) + + +class Fp8Config(QuantizationConfig): + """Config class for FP8.""" + + def __init__( + self, + is_checkpoint_fp8_serialized: bool = False, + activation_scheme: str = "dynamic", + ignored_layers: list[str] | None = None, + weight_block_size: list[int] | None = None, + store_dtype: str | None = None, + ) -> None: + super().__init__() + + self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized + + if activation_scheme not in ACTIVATION_SCHEMES: + raise ValueError(f"Unsupported activation scheme {activation_scheme}") + self.activation_scheme = activation_scheme + self.ignored_layers = ignored_layers or [] + self.store_dtype = store_dtype + if weight_block_size is not None: + if not is_checkpoint_fp8_serialized: + raise ValueError( + "The block-wise quantization only supports fp8-serialized " + "checkpoint for now." + ) + if len(weight_block_size) != 2: + raise ValueError( + "The quantization block size of weight must have 2 " + f"dimensions, but got {len(weight_block_size)} dimensions" + ) + if activation_scheme != "dynamic": + raise ValueError( + "The block-wise quantization only supports " + "dynamic activation scheme for now, but got " + f"{activation_scheme} activation scheme." + ) + self.weight_block_size = weight_block_size + self.use_deep_gemm: bool | None = None + + @classmethod + def get_name(cls) -> QuantizationMethods: + return "fp8" + + @classmethod + def get_supported_act_dtypes(cls) -> list[torch.dtype]: + return [torch.bfloat16, torch.half] + + @classmethod + def get_min_capability(cls) -> int: + return 75 + + @classmethod + def get_config_filenames(cls) -> list[str]: + return [] + + def apply_vllm_mapper(self, hf_to_vllm_mapper: "WeightsMapper"): + if self.ignored_layers is not None: + self.ignored_layers = hf_to_vllm_mapper.apply_list(self.ignored_layers) + + @classmethod + def from_config(cls, config: dict[str, Any]) -> "Fp8Config": + quant_method = cls.get_from_keys(config, ["quant_method"]) + is_checkpoint_fp8_serialized = "fp8" in quant_method + activation_scheme = cls.get_from_keys(config, ["activation_scheme"]) + ignored_layers = cls.get_from_keys_or(config, ["ignored_layers"], None) + weight_block_size = cls.get_from_keys_or(config, ["weight_block_size"], None) + store_dtype = cls.get_from_keys_or(config, ["store_dtype"], None) + if not ignored_layers: + ignored_layers = cls.get_from_keys_or( + config, ["modules_to_not_convert"], None + ) + return cls( + is_checkpoint_fp8_serialized=is_checkpoint_fp8_serialized, + activation_scheme=activation_scheme, + ignored_layers=ignored_layers, + weight_block_size=weight_block_size, + store_dtype=store_dtype, + ) + + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> "QuantizeMethodBase | None": + if isinstance(layer, LinearBase): + if is_layer_skipped( + prefix=prefix, + ignored_layers=self.ignored_layers, + fused_mapping=self.packed_modules_mapping, + ): + return UnquantizedLinearMethod() + if not self.is_checkpoint_fp8_serialized: + from vllm.model_executor.layers.quantization.online.fp8 import ( + Fp8PerTensorOnlineLinearMethod, + ) + + online_method = Fp8PerTensorOnlineLinearMethod() + online_method.marlin_input_dtype = get_marlin_input_dtype(prefix) + return online_method + else: + offline_method = Fp8LinearMethod(self) + offline_method.marlin_input_dtype = get_marlin_input_dtype(prefix) + return offline_method + elif isinstance(layer, RoutedExperts): + if is_layer_skipped( + prefix=prefix, + ignored_layers=self.ignored_layers, + fused_mapping=self.packed_modules_mapping, + ): + return UnquantizedFusedMoEMethod(layer.moe_config) + if self.store_dtype == "mxfp4": + from vllm.model_executor.layers.quantization.mxfp4 import ( + Mxfp4MoEMethod, + ) + + return Mxfp4MoEMethod(layer.moe_config) + if self.is_checkpoint_fp8_serialized: + return Fp8MoEMethod(self, layer) + else: + from vllm.model_executor.layers.quantization.online.fp8 import ( + Fp8PerTensorOnlineMoEMethod, + ) + + return Fp8PerTensorOnlineMoEMethod(layer=layer) + elif isinstance(layer, Attention): + return Fp8KVCacheMethod(self) + return None + + @staticmethod + def get_cache_scale_mapper() -> "WeightsMapper": + """Map compressed-tensors KV-cache scale names to vLLM names.""" + from vllm.model_executor.models.utils import WeightsMapper + + orig_to_new_suffix = { + ".k_proj.output_scale": ".attn.k_scale", + ".v_proj.output_scale": ".attn.v_scale", + ".q_proj.output_scale": ".attn.q_scale", + ".self_attn.prob_output_scale": ".self_attn.attn.prob_scale", + } + cache_scale_mapper = WeightsMapper(orig_to_new_suffix=orig_to_new_suffix) + return cache_scale_mapper | QuantizationConfig.get_cache_scale_mapper() + + +class CopyNumelCounter(TorchDispatchMode): + """ + Tracks total number of elements modified with `copy_`. Useful for keeping + track of weight loading where underlying weights can be arbitrarily + transformed (such as with `narrow`) before calling copy. + """ + + def __init__(self): + super().__init__() + self.copied_numel = 0 + + def __torch_dispatch__(self, func, types, args=(), kwargs=None): + if kwargs is None: + kwargs = {} + out = func(*args, **kwargs) + if func == torch.ops.aten.copy_.default: + self.copied_numel += args[0].numel() + return out + + +def _copy_missing_attrs(old: torch.Tensor, new: torch.Tensor) -> None: + """Copies any attrs present in `old` but not in `new` to `new`""" + new_attrs = set(dir(new)) + attrs_to_set = {} + for attr in dir(old): + if attr not in new_attrs: + attrs_to_set[attr] = getattr(old, attr) + set_weight_attrs(new, attrs_to_set) + + +class Fp8LinearMethod(LinearMethodBase): + """Linear method for FP8. + Supports loading FP8 checkpoints with static weight scale and + dynamic/static activation scale. + + Limitations: + 1. Only support float8_e4m3fn data type due to the limitation of + torch._scaled_mm (https://github.com/pytorch/pytorch/blob/2e48b39603411a41c5025efbe52f89560b827825/aten/src/ATen/native/cuda/Blas.cpp#L854-L856) + + Args: + quant_config: The quantization config. + """ + + def __init__(self, quant_config: Fp8Config): + self.quant_config = quant_config + self.is_scale_e8m0 = getattr(quant_config, "is_scale_e8m0", False) + self.cutlass_block_fp8_supported = cutlass_block_fp8_supported() + self.out_dtype = torch.get_default_dtype() + self.input_dtype = get_current_vllm_config().model_config.dtype + + # For GPUs that lack FP8 hardware support, we can leverage the Marlin + # kernel for fast weight-only FP8 quantization + self.marlin_input_dtype = None + self.use_marlin = False + + if self.quant_config.use_deep_gemm is not None: + self.use_deep_gemm = self.quant_config.use_deep_gemm + else: + self.use_deep_gemm = is_deep_gemm_supported() + + self.weight_block_size = self.quant_config.weight_block_size + self.block_quant = self.weight_block_size is not None + self.act_q_static = self.quant_config.activation_scheme == "static" + + if self.block_quant: + assert not self.act_q_static + assert self.weight_block_size is not None + + self.activation_quant_key = create_fp8_quant_key( + static=self.act_q_static, + group_shape=GroupShape(1, self.weight_block_size[0]), + ) + self.weight_quant_key = create_fp8_quant_key( + static=True, group_shape=GroupShape(*self.weight_block_size) + ) + else: + self.weight_quant_key = kFp8StaticTensorSym + # Use per-token quantization for better perf if dynamic and cutlass + if self.act_q_static: + self.activation_quant_key = kFp8StaticTensorSym + elif cutlass_fp8_supported(): + self.activation_quant_key = kFp8DynamicTokenSym + else: + self.activation_quant_key = kFp8DynamicTensorSym + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: list[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + output_size_per_partition = sum(output_partition_sizes) + weight_loader = extra_weight_attrs.get("weight_loader") + layer.logical_widths = output_partition_sizes + layer.input_size_per_partition = input_size_per_partition + layer.output_size_per_partition = output_size_per_partition + layer.orig_dtype = params_dtype + layer.weight_block_size = None + + if self.block_quant: + assert self.weight_block_size is not None + layer.weight_block_size = self.weight_block_size + validate_fp8_block_shape( + layer, + input_size, + output_size, + input_size_per_partition, + output_partition_sizes, + self.weight_block_size, + ) + + weight = create_fp8_weight_parameter( + output_size_per_partition, input_size_per_partition, weight_loader + ) + layer.register_parameter("weight", weight) + + # WEIGHT SCALE + if not self.block_quant: + scale = create_fp8_scale_parameter( + PerTensorScaleParameter, + output_partition_sizes, + input_size_per_partition, + None, + weight_loader, + ) + layer.register_parameter("weight_scale", scale) + else: + assert not self.act_q_static + assert self.weight_block_size is not None + scale = create_fp8_scale_parameter( + BlockQuantScaleParameter, + output_partition_sizes, + input_size_per_partition, + self.weight_block_size, + weight_loader, + scale_dtype=(torch.float8_e8m0fnu if self.is_scale_e8m0 else None), + ) + # The weight_scale_inv name is intentional for deepseekv3 + layer.register_parameter("weight_scale_inv", scale) + + # INPUT ACTIVATION SCALE + if self.act_q_static: + scale = create_fp8_input_scale(output_partition_sizes, weight_loader) + set_weight_attrs(scale, {"scale_type": "input_scale"}) + layer.register_parameter("input_scale", scale) + + self.fp8_linear = init_fp8_linear_kernel( + activation_quant_key=self.activation_quant_key, + weight_quant_key=self.weight_quant_key, + weight_shape=layer.weight.shape, + input_dtype=self.input_dtype, + out_dtype=self.out_dtype, + module_name=self.__class__.__name__, + ) + + self.use_marlin = isinstance(self.fp8_linear, MarlinFP8ScaledMMLinearKernel) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + # Opt-in fallback before Marlin repacks wo_a. + import os + if (os.environ.get("VLLM_PPU_BF16_O_PROJ") == "1" + and current_platform.is_ppu() + and getattr(layer, "is_bmm", False) + and self.block_quant): + from vllm.models.deepseek_v4.ppu.ops.o_proj import prepare_bf16_o_proj_weight + prepare_bf16_o_proj_weight(layer) + return + if self.use_marlin: + if not self.block_quant: + # Canonicalize to (K, N) for the kernel. + replace_parameter(layer, "weight", layer.weight.t()) + # Only Marlin kernels support `marlin_input_dtype`; guard to avoid + # AttributeError if backend selection changes. + if hasattr(self.fp8_linear, "marlin_input_dtype"): + self.fp8_linear.marlin_input_dtype = self.marlin_input_dtype + self.fp8_linear.process_weights_after_loading(layer) + return + + input_scale = None + # TODO(rob): refactor block quant into separate class. + if self.block_quant: + assert not self.act_q_static + + # If checkpoint not serialized fp8, quantize the weights. + else: + # If checkpoint is fp8 per-tensor, handle that there are N scales for N + # shards in a fused module + weight = layer.weight + weight_scale = layer.weight_scale + + # If using w8a8, torch._scaled_mm needs per tensor, so + # requantize the logical shards as a single weight. + weight, weight_scale, input_scale = process_fp8_weight_tensor_strategy( + weight, + weight_scale, + layer.logical_widths, + getattr(layer, "input_scale", None), + ) + if self.act_q_static: + assert input_scale is not None + input_scale = input_scale.max() + weight = weight.t() + + # Update layer with new values. + replace_parameter(layer, "weight", weight.data) + replace_parameter(layer, "weight_scale", weight_scale.data) + + if input_scale is not None: + replace_parameter(layer, "input_scale", input_scale) + else: + layer.input_scale = None + + self.fp8_linear.process_weights_after_loading(layer) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + # if batch invariant mode is enabled, prefer direct FP8 path + # we will use BF16 dequant when direct FP8 is not supported. + if envs.VLLM_BATCH_INVARIANT: + if self.block_quant: + assert self.weight_block_size is not None + return self.fp8_linear.apply_weights( + layer, + x, + bias, + ) + else: + if isinstance(self.fp8_linear, CutlassFP8ScaledMMLinearKernel): + return self.fp8_linear.apply_weights(layer, x, bias) + + # per-tensor/channel: dequant to BF16 and run GEMM + weight_fp8 = layer.weight.to(torch.bfloat16) + weight_scale = layer.weight_scale.to(torch.bfloat16) + if weight_scale.numel() == 1: + # Per-tensor: simple scalar multiplication + weight_bf16 = weight_fp8 * weight_scale + else: + # Multiple scales (fused modules like QKV) + # Try to infer correct broadcasting + # weight is [K, N], scale could be [num_logical_weights] + # Need to figure out how to broadcast - for now just try + # direct multiplication + if ( + weight_scale.dim() == 1 + and weight_scale.shape[0] == weight_fp8.shape[0] + ): + # Per-row scaling + weight_bf16 = weight_fp8 * weight_scale.unsqueeze(1) + else: + # Fallback + weight_bf16 = weight_fp8 * weight_scale + return torch.nn.functional.linear(x, weight_bf16.t(), bias) + + return self.fp8_linear.apply_weights(layer, x, bias) + + +class Fp8MoEMethod(FusedMoEMethodBase): + """MoE method for FP8. + Supports loading FP8 checkpoints with static weight scale and + dynamic/static activation scale. + + Also supports loading quantized FP16/BF16 model checkpoints with dynamic + activation scaling. The weight scaling factor will be initialized after + the model weights are loaded. + + Args: + quant_config: The quantization config. + """ + + def __init__(self, quant_config: Fp8Config, layer: RoutedExperts): + super().__init__(layer.moe_config) + self.quant_config = quant_config + self.weight_block_size = self.quant_config.weight_block_size + self.block_quant: bool = self.weight_block_size is not None + self.weight_scale_name = ( + "weight_scale_inv" if self.block_quant else "weight_scale" + ) + + # Set weight key and activation key for kernel compatibility + if self.block_quant: + weight_key = kFp8Static128BlockSym + activation_key = kFp8Dynamic128Sym + else: + weight_key = kFp8StaticTensorSym + activation_key = ( + kFp8StaticTensorSym + if self.quant_config.activation_scheme == "static" + else kFp8DynamicTensorSym + ) + + # Select Fp8 MoE backend + self.fp8_backend, self.experts_cls = select_fp8_moe_backend( + config=self.moe, + weight_key=weight_key, + activation_key=activation_key, + allow_vllm_cutlass=False, + ) + + def create_weights( + self, + layer: RoutedExperts, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + layer.num_experts = num_experts + layer.orig_dtype = params_dtype + layer.weight_block_size = None + + assert self.quant_config.is_checkpoint_fp8_serialized + params_dtype = torch.float8_e4m3fn + + if self.block_quant: + assert self.weight_block_size is not None + layer.weight_block_size = self.weight_block_size + tp_size = get_tensor_model_parallel_world_size() + block_n, block_k = ( + self.weight_block_size[0], + self.weight_block_size[1], + ) + # NOTE: To ensure proper alignment of the block-wise quantization + # scales, the output_size of the weights for both the gate and up + # layers must be divisible by block_n. + # Required by column parallel or enabling merged weights + if intermediate_size_per_partition % block_n != 0: + raise ValueError( + f"The output_size of gate's and up's weight = " + f"{intermediate_size_per_partition} is not divisible by " + f"weight quantization block_n = {block_n}." + ) + if tp_size > 1 and intermediate_size_per_partition % block_k != 0: + # Required by row parallel + raise ValueError( + f"The input_size of down's weight = " + f"{intermediate_size_per_partition} is not divisible by " + f"weight quantization block_k = {block_k}." + ) + + # WEIGHTS + w13_weight = torch.nn.Parameter( + torch.empty( + num_experts, + 2 * intermediate_size_per_partition, + hidden_size, + dtype=params_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight", w13_weight) + set_weight_attrs(w13_weight, extra_weight_attrs) + + w2_weight = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + intermediate_size_per_partition, + dtype=params_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight", w2_weight) + set_weight_attrs(w2_weight, extra_weight_attrs) + + # BIASES (for models like GPT-OSS that have biased MoE) + if self.moe.has_bias: + w13_bias = torch.nn.Parameter( + torch.zeros( + num_experts, + 2 * intermediate_size_per_partition, + dtype=layer.orig_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w13_bias", w13_bias) + set_weight_attrs(w13_bias, extra_weight_attrs) + w2_bias = torch.nn.Parameter( + torch.zeros(num_experts, hidden_size, dtype=layer.orig_dtype), + requires_grad=False, + ) + layer.register_parameter("w2_bias", w2_bias) + set_weight_attrs(w2_bias, extra_weight_attrs) + + # WEIGHT_SCALES + if not self.block_quant: + # For per-tensor quant, the scales are per expert and weight. + w13_scale_data = torch.ones(num_experts, 2, dtype=torch.float32) + w2_scale_data = torch.ones(num_experts, dtype=torch.float32) + else: + # For block quant, the scales are per block (typically 128x128). + w13_scale_data = torch.ones( + num_experts, + 2 * ((intermediate_size_per_partition + block_n - 1) // block_n), + (hidden_size + block_k - 1) // block_k, + dtype=torch.float32, + ) + w2_scale_data = torch.ones( + num_experts, + (hidden_size + block_n - 1) // block_n, + (intermediate_size_per_partition + block_k - 1) // block_k, + dtype=torch.float32, + ) + w13_weight_scale = torch.nn.Parameter(w13_scale_data, requires_grad=False) + w2_weight_scale = torch.nn.Parameter(w2_scale_data, requires_grad=False) + # Note: name is weight_scale for tensor, weight_scale_inv for block. + layer.register_parameter(f"w13_{self.weight_scale_name}", w13_weight_scale) + layer.register_parameter(f"w2_{self.weight_scale_name}", w2_weight_scale) + + # Add the quantization method used (per tensor/grouped/channel) + # to ensure the weight scales are loaded in properly + extra_weight_attrs.update( + {"quant_method": FusedMoeWeightScaleSupported.BLOCK.value} + if self.block_quant + else {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value} + ) + set_weight_attrs(w13_weight_scale, extra_weight_attrs) + set_weight_attrs(w2_weight_scale, extra_weight_attrs) + + # INPUT_SCALES + if self.quant_config.activation_scheme == "static": + assert not self.block_quant + w13_input_scale = torch.nn.Parameter( + torch.ones(num_experts, dtype=torch.float32), requires_grad=False + ) + layer.register_parameter("w13_input_scale", w13_input_scale) + set_weight_attrs(w13_input_scale, extra_weight_attrs) + + w2_input_scale = torch.nn.Parameter( + torch.ones(num_experts, dtype=torch.float32), requires_grad=False + ) + layer.register_parameter("w2_input_scale", w2_input_scale) + set_weight_attrs(w2_input_scale, extra_weight_attrs) + + else: + layer.w13_input_scale = None + layer.w2_input_scale = None + + def _setup_kernel( + self, + layer: RoutedExperts, + w13: torch.Tensor, + w2: torch.Tensor, + w13_scale: torch.Tensor, + w2_scale: torch.Tensor, + w13_input_scale: torch.Tensor | None, + w2_input_scale: torch.Tensor | None, + ) -> None: + # Shuffle weights to runtime format. + w13, w2, w13_scale, w2_scale = convert_to_fp8_moe_kernel_format( + fp8_backend=self.fp8_backend, + layer=layer, + w13=w13, + w2=w2, + w13_scale=w13_scale, + w2_scale=w2_scale, + w13_input_scale=w13_input_scale, + w2_input_scale=w2_input_scale, + ) + + # Replace parameters with updated versions. Note that this helper + # function ensures the replacement is compatible with RL weight reloads. + replace_parameter(layer, "w13_weight", w13) + replace_parameter(layer, "w2_weight", w2) + replace_parameter(layer, f"w13_{self.weight_scale_name}", w13_scale) + replace_parameter(layer, f"w2_{self.weight_scale_name}", w2_scale) + + # AITER backend requires weights to be marked as shuffled. + if self.fp8_backend == Fp8MoeBackend.AITER: + layer.w13_weight.is_shuffled = True + layer.w2_weight.is_shuffled = True + + self.moe_quant_config = self.get_fused_moe_quant_config(layer) + assert self.moe_quant_config is not None + assert self.experts_cls is not None + self.moe_kernel = make_fp8_moe_kernel( + moe_quant_config=self.moe_quant_config, + moe_config=self.moe, + fp8_backend=self.fp8_backend, + experts_cls=self.experts_cls, + routing_tables=layer._expert_routing_tables(), + layer=layer, + ) + + def process_weights_after_loading(self, layer: RoutedExperts) -> None: + # Allow for accessing weights and scales in standard way. + w13 = layer.w13_weight + w2 = layer.w2_weight + w13_scale = getattr(layer, f"w13_{self.weight_scale_name}") + w2_scale = getattr(layer, f"w2_{self.weight_scale_name}") + w13_input_scale = layer.w13_input_scale + w2_input_scale = layer.w2_input_scale + + # MI300x and MI325x use FNUZ format for FP8. Convert if needed. + if current_platform.is_fp8_fnuz(): + w13, w13_scale, w13_input_scale = normalize_e4m3fn_to_e4m3fnuz( + w13, + w13_scale, + w13_input_scale, + ) + w2, w2_scale, w2_input_scale = normalize_e4m3fn_to_e4m3fnuz( + w2, + w2_scale, + w2_input_scale, + ) + + # Per tensor kernels require single activation scale. Use the max. + if self.quant_config.activation_scheme == "static": + assert not self.block_quant + assert w13_input_scale is not None and w2_input_scale is not None + w13_input_scale, w2_input_scale = process_fp8_input_tensor_strategy_moe( + w13_input_scale, w2_input_scale + ) + replace_parameter(layer, "w13_input_scale", w13_input_scale) + replace_parameter(layer, "w2_input_scale", w2_input_scale) + + # Per tensor kernels require single weight scale for w13 per expert, but + # on disk there is a scale for w1 and w3. Use the max to requantize. + if not self.block_quant: + shard_size = layer.intermediate_size_per_partition + w13, w13_scale = process_fp8_weight_tensor_strategy_moe( + w13, w13_scale, shard_size, layer.local_num_experts + ) + + # Shuffle weights to runtime format and setup kernel. + self._setup_kernel( + layer, w13, w2, w13_scale, w2_scale, w13_input_scale, w2_input_scale + ) + + def maybe_make_prepare_finalize( + self, + routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, + ) -> mk.FusedMoEPrepareAndFinalizeModular | None: + raise ValueError( + f"{self.__class__.__name__} uses the new modular kernel initialization " + "logic. This function should not be called." + ) + + def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfig: + w1_scale = getattr(layer, f"w13_{self.weight_scale_name}") + w2_scale = getattr(layer, f"w2_{self.weight_scale_name}") + a1_scale = layer.w13_input_scale + a2_scale = layer.w2_input_scale + + quant_config = make_fp8_moe_quant_config( + fp8_backend=self.fp8_backend, + w1_scale=w1_scale, + w2_scale=w2_scale, + a1_scale=a1_scale, + a2_scale=a2_scale, + block_shape=self.weight_block_size, + swiglu_limit=getattr(layer, "swiglu_limit", None), + gemm1_alpha=getattr(layer, "swiglu_alpha", None), + gemm1_beta=getattr(layer, "swiglu_beta", None), + layer=layer, + ) + + # Inject biases into the quant config if the model has them + # (e.g. GPT-OSS biased MoE) + if quant_config is not None and self.moe.has_bias: + w13_bias = getattr(layer, "w13_bias", None) + w2_bias = getattr(layer, "w2_bias", None) + if w13_bias is not None: + quant_config._w1.bias = w13_bias + if w2_bias is not None: + quant_config._w2.bias = w2_bias + + return quant_config + + @property + def supports_eplb(self) -> bool: + return True + + def apply_monolithic( + self, + layer: RoutedExperts, + x: torch.Tensor, + router_logits: torch.Tensor, + input_ids: torch.Tensor | None = None, + ) -> torch.Tensor: + assert self.is_monolithic + assert self.moe_kernel is not None + return self.moe_kernel.apply_monolithic( + x, + layer.w13_weight, + layer.w2_weight, + router_logits, + activation=layer.activation, + global_num_experts=layer.global_num_experts, + expert_map=layer.expert_map, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + num_expert_group=layer.num_expert_group, + topk_group=layer.topk_group, + e_score_correction_bias=layer.e_score_correction_bias, + routed_scaling_factor=layer.routed_scaling_factor, + ) + + def apply( + self, + layer: RoutedExperts, + x: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + shared_experts: SharedExperts | None, + shared_experts_input: torch.Tensor | None, + ) -> torch.Tensor: + assert not self.is_monolithic + assert self.moe_kernel is not None + return self.moe_kernel.apply( + x, + layer.w13_weight, + layer.w2_weight, + topk_weights, + topk_ids, + activation=layer.activation, + global_num_experts=layer.global_num_experts, + expert_map=layer.expert_map, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + shared_experts=shared_experts, + shared_experts_input=shared_experts_input, + ) + + +class Fp8KVCacheMethod(BaseKVCacheMethod): + """ + Supports loading kv-cache scaling factors from FP8 checkpoints. + """ + + def __init__(self, quant_config: Fp8Config): + super().__init__(quant_config) diff --git a/src/vllm/models/deepseek_v4/attention.py b/src/vllm/models/deepseek_v4/attention.py new file mode 100644 index 0000000..985654a --- /dev/null +++ b/src/vllm/models/deepseek_v4/attention.py @@ -0,0 +1,844 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +DeepseekV4 MLA Attention Layer +""" + +from abc import ABC, abstractmethod +from collections.abc import Callable +from typing import TYPE_CHECKING, Any, ClassVar, cast + +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers import DeepseekV2Config, DeepseekV3Config + +import vllm.envs as envs +from vllm.compilation.breakable_cudagraph import eager_break_during_capture +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer +from vllm.models.deepseek_v4.common.ops import ( + fused_indexer_q_rope_quant, + fused_q_kv_rmsnorm, +) +from vllm.utils.deep_gemm import fp8_einsum +from vllm.utils.torch_utils import direct_register_custom_op + +if TYPE_CHECKING: + from vllm.v1.attention.backends.mla.sparse_swa import ( + DeepseekSparseSWAMetadata, + ) + +from vllm.config import ( + CacheConfig, + VllmConfig, + get_current_vllm_config, +) +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.forward_context import get_forward_context +from vllm.logger import init_logger +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.models.utils import extract_layer_index +from vllm.models.deepseek_v4.common.rope import build_deepseek_v4_rope +from vllm.models.deepseek_v4.compressor import DeepseekCompressor +from vllm.utils.multi_stream_utils import ( + execute_in_parallel, + maybe_execute_in_parallel, +) +from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata +from vllm.v1.attention.backends.mla.indexer import ( + DeepseekV4IndexerBackend, + get_max_prefill_buffer_size, +) +from vllm.v1.attention.backends.mla.sparse_swa import DeepseekV4SWACache +from vllm.v1.kv_cache_interface import ( + KVCacheSpec, + MLAAttentionSpec, + get_kv_quant_mode, +) + +logger = init_logger(__name__) + + +def _resolve_dsv4_kv_cache_dtype( + use_fp8_ds_mla_layout: bool, + kv_cache_dtype: str, + cache_config: CacheConfig | None, +) -> tuple[str, torch.dtype]: + """Map ``(layout, --kv-cache-dtype)`` to ``(cache_dtype_str, torch_dtype)``. + + Both layouts are paged; they differ in the per-token block format. The + ``fp8_ds_mla`` format is UE8M0 block-scaled fp8 packed as ``uint8`` (the + canonical ``fp8_ds_mla`` string is written back onto ``cache_config`` so the + page-size specs pick the 576B per-token slot). Plain-row backends store each + token's KV row in its element dtype: bf16 or per-tensor FP8 E4M3. + """ + if use_fp8_ds_mla_layout: + # fp8_ds_mla block format: UE8M0 block-scaled fp8 packed as uint8. + assert kv_cache_dtype.startswith("fp8"), ( + f"DeepseekV4 fp8_ds_mla layout only supports fp8 kv-cache, " + f"got {kv_cache_dtype}" + ) + if kv_cache_dtype != "fp8_ds_mla": + if cache_config is not None: + cache_config.cache_dtype = "fp8_ds_mla" + kv_cache_dtype = "fp8_ds_mla" + logger.info_once("Using DeepSeek's fp8_ds_mla KV cache format.") + return kv_cache_dtype, torch.uint8 + + # Plain bf16 / per-tensor fp8 KV row (FlashInfer). + if kv_cache_dtype.startswith("fp8"): + return kv_cache_dtype, torch.float8_e4m3fn + # auto / bfloat16 -> plain bf16 KV row. + return kv_cache_dtype, torch.bfloat16 + + +class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): + """DeepseekV4 MLA attention layer. + + The platform-specific sparse-MLA forward (``forward_mqa`` / + ``get_padded_num_q_heads`` / ``_o_proj`` / ``backend_cls``) is provided by a + subclass — ``DeepseekV4FlashMLAAttention`` / + ``DeepseekV4FlashInferSM120Attention`` / + ``DeepseekV4FlashInferMLAAttention`` (CUDA) or + ``DeepseekV4ROCMAiterMLAAttention`` (ROCm) — selected by the platform-specific + deepseek_v4 model module. The base is never instantiated directly. + """ + + # Provided by the platform subclass. + backend_cls: ClassVar[type[AttentionBackend]] + # KV-cache per-token block format (both layouts are paged). True (default) + # = fp8_ds_mla (UE8M0 block-scaled fp8 packed as uint8); False = plain + # bf16 / per-tensor fp8 KV row. Backends can override the instance hook when + # a single attention class dispatches across arch-specific layouts. + use_fp8_ds_mla_layout: ClassVar[bool] = True + # Prefill is processed in fixed-size chunks; this bounds the bf16 kv-gather + # workspace allocated in _forward_prefill and is also read by the dummy-run + # path to pre-reserve that workspace. + PREFILL_CHUNK_SIZE: ClassVar[int] = 4 + + @classmethod + @abstractmethod + def get_padded_num_q_heads(cls, num_heads: int) -> int: + """Q head count the q/output buffers are allocated at. + + The layer allocates the q/output buffers at + ``[N, get_padded_num_q_heads(n_local_heads), head_dim]``. Must satisfy + ``result >= num_heads``. Backends with no padding constraint return + ``num_heads``. + """ + raise NotImplementedError + + @abstractmethod + def forward_mqa( + self, + q: torch.Tensor, + kv: torch.Tensor, + positions: torch.Tensor, + output: torch.Tensor, + ) -> None: + """Platform-specific sparse MLA forward; writes attention into ``output``.""" + raise NotImplementedError + + @abstractmethod + def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + """Inverse-RoPE + wo_a + wo_b output projection (platform-specific).""" + raise NotImplementedError + + def _uses_fp8_ds_mla_layout(self) -> bool: + """Return whether this instance stores fp8 KV in fp8_ds_mla layout.""" + return self.use_fp8_ds_mla_layout + + def __init__( + self, + vllm_config: VllmConfig, + prefix: str, + topk_indices_buffer: torch.Tensor | None = None, + aux_stream_list: list[torch.cuda.Stream] | None = None, + ) -> None: + super().__init__() + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + cache_config = vllm_config.cache_config + tp_size = get_tensor_model_parallel_world_size() + layer_id = extract_layer_index(prefix) + # Diagnostic PPU path: run attention/indexer/compressor on one stream. + import os + from vllm.platforms import current_platform + if (current_platform.is_ppu() + and os.environ.get("VLLM_PPU_SERIAL_ATTENTION") == "1"): + aux_stream_list = None + if layer_id == 0: + logger.info("PPU attention auxiliary streams disabled") + + self.prefix = prefix # Alias for compatibility with compressor + self.hidden_size = config.hidden_size + self.n_heads = config.num_attention_heads + assert self.n_heads % tp_size == 0 + self.n_local_heads = self.n_heads // tp_size + self.q_lora_rank = config.q_lora_rank + self.o_lora_rank = config.o_lora_rank + self.head_dim = config.head_dim + self.rope_head_dim = config.qk_rope_head_dim + self.nope_head_dim = self.head_dim - self.rope_head_dim + self.n_groups = config.o_groups + self.n_local_groups = self.n_groups // tp_size + self.window_size = config.sliding_window + # NOTE(zyongye) Compress ratio can't be 0 + # we do this for because MTP layer is not included + # in the compress ratio list + if layer_id < config.num_hidden_layers: + self.compress_ratio = max(1, config.compress_ratios[layer_id]) + else: + self.compress_ratio = 1 + self.eps = config.rms_norm_eps + self.scale = self.head_dim**-0.5 + + # Padded Q head count is dictated by the platform subclass. + self.padded_heads = self.get_padded_num_q_heads(self.n_local_heads) + # Sink padded to the same head count, initialized to -inf (no sink + # effect). Weight loading fills the first n_local_heads slots. + self.attn_sink = nn.Parameter( + torch.full((self.padded_heads,), -float("inf"), dtype=torch.float32), + requires_grad=False, + ) + + self.fused_wqa_wkv = MergedColumnParallelLinear( + self.hidden_size, + [self.q_lora_rank, self.head_dim], + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.fused_wqa_wkv", + disable_tp=True, # fused ReplicatedLinear + ) + self.q_norm = RMSNorm(self.q_lora_rank, self.eps) + self.wq_b = ColumnParallelLinear( + self.q_lora_rank, + self.n_heads * self.head_dim, + bias=False, + quant_config=quant_config, + return_bias=False, + prefix=f"{prefix}.wq_b", + ) + + self.kv_norm = RMSNorm(self.head_dim, self.eps) + # wo_a uses the standard quant_config for all quantization schemes + # (FP8 blockwise, FP8 channelwise, INT8 channelwise). The _o_proj + # method in the platform subclass dispatches at runtime based on + # weight dtype and scale attribute: + # - int8 + weight_scale → INT8 channelwise einsum + # - fp8 + weight_scale → FP8 channelwise einsum + # - fp8 + weight_scale_inv → FP8 blockwise einsum + self.wo_a = ColumnParallelLinear( + self.n_heads * self.head_dim // self.n_groups, + self.n_groups * self.o_lora_rank, + bias=False, + quant_config=quant_config, + return_bias=False, + prefix=f"{prefix}.wo_a", + ) + self.wo_a.is_bmm = True + self.wo_a.bmm_batch_size = self.n_local_groups + self.wo_b = RowParallelLinear( + self.n_groups * self.o_lora_rank, + self.hidden_size, + bias=False, + quant_config=quant_config, + return_bias=False, + prefix=f"{prefix}.wo_b", + ) + + # Initialize rotary embedding before the indexer/compressor consume it. + self.rotary_emb = build_deepseek_v4_rope( + config, + head_dim=self.head_dim, + rope_head_dim=self.rope_head_dim, + max_position_embeddings=config.max_position_embeddings, + compress_ratio=self.compress_ratio, + ) + self.indexer_rotary_emb = self.rotary_emb + self.topk_indices_buffer = topk_indices_buffer + + self.indexer = None + if self.compress_ratio == 4: + # Only C4A uses sparse attention and hence has indexer. + # aux_stream_list[2] is free here (outer GEMMs joined) for the inner + # overlap of wq_b+fused_indexer_q_rope_quant vs compressor. None on + # ROCm, where aux_stream_list is None. + indexer_aux_stream = ( + aux_stream_list[2] if aux_stream_list is not None else None + ) + self.indexer = DeepseekV4Indexer( + vllm_config, + config=config, + hidden_size=self.hidden_size, + q_lora_rank=self.q_lora_rank, + quant_config=quant_config, + cache_config=cache_config, + topk_indices_buffer=topk_indices_buffer, + compress_ratio=self.compress_ratio, + prefix=f"{prefix}.indexer", + aux_stream=indexer_aux_stream, + ) + + # Will be None on ROCm for now. + self.aux_stream_list = aux_stream_list + # [0]: GEMM start / post-GEMM event0. [1..3]: GEMM done events; + # [1] doubles as post-GEMM event1. Reuse is safe: GEMM fully joins + # before post-GEMM starts. + self.ln_events = [torch.cuda.Event() for _ in range(4)] + + assert cache_config is not None, "DeepseekV4 attention requires cache_config" + # ---- Attention / KV-cache setup ---- + self.max_num_batched_tokens = ( + vllm_config.scheduler_config.max_num_batched_tokens + ) + self.max_model_len = vllm_config.model_config.max_model_len + + # Resolve the kv-cache dtype from this backend's block format. The same + # resolution drives the SWA cache tensor dtype below. + self.kv_cache_dtype, self.kv_cache_torch_dtype = _resolve_dsv4_kv_cache_dtype( + self._uses_fp8_ds_mla_layout(), cache_config.cache_dtype, cache_config + ) + + self.swa_cache_layer = DeepseekV4SWACache( + head_dim=self.head_dim, + window_size=self.window_size, + dtype=self.kv_cache_torch_dtype, + prefix=f"{prefix}.swa_cache", + cache_config=cache_config, + ) + + # Register with compilation context for metadata lookup. + compilation_config = vllm_config.compilation_config + if prefix and prefix in compilation_config.static_forward_context: + raise ValueError(f"Duplicate layer name: {prefix}") + if prefix: + compilation_config.static_forward_context[prefix] = self + self.kv_cache = torch.tensor([]) + + # Create the compressor for layers with compress_ratio > 1; after the + # attention setup above so its KV-cache prefix (self.prefix) is set. + self.compressor = None + if self.compress_ratio > 1: + self.compressor = DeepseekCompressor( + vllm_config=vllm_config, + compress_ratio=self.compress_ratio, + hidden_size=self.hidden_size, + head_dim=self.head_dim, + rotate=True, + prefix=f"{prefix}.compressor", + k_cache_prefix=self.prefix, + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + llama_4_scaling: torch.Tensor | None = None, + ) -> torch.Tensor: + # Pre-allocate attention output with FlashMLA-padded head count. + # The op writes into `o_padded`; we slice to n_local_heads after. + num_tokens = hidden_states.shape[0] + o_padded = torch.empty( + (num_tokens, self.padded_heads, self.head_dim), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + + # Metadata-independent input GEMMs + RMSNorm stay in the captured + # graph; the metadata-dependent rest (q up-proj + kv-insert, indexer, + # compressor, MLA attention) runs in the eager break. + qr_kv, kv_score, indexer_kv_score, indexer_weights = ( + self.attn_gemm_parallel_execute(hidden_states) + ) + qr, kv = qr_kv.split([self.q_lora_rank, self.head_dim], dim=-1) + qr, kv = fused_q_kv_rmsnorm( + qr, + kv, + self.q_norm.weight.data, + self.kv_norm.weight.data, + self.eps, + ) + + # attention_impl is wrapped with @eager_break_during_capture: this is + # where the breakable cudagraph capture breaks (the attention op runs + # eagerly between captured graph segments). + self.attention_impl( + hidden_states, + qr, + kv, + kv_score, + indexer_kv_score, + indexer_weights, + positions, + o_padded, + ) + o = o_padded[:, : self.n_local_heads, :] + + # Inverse-RoPE + wo_a + wo_b output projection (platform-specific). + return self._o_proj(o, positions) + + def attn_gemm_parallel_execute(self, hidden_states) -> tuple[Any, ...]: + aux_streams = self.aux_stream_list + if aux_streams is not None: + assert len(aux_streams) >= 3 + aux_streams = aux_streams[:3] + + # fused_wqa_wkv (heaviest) on default; the three lighter input GEMMs + # on aux streams 0..2 when their owning module exists. ln_events[0] + # is the fan-out start event; ln_events[1..3] are per-aux done events. + # On ROCm, aux_streams is None and execute_in_parallel runs serially. + aux_fns: list[Callable[[], Any] | None] = [None, None, None] + + if self.compressor is not None: + # Local ref so the closure keeps a non-None type for mypy. + compressor = self.compressor + + def compressor_kv_score() -> torch.Tensor: + import os + from vllm.platforms import current_platform + if (current_platform.is_ppu() + and os.environ.get("VLLM_PPU_BF16_O_PROJ") == "1"): + from vllm.models.deepseek_v4.ppu.ops.o_proj import ppu_compressor_mm + return ppu_compressor_mm(hidden_states, compressor.fused_wkv_wgate.weight.T) + return torch.mm( + hidden_states, + compressor.fused_wkv_wgate.weight.T, + out_dtype=torch.float32, + ) + + aux_fns[0] = compressor_kv_score + + if self.indexer is not None: + indexer = self.indexer + + def indexer_weights_proj() -> torch.Tensor: + # ReplicatedLinear returns (output, bias); bias is None. + weights, _ = indexer.weights_proj(hidden_states) + return weights + + def indexer_compressor_kv_score() -> torch.Tensor: + import os + from vllm.platforms import current_platform + if (current_platform.is_ppu() + and os.environ.get("VLLM_PPU_BF16_O_PROJ") == "1"): + from vllm.models.deepseek_v4.ppu.ops.o_proj import ppu_compressor_mm + return ppu_compressor_mm(hidden_states, indexer.compressor.fused_wkv_wgate.weight.T) + return torch.mm( + hidden_states, + indexer.compressor.fused_wkv_wgate.weight.T, + out_dtype=torch.float32, + ) + + aux_fns[1] = indexer_weights_proj + aux_fns[2] = indexer_compressor_kv_score + + def fused_wqa_wkv() -> torch.Tensor: + # MergedColumnParallelLinear returns (output, bias); bias is None. + qr_kv, _ = self.fused_wqa_wkv(hidden_states) + return qr_kv + + qr_kv, (kv_score, indexer_weights, indexer_kv_score) = execute_in_parallel( + fused_wqa_wkv, + aux_fns, + self.ln_events[0], + self.ln_events[1:4], + aux_streams, + enable=hidden_states.shape[0] + <= envs.VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD, + ) + + return qr_kv, kv_score, indexer_kv_score, indexer_weights + + @eager_break_during_capture + def attention_impl( + self, + hidden_states: torch.Tensor, + qr: torch.Tensor, + kv: torch.Tensor, + kv_score: torch.Tensor, + indexer_kv_score: torch.Tensor, + indexer_weights: torch.Tensor, + positions: torch.Tensor, + out: torch.Tensor, # [num_tokens, padded_heads, head_dim], written in place + ) -> None: + forward_context = get_forward_context() + attn_metadata = forward_context.attn_metadata + + # wq_b + kv_insert (+ MLA compressor when an indexer is present) ride + # on the default stream so q stays on its consumer stream (forward_mqa + # downstream reads q on default). Indexer/compressor go on aux for + # overlap with default's GEMM + cache write. + if self.indexer is not None: + aux_streams = self.aux_stream_list + indexer = self.indexer + # Local ref so the closure keeps a non-None type for mypy. + assert self.compressor is not None + compressor = self.compressor + + def wq_b_kv_insert() -> torch.Tensor: + q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim) + q = self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata) + return q + + # 3-way overlap (matches TRT-LLM PR #14142 Level 1): default runs + # wq_b+kv_insert; slot [0] runs the full indexer; slot [1] runs the + # MLA compressor. Slot [2] is reserved for the indexer's inner + # overlap. ROCm (aux_streams is None) falls back to sequential. + q, _ = execute_in_parallel( + wq_b_kv_insert, + [ + lambda: indexer( + hidden_states, + qr, + indexer_kv_score, + indexer_weights, + positions, + self.indexer_rotary_emb, + ), + lambda: compressor(kv_score, positions, self.rotary_emb), + ], + self.ln_events[0], + [self.ln_events[1], self.ln_events[2]], + [aux_streams[0], aux_streams[1]] if aux_streams is not None else None, + enable=aux_streams is not None, + ) + elif self.compressor is not None: + # wq_b + kv_insert on default, compressor on aux. + aux_stream = ( + self.aux_stream_list[0] if self.aux_stream_list is not None else None + ) + compressor = self.compressor + + def wq_b_kv_insert() -> torch.Tensor: + q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim) + q = self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata) + return q + + q, _ = maybe_execute_in_parallel( + wq_b_kv_insert, + lambda: compressor(kv_score, positions, self.rotary_emb), + self.ln_events[0], + self.ln_events[1], + aux_stream, + ) + else: + # SWA-only layer: no compressor, no overlap. + q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim) + q = self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata) + + # MLA attention writes into the pre-allocated `out` buffer + # ([num_tokens, padded_heads, head_dim]). + self.forward_mqa(q, kv, positions, out) + + def _fused_qnorm_rope_kv_insert( + self, + q: torch.Tensor, + kv: torch.Tensor, + positions: torch.Tensor, + attn_metadata: ( + dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None + ), + ) -> torch.Tensor: + if not isinstance(attn_metadata, dict): + # Profile run: kernel doesn't fire; produce a padded tensor so + # downstream FlashMLA gets the right shape. + if self.n_local_heads < self.padded_heads: + return F.pad( + q, + (0, 0, 0, self.padded_heads - self.n_local_heads), + value=0.0, + ) + return q + + swa_metadata = cast( + "DeepseekSparseSWAMetadata | None", + attn_metadata.get(self.swa_cache_layer.prefix), + ) + assert swa_metadata is not None + + swa_kv_cache = self.swa_cache_layer.kv_cache + # The fused insert ops require int64 position_ids; the runner's positions + # buffer is already int64, so no cast is needed. + assert positions.dtype == torch.int64 + cos_sin_cache = self.rotary_emb.cos_sin_cache + cache_dtype = swa_kv_cache.dtype + + # kv is unchanged; attention reads kv solely via swa_kv_cache. + if cache_dtype == torch.uint8: + # fp8_ds_mla UE8M0 paged path. Horizontally fused: + # Q side: per-head RMSNorm (no weight) + GPT-J RoPE, zero-filling + # the padding head slots; the kernel allocates and returns + # the padded q tensor. + # KV side: GPT-J RoPE + UE8M0 FP8 quant + paged cache insert. + swa_kv_cache_2d = swa_kv_cache.view(swa_kv_cache.shape[0], -1) + return torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + q, + kv, + swa_kv_cache_2d, + swa_metadata.slot_mapping, + positions, + cos_sin_cache, + self.padded_heads, + self.eps, + swa_metadata.block_size, + ) + + # Plain-row path: the [num_blocks, block_size, 512] cache stores the KV + # row in its element dtype (no Q padding). bf16 rewrites q in place; + # per-tensor fp8 writes a separately-allocated fp8 q and quantizes the + # KV row. + block_size = swa_metadata.block_size + swa_kv_cache_3d = swa_kv_cache.view(-1, block_size, self.head_dim) + if cache_dtype == torch.bfloat16: + torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert( + q, + kv, + swa_kv_cache_3d, + swa_metadata.slot_mapping, + positions, + cos_sin_cache, + self.eps, + block_size, + ) + return q + + # per-tensor fp8 (torch.float8_e4m3fn) + q_fp8 = torch.empty_like(q, dtype=torch.float8_e4m3fn) + torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert( + q, + kv, + q_fp8, + swa_kv_cache_3d, + swa_metadata.slot_mapping, + positions, + cos_sin_cache, + self._flashinfer_fp8_kv_scale, + self._flashinfer_fp8_q_scale_inv, + self.eps, + block_size, + ) + return q_fp8 + + def get_attn_backend(self) -> type[AttentionBackend]: + return self.backend_cls + + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: + if ( + self.compress_ratio <= 1 + ): # SWA part. Allocated separately as DeepseekV4SWACache. + return None + # fp8_ds_mla is a UE8M0 block-scaled uint8 layout and needs 576B + # alignment; plain bf16 / per-tensor fp8 rows use natural element-size + # pages. + uses_fp8_ds_mla_layout = self.kv_cache_dtype == "fp8_ds_mla" + return MLAAttentionSpec( + block_size=vllm_config.cache_config.block_size, + num_kv_heads=1, + head_size=self.head_dim, + dtype=torch.uint8 if uses_fp8_ds_mla_layout else self.kv_cache_torch_dtype, + compress_ratio=self.compress_ratio, + cache_dtype_str=self.kv_cache_dtype, + alignment=576 if uses_fp8_ds_mla_layout else 512, + model_version="deepseek_v4", + kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype), + ) + + +class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase): + def __init__( + self, + head_dim: int, + dtype: torch.dtype, + prefix: str, + cache_config: CacheConfig, + compress_ratio: int = 1, + n_head: int | None = None, + q_head_dim: int | None = None, + ): + super().__init__() + self.kv_cache = torch.tensor([]) + self.head_dim = head_dim + self.prefix = prefix + self.cache_config = cache_config + self.dtype = dtype + self.compress_ratio = compress_ratio + self.n_head = n_head + self.q_head_dim = q_head_dim + compilation_config = get_current_vllm_config().compilation_config + if prefix in compilation_config.static_forward_context: + raise ValueError(f"Duplicate layer name: {prefix}") + compilation_config.static_forward_context[prefix] = self + + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: + # head_dim already carries the fp8 scale padding + # compress_ratio=1 for V3.2, >1 for DeepseekV4; both use the same cache layout. + uses_fp8_ds_mla_layout = vllm_config.cache_config.cache_dtype == "fp8_ds_mla" + return MLAAttentionSpec( + block_size=self.cache_config.block_size, + num_kv_heads=1, + head_size=self.head_dim, + dtype=self.dtype, + compress_ratio=self.compress_ratio, + # 576B for FlashMLA packing; 512B for FlashInfer sparse (#44577). + alignment=576 if uses_fp8_ds_mla_layout else 512, + indexer_n_head=self.n_head, + indexer_q_head_dim=self.q_head_dim, + ) + + def forward(self): ... + + def get_attn_backend(self) -> type[AttentionBackend]: + return DeepseekV4IndexerBackend + + +class DeepseekV4Indexer(nn.Module): + def __init__( + self, + vllm_config: VllmConfig, + config: DeepseekV2Config | DeepseekV3Config, + hidden_size: int, + q_lora_rank: int, + quant_config: QuantizationConfig | None, + cache_config: CacheConfig | None, + topk_indices_buffer: torch.Tensor | None, + compress_ratio: int = 1, + prefix: str = "", + aux_stream: torch.cuda.Stream | None = None, + ): + super().__init__() + self.vllm_config = vllm_config + self.config = config + self.quant_config = quant_config + # self.indexer_cfg = config.attn_module_list_cfg[0]["attn_index"] + self.topk_tokens = config.index_topk + self.n_head = config.index_n_heads # 64 + self.head_dim = config.index_head_dim # 128 + self.rope_dim = config.qk_rope_head_dim # 64 + self.q_lora_rank = q_lora_rank # 1536 + self.compress_ratio = compress_ratio + self.use_fp4_kv = self.vllm_config.attention_config.use_fp4_indexer_cache + logger.info_once( + "Using %s indexer cache for Lightning Indexer.", + "MXFP4" if self.use_fp4_kv else "FP8", + ) + + # no tensor parallel, just replicated + self.wq_b = ReplicatedLinear( + self.q_lora_rank, + self.head_dim * self.n_head, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.wq_b", + ) + self.weights_proj = ReplicatedLinear( + hidden_size, + self.n_head, + bias=False, + quant_config=None, + prefix=f"{prefix}.weights_proj", + ) + self.softmax_scale = self.head_dim**-0.5 + + self.scale_fmt = "ue8m0" + self.quant_block_size = 128 # TODO: get from config + self.topk_indices_buffer = topk_indices_buffer + + self.max_model_len = ( + vllm_config.model_config.max_model_len // self.compress_ratio + ) + self.prefix = prefix + + self.max_total_seq_len = ( + get_max_prefill_buffer_size(vllm_config) // self.compress_ratio + ) + + assert cache_config is not None, "Deepseek V4 indexer requires cache_config" + # NOTE(yifan): FP8 indxer cache use the same layout as V3.2: + # head_dim bytes = 128 fp8 + 4 fp32 scale = 132. + # For FP4 indexer cache, we still allocate the same amount of memory as FP8, + # but only use the first half of the memory. + k_cache_head_dim = self.head_dim + self.head_dim // self.quant_block_size * 4 + self.k_cache = DeepseekV4IndexerCache( + head_dim=k_cache_head_dim, + dtype=torch.uint8, + prefix=f"{prefix}.k_cache", + cache_config=cache_config, + compress_ratio=self.compress_ratio, + n_head=self.n_head, + q_head_dim=self.head_dim, + ) + self.compressor = DeepseekCompressor( + vllm_config=vllm_config, + compress_ratio=self.compress_ratio, + hidden_size=hidden_size, + head_dim=self.head_dim, + rotate=True, + prefix=f"{prefix}.compressor", + k_cache_prefix=self.k_cache.prefix, + use_fp4_cache=self.use_fp4_kv, + ) + + self.indexer_op = SparseAttnIndexer( + self.k_cache, + self.quant_block_size, + self.scale_fmt, + self.topk_tokens, + self.head_dim, + self.max_model_len, + self.max_total_seq_len, + self.topk_indices_buffer, + skip_k_cache_insert=True, + use_fp4_cache=self.use_fp4_kv, + ) + + # None on ROCm — maybe_execute_in_parallel falls back to sequential. + self.aux_stream = aux_stream + self.ln_events: list[torch.cuda.Event] = [ + torch.cuda.Event(), + torch.cuda.Event(), + ] + + def forward( + self, + hidden_states: torch.Tensor, + qr: torch.Tensor, + compressed_kv_score: torch.Tensor, + indexer_weights: torch.Tensor, + positions: torch.Tensor, + rotary_emb: nn.Module, + ) -> torch.Tensor: + compressor = self.compressor + + def wq_b_and_q_quant(): + # ReplicatedLinear returns (output, bias); bias is None. + q, _ = self.wq_b(qr) + q = q.view(-1, self.n_head, self.head_dim) + return fused_indexer_q_rope_quant( + positions, + q, + rotary_emb.cos_sin_cache, + indexer_weights, + self.softmax_scale, + self.n_head**-0.5, + use_fp4=self.use_fp4_kv, + ) + + # compressor returns None and writes K to the indexer KV cache; the + # join orders that write before indexer_op (skip_k_cache_insert=True). + (q_quant, weights), k = maybe_execute_in_parallel( + wq_b_and_q_quant, + lambda: compressor(compressed_kv_score, positions, rotary_emb), + self.ln_events[0], + self.ln_events[1], + self.aux_stream, + ) + return self.indexer_op(hidden_states, q_quant, k, weights) diff --git a/src/vllm/models/deepseek_v4/ppu/flashmla.py b/src/vllm/models/deepseek_v4/ppu/flashmla.py new file mode 100644 index 0000000..20cbfe8 --- /dev/null +++ b/src/vllm/models/deepseek_v4/ppu/flashmla.py @@ -0,0 +1,403 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import TYPE_CHECKING, cast + +import torch + +from vllm.forward_context import get_forward_context +from vllm.models.deepseek_v4.attention import DeepseekV4Attention +from vllm.platforms import current_platform +from vllm.models.deepseek_v4.common.ops import ( + combine_topk_swa_indices, + compute_global_topk_indices_and_lens, + dequantize_and_gather_k_cache, +) +from vllm.models.deepseek_v4.ppu.ops.o_proj import ( + compute_fp8_einsum_recipe, + deep_gemm_fp8_o_proj, + deep_gemm_fp8_channel_o_proj, + deep_gemm_int8_o_proj, +) +from vllm.models.deepseek_v4.sparse_mla import ( + DeepseekV4FlashMLABackend, + DeepseekV4FlashMLAMetadata, +) +from vllm.v1.attention.ops.flashmla import ( + flash_mla_sparse_fwd, + flash_mla_with_kvcache, +) +from vllm.v1.worker.workspace import current_workspace_manager + +if TYPE_CHECKING: + from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata + + +class DeepseekV4FlashMLAAttention(DeepseekV4Attention): + """FlashMLA sparse MLA attention layer for DeepSeek V4 (CUDA).""" + + backend_cls = DeepseekV4FlashMLABackend + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + self._einsum_recipe, self._tma_aligned_scales = compute_fp8_einsum_recipe() + + def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + import os + if os.environ.get("VLLM_PPU_BF16_O_PROJ") == "1": + from vllm.models.deepseek_v4.ppu.ops.o_proj import bf16_o_proj + return bf16_o_proj( + o, positions, self.rotary_emb.cos_sin_cache, self.wo_a, self.wo_b, + n_groups=self.n_local_groups, + heads_per_group=self.n_local_heads // self.n_local_groups, + nope_dim=self.nope_head_dim, rope_dim=self.rope_head_dim, + o_lora_rank=self.o_lora_rank, + ) + if self.wo_a.weight.dtype == torch.int8: + return deep_gemm_int8_o_proj( + o, + positions, + self.rotary_emb.cos_sin_cache, + self.wo_a, + self.wo_b, + n_groups=self.n_local_groups, + heads_per_group=self.n_local_heads // self.n_local_groups, + nope_dim=self.nope_head_dim, + rope_dim=self.rope_head_dim, + o_lora_rank=self.o_lora_rank, + einsum_recipe=self._einsum_recipe, + ) + elif self.wo_a.weight.dtype == torch.float8_e4m3fn and hasattr( + self.wo_a, "weight_scale" + ): + # Channelwise FP8 (PPU): wo_a kept as FP8 + per-output-channel + # weight_scale (no pre-dequant). + return deep_gemm_fp8_channel_o_proj( + o, + positions, + self.rotary_emb.cos_sin_cache, + self.wo_a, + self.wo_b, + n_groups=self.n_local_groups, + heads_per_group=self.n_local_heads // self.n_local_groups, + nope_dim=self.nope_head_dim, + rope_dim=self.rope_head_dim, + o_lora_rank=self.o_lora_rank, + einsum_recipe=self._einsum_recipe, + ) + return deep_gemm_fp8_o_proj( + o, + positions, + self.rotary_emb.cos_sin_cache, + self.wo_a, + self.wo_b, + n_groups=self.n_local_groups, + heads_per_group=self.n_local_heads // self.n_local_groups, + nope_dim=self.nope_head_dim, + rope_dim=self.rope_head_dim, + o_lora_rank=self.o_lora_rank, + einsum_recipe=self._einsum_recipe, + tma_aligned_scales=self._tma_aligned_scales, + ) + + @classmethod + def get_padded_num_q_heads(cls, num_heads: int) -> int: + # FP8 decode kernel only supports h_q = 64 or 128. + if num_heads > 128: + raise ValueError( + f"DeepseekV4 FlashMLA does not support {num_heads} heads " + "(FP8 decode kernel requires h_q in {64, 128})." + ) + if current_platform.is_ppu(): + # PPU FlashMLA supports arbitrary head counts; no padding needed. + return num_heads + return 64 if num_heads <= 64 else 128 + + def forward_mqa( + self, + q: torch.Tensor, + kv: torch.Tensor, + positions: torch.Tensor, + output: torch.Tensor, + ) -> None: + assert output.shape == q.shape, ( + f"output buffer shape {output.shape} must match q shape {q.shape}" + ) + assert output.dtype == q.dtype, ( + f"output buffer dtype {output.dtype} must match q dtype {q.dtype}" + ) + + # Get SWA and indexer metadata from forward context + forward_context = get_forward_context() + attn_metadata = forward_context.attn_metadata + + if attn_metadata is None: + # Warmup dummy run: no real metadata. Reserve the same bf16 + # gather workspace _forward_prefill would; the dequantize / topk + # / sparse_fwd kernels are skipped this step. + swa_only = self.compress_ratio <= 1 + N = ( + 0 + if swa_only + else (self.max_model_len + self.compress_ratio - 1) + // self.compress_ratio + ) + M = N + self.window_size + self.max_num_batched_tokens + current_workspace_manager().get_simultaneous( + ((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16), + ) + output.zero_() + return + + assert isinstance(attn_metadata, dict) + flashmla_metadata = cast( + DeepseekV4FlashMLAMetadata | None, attn_metadata.get(self.prefix) + ) + swa_metadata = cast( + "DeepseekSparseSWAMetadata | None", + attn_metadata.get(self.swa_cache_layer.prefix), + ) + assert swa_metadata is not None + + swa_only = self.compress_ratio <= 1 + # SWA-only layers (compress_ratio <= 1) don't have their own KV cache + # allocation, so self.kv_cache may be empty after profiling cleanup. + self_kv_cache = self.kv_cache if not swa_only else None + swa_kv_cache = self.swa_cache_layer.kv_cache + + # Split prefill and decode + num_decodes = swa_metadata.num_decodes + num_prefills = swa_metadata.num_prefills + num_decode_tokens = swa_metadata.num_decode_tokens + + if num_prefills > 0: + self._forward_prefill( + q=q[num_decode_tokens:], + positions=positions[num_decode_tokens:], + compressed_k_cache=self_kv_cache, + swa_k_cache=swa_kv_cache, + output=output[num_decode_tokens:], + attn_metadata=flashmla_metadata, + swa_metadata=swa_metadata, + ) + if num_decodes > 0: + self._forward_decode( + q=q[:num_decode_tokens], + kv_cache=self_kv_cache, + swa_metadata=swa_metadata, + attn_metadata=flashmla_metadata, + swa_only=swa_only, + output=output[:num_decode_tokens], + ) + + def _forward_decode( + self, + q: torch.Tensor, + kv_cache: torch.Tensor | None, # Only used when compress_ratio > 1 + swa_metadata: "DeepseekSparseSWAMetadata", + attn_metadata: DeepseekV4FlashMLAMetadata | None, + swa_only: bool, + output: torch.Tensor, + ) -> None: + num_decodes = swa_metadata.num_decodes + num_decode_tokens = swa_metadata.num_decode_tokens + + topk_indices = None + topk_lens = None + if not swa_only: + assert attn_metadata is not None + assert swa_metadata.is_valid_token is not None + block_size = attn_metadata.block_size // self.compress_ratio + is_valid = swa_metadata.is_valid_token[:num_decode_tokens] + if self.compress_ratio == 4: + # C4A: local indices differ per layer (filled by Indexer). + assert self.topk_indices_buffer is not None + global_indices, topk_lens = compute_global_topk_indices_and_lens( + self.topk_indices_buffer[:num_decode_tokens], + swa_metadata.token_to_req_indices, + attn_metadata.block_table[:num_decodes], + block_size, + is_valid, + ) + topk_indices = global_indices.view(num_decode_tokens, 1, -1) + else: + # C128A: pre-computed during metadata build. + topk_indices = attn_metadata.c128a_global_decode_topk_indices + topk_lens = attn_metadata.c128a_decode_topk_lens + + swa_indices = swa_metadata.decode_swa_indices + swa_lens = swa_metadata.decode_swa_lens + + # We treat queries in the same seq as different queries + # and later we only attend by generated indices. + # q arrives pre-padded to self.padded_heads by the outer wrapper. + q = q.unsqueeze(1) + + # Prepare SWA cache (num_blocks, swa_block_size, 1, head_bytes) + # Use unsqueeze to preserve strides (handles padded blocks correctly) + swa_cache = self.swa_cache_layer.kv_cache.unsqueeze(-2) + # Reshape KV cache to (num_blocks, block_size, 1, head_bytes) + if kv_cache is not None: + kv_cache = kv_cache.unsqueeze(-2) + + # One FlashMLASchedMeta per layer type, shared across all same-type + # layers within this decode step. The first forward call per type + # triggers the in-kernel planner (allocating tile_scheduler_metadata + # and num_splits via PyTorch's graph-aware allocator so CUDA graph + # capture reuses the same addresses on replay); subsequent same-type + # layers see have_initialized=True and skip the planner. + if self.compress_ratio <= 1: + tile_metadata = swa_metadata.tile_sched_swaonly + elif self.compress_ratio == 4: + tile_metadata = swa_metadata.tile_sched_c4a + elif self.compress_ratio == 128: + tile_metadata = swa_metadata.tile_sched_c128a + else: + raise ValueError( + f"Unsupported compress_ratio={self.compress_ratio}; " + "expected 1, 4, or 128." + ) + assert tile_metadata is not None, ( + "swa_metadata missing tile_sched entry for " + f"compress_ratio={self.compress_ratio}; " + "DeepseekSparseSWAMetadataBuilder.build_tile_scheduler did not " + "allocate one for this layer type." + ) + + out, _ = flash_mla_with_kvcache( + q=q, + k_cache=swa_cache, + block_table=None, + head_dim_v=512, + tile_scheduler_metadata=tile_metadata, + cache_seqlens=None, + is_fp8_kvcache=True, + indices=swa_indices, + topk_length=swa_lens, + softmax_scale=self.scale, + attn_sink=self.attn_sink, + extra_k_cache=kv_cache if not swa_only else None, + extra_indices_in_kvcache=topk_indices, + extra_topk_length=topk_lens, + out=output.unsqueeze(1), + ) + + def _forward_prefill( + self, + q: torch.Tensor, + positions: torch.Tensor, + compressed_k_cache: torch.Tensor | None, # Only used when compress_ratio > 1 + swa_k_cache: torch.Tensor, + output: torch.Tensor, + attn_metadata: DeepseekV4FlashMLAMetadata | None, + swa_metadata: "DeepseekSparseSWAMetadata", + ) -> None: + swa_only = attn_metadata is None + + num_prefills = swa_metadata.num_prefills + num_prefill_tokens = swa_metadata.num_prefill_tokens + num_decodes = swa_metadata.num_decodes + num_decode_tokens = swa_metadata.num_decode_tokens + + # Use pre-computed prefill metadata. + seq_lens = swa_metadata.prefill_seq_lens + gather_lens = swa_metadata.prefill_gather_lens + assert seq_lens is not None + assert gather_lens is not None + + # Derive prefill-local token offsets from the full query_start_loc_cpu. + query_start_loc_cpu = swa_metadata.query_start_loc_cpu + query_start_loc = swa_metadata.query_start_loc + assert query_start_loc_cpu is not None + assert query_start_loc is not None + prefill_token_base = query_start_loc_cpu[num_decodes] + + if not swa_only: + if self.compress_ratio == 4: + assert self.topk_indices_buffer is not None + topk_indices = self.topk_indices_buffer[num_decode_tokens:] + topk_indices = topk_indices[:num_prefill_tokens] + else: + # C128A: pre-computed during metadata build. + assert attn_metadata is not None + topk_indices = attn_metadata.c128a_prefill_topk_indices + top_k = topk_indices.shape[-1] + # Compressed region must fit the full compressed pool (seq_len // + # compress_ratio), not just top_k. top_k bounds how many indices + # the indexer selects, not the pool size it indexes into. + N = (self.max_model_len + self.compress_ratio - 1) // self.compress_ratio + else: + # NOTE(woosuk): topk_indices will not be used for SWA-only layers. + assert self.topk_indices_buffer is not None + topk_indices = self.topk_indices_buffer[num_decode_tokens:] + top_k = 0 + N = 0 + + M = N + self.window_size + self.max_num_batched_tokens + chunk_size_const = self.PREFILL_CHUNK_SIZE + num_chunks = (num_prefills + chunk_size_const - 1) // chunk_size_const + + workspace_manager = current_workspace_manager() + kv = workspace_manager.get_simultaneous( + ((chunk_size_const, M, q.shape[-1]), torch.bfloat16), + )[0] + for chunk_idx in range(num_chunks): + chunk_start = chunk_idx * chunk_size_const + chunk_end = min(chunk_start + chunk_size_const, num_prefills) + chunk_size = chunk_end - chunk_start + if not swa_only: + # Gather compressed KV + assert attn_metadata is not None + block_table = attn_metadata.block_table[num_decodes:] + dequantize_and_gather_k_cache( + kv[:chunk_size], + compressed_k_cache, + seq_lens=seq_lens[chunk_start:chunk_end] // self.compress_ratio, + gather_lens=None, + block_table=block_table[chunk_start:chunk_end], + block_size=attn_metadata.block_size // self.compress_ratio, + offset=0, + ) + + # Gather SWA KV + swa_block_table = swa_metadata.block_table[num_decodes:] + dequantize_and_gather_k_cache( + kv[:chunk_size], + swa_k_cache, + seq_lens=seq_lens[chunk_start:chunk_end], + gather_lens=gather_lens[chunk_start:chunk_end], + block_table=swa_block_table[chunk_start:chunk_end], + block_size=swa_metadata.block_size, + offset=N, + ) + + # Combine the topk indices and SWA indices for gathered KV cache + query_start = ( + query_start_loc_cpu[num_decodes + chunk_start] - prefill_token_base + ) + query_end = ( + query_start_loc_cpu[num_decodes + chunk_end] - prefill_token_base + ) + + combined_indices, combined_lens = combine_topk_swa_indices( + topk_indices[query_start:query_end], + query_start_loc[ + num_decodes + chunk_start : num_decodes + chunk_end + 1 + ], + seq_lens[chunk_start:chunk_end], + gather_lens[chunk_start:chunk_end], + self.window_size, + self.compress_ratio, + top_k, + M, + N, + ) + flash_mla_sparse_fwd( + q=q[query_start:query_end], + kv=kv.view(-1, 1, q.shape[-1]), + indices=combined_indices.unsqueeze(1), + sm_scale=self.scale, + attn_sink=self.attn_sink, + topk_length=combined_lens, + out=output[query_start:query_end], + ) diff --git a/src/vllm/models/deepseek_v4/ppu/ops/o_proj.py b/src/vllm/models/deepseek_v4/ppu/ops/o_proj.py new file mode 100644 index 0000000..6941f10 --- /dev/null +++ b/src/vllm/models/deepseek_v4/ppu/ops/o_proj.py @@ -0,0 +1,300 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import torch +import torch.nn as nn + +from vllm.models.deepseek_v4.common.ops import ( + fused_inv_rope_fp8_quant, + fused_inv_rope_fp8_quant_channelwise, + fused_inv_rope_int8_quant_channelwise, +) +from vllm.platforms import current_platform +from vllm.utils.ppu_deep_gemm import fp8_einsum +from vllm.utils.ppu_deep_gemm import int8_einsum +from vllm.utils.torch_utils import direct_register_custom_op + + +def compute_fp8_einsum_recipe() -> tuple[tuple[int, int, int], bool]: + """fp8_einsum recipe + scale layout for the current GPU arch. + + SM90: FP32 block scales stay [g, r/128, d/128] → sfb_gran_mn=128. + SM100: INT32 packed scales become [g, r, ...] → sfb_gran_mn=1. + + Returns ``(einsum_recipe, tma_aligned_scales)`` for ``deep_gemm_fp8_o_proj``. + """ + cap = current_platform.get_device_capability() + assert cap is not None, "DeepseekV4 attention requires a CUDA device" + einsum_recipe = (1, 128, 128) if cap.major <= 9 else (1, 1, 128) + tma_aligned_scales = cap.major >= 10 + return einsum_recipe, tma_aligned_scales + + +def deep_gemm_fp8_o_proj( + o: torch.Tensor, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, + wo_a: nn.Module, + wo_b: nn.Module, + *, + n_groups: int, + heads_per_group: int, + nope_dim: int, + rope_dim: int, + o_lora_rank: int, + einsum_recipe: tuple[int, int, int], + tma_aligned_scales: bool, +) -> torch.Tensor: + """O projection: inverse RoPE + FP8 quant + einsum + wo_b. + + Shared by the FlashMLA and FlashInfer CUDA backends. ``einsum_recipe`` / + ``tma_aligned_scales`` come from ``compute_fp8_einsum_recipe``. + """ + o_fp8, o_scale = fused_inv_rope_fp8_quant( + o, + positions, + cos_sin_cache, + n_groups=n_groups, + heads_per_group=heads_per_group, + nope_dim=nope_dim, + rope_dim=rope_dim, + tma_aligned_scales=tma_aligned_scales, + ) + z = torch.empty( + (o.shape[0], n_groups, o_lora_rank), + device=o.device, + dtype=torch.bfloat16, + ) + fp8_einsum( + "bhr,hdr->bhd", + (o_fp8, o_scale), + (wo_a.weight, wo_a.weight_scale_inv), + z, + recipe=einsum_recipe, + ) + return wo_b(z.flatten(1)) + + +def deep_gemm_int8_o_proj( + o: torch.Tensor, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, + wo_a: nn.Module, + wo_b: nn.Module, + *, + n_groups: int, + heads_per_group: int, + nope_dim: int, + rope_dim: int, + o_lora_rank: int, + einsum_recipe: tuple[int, int, int], +): + # Channelwise INT8 GEMM via PPU DeepGEMM int8_einsum: + # - activation: per-(token, group) symmetric INT8 quant. + # o_int8 has shape (b=T, h=G, r=R) so o_scale is (b, h, 1). + # - weight: per-output-channel INT8. compressed-tensors + # W8A8 channel scheme stores wo_a as 2D [G*d, R] / scale as + # [G*d, 1]; reshape to 3D so int8_bmm sees (h=G, d=lora_rank, + # r=R) for `hdr` and a (h, d, 1) scale broadcasting along r. + o_int8, o_scale = fused_inv_rope_int8_quant_channelwise( + o, + positions, + cos_sin_cache, + n_groups, + heads_per_group, + nope_dim, + rope_dim, + ) + wo_a_view = wo_a.weight.reshape( + n_groups, o_lora_rank, -1 + ) + wo_a_scale_view = wo_a.weight_scale.reshape( + n_groups, o_lora_rank, 1 + ) + z = torch.empty( + (o.shape[0], n_groups, o_lora_rank), + device=o.device, + dtype=torch.bfloat16, + ) + torch.ops.vllm.deepseek_v4_int8_einsum( + o_int8, + o_scale, + wo_a_view, + wo_a_scale_view, + z, + "bhr,hdr->bhd", + [1, 1, o_int8.shape[-1]], + ) + return wo_b(z.flatten(1)) + + +def deep_gemm_fp8_channel_o_proj( + o: torch.Tensor, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, + wo_a: nn.Module, + wo_b: nn.Module, + *, + n_groups: int, + heads_per_group: int, + nope_dim: int, + rope_dim: int, + o_lora_rank: int, + einsum_recipe: tuple[int, int, int], +): + # Channelwise FP8 GEMM via PPU DeepGEMM fp8_einsum: + # - activation: per-(token, group) symmetric FP8 quant. + # o_fp8 has shape (b=T, h=G, r=R) so o_scale is (b, h, 1). + # - weight: per-output-channel FP8. Channelwise scheme stores + # wo_a as 2D [G*d, R] / scale as [G*d, 1]; reshape to 3D so + # fp8_bmm sees (h=G, d=lora_rank, r=R) for `hdr` and a + # (h, d, 1) scale broadcasting along r. + o_fp8, o_scale = fused_inv_rope_fp8_quant_channelwise( + o, + positions, + cos_sin_cache, + n_groups, + heads_per_group, + nope_dim, + rope_dim, + ) + wo_a_view = wo_a.weight.reshape( + n_groups, o_lora_rank, -1 + ) + wo_a_scale_view = wo_a.weight_scale.reshape( + n_groups, o_lora_rank, 1 + ) + z = torch.empty( + (o.shape[0], n_groups, o_lora_rank), + device=o.device, + dtype=torch.bfloat16, + ) + torch.ops.vllm.deepseek_v4_fp8_channel_einsum( + o_fp8, + o_scale, + wo_a_view, + wo_a_scale_view, + z, + "bhr,hdr->bhd", + [1, 1, o_fp8.shape[-1]], + ) + return wo_b(z.flatten(1)) + + +def deepseek_v4_int8_einsum( + a: torch.Tensor, + a_scale: torch.Tensor, + b: torch.Tensor, + b_scale: torch.Tensor, + out: torch.Tensor, + equation: str, + recipe: list[int], +) -> None: + int8_einsum(equation, (a, a_scale), (b, b_scale), out, recipe=tuple(recipe)) + + +def deepseek_v4_int8_einsum_fake( + a: torch.Tensor, + a_scale: torch.Tensor, + b: torch.Tensor, + b_scale: torch.Tensor, + out: torch.Tensor, + equation: str, + recipe: list[int], +) -> None: + return None + + +direct_register_custom_op( + op_name="deepseek_v4_int8_einsum", + op_func=deepseek_v4_int8_einsum, + mutates_args=["out"], + fake_impl=deepseek_v4_int8_einsum_fake, +) + + +def deepseek_v4_fp8_channel_einsum( + a: torch.Tensor, + a_scale: torch.Tensor, + b: torch.Tensor, + b_scale: torch.Tensor, + out: torch.Tensor, + equation: str, + recipe: list[int], +) -> None: + fp8_einsum(equation, (a, a_scale), (b, b_scale), out, recipe=tuple(recipe)) + + +def deepseek_v4_fp8_channel_einsum_fake( + a: torch.Tensor, + a_scale: torch.Tensor, + b: torch.Tensor, + b_scale: torch.Tensor, + out: torch.Tensor, + equation: str, + recipe: list[int], +) -> None: + return None + + +direct_register_custom_op( + op_name="deepseek_v4_fp8_channel_einsum", + op_func=deepseek_v4_fp8_channel_einsum, + mutates_args=["out"], + fake_impl=deepseek_v4_fp8_channel_einsum_fake, +) +# Experimental ZW810E BF16 O-projection fallback (opt-in). +def prepare_bf16_o_proj_weight(layer): + """Decode checkpoint block FP8 before any Marlin/DeepGEMM repacking.""" + from vllm.model_executor.utils import replace_parameter + w = layer.weight + s = layer.weight_scale_inv + if w.dtype != torch.float8_e4m3fn or w.ndim != 2: + raise ValueError(f"BF16 O-proj expects raw 2D FP8 weight, got {w.dtype} {w.shape}") + bn, bk = layer.weight_block_size + n, k = w.shape + expected = ((n + bn - 1) // bn, (k + bk - 1) // bk) + if tuple(s.shape) != expected: + raise ValueError(f"BF16 O-proj scale shape {s.shape}, expected {expected}") + # Convert row blocks to keep temporary FP32 memory bounded. + decoded = torch.empty((n, k), device=w.device, dtype=torch.bfloat16) + sf = s.float() + for row in range(0, n, bn): + expanded = sf[row // bn].repeat_interleave(bk)[:k] + decoded[row:row + bn] = ( + w[row:row + bn].float() * expanded[None, :] + ).to(torch.bfloat16) + replace_parameter(layer, "weight", decoded) + layer._ppu_bf16_o_proj_ready = True + + +def bf16_o_proj( + o, positions, cos_sin_cache, wo_a, wo_b, *, + n_groups, heads_per_group, nope_dim, rope_dim, o_lora_rank, **kwargs +): + """Inverse GPT-J RoPE followed by grouped BF16 GEMM and existing wo_b.""" + if not getattr(wo_a, "_ppu_bf16_o_proj_ready", False): + raise RuntimeError("BF16 O-proj weight was not prepared before repacking") + t, h, d = o.shape + if h != n_groups * heads_per_group or d != nope_dim + rope_dim: + raise ValueError(f"Unexpected O-proj input shape: {o.shape}") + # Match fused kernel: even = x*cos + partner*sin; odd = x*cos - partner*sin. + cs = cos_sin_cache.index_select(0, positions.reshape(-1).long()).float() + half = rope_dim // 2 + cos = cs[:, None, :half] + sin = cs[:, None, half:] + tail = o[..., nope_dim:].float().reshape(t, h, half, 2) + even, odd = tail[..., 0], tail[..., 1] + rotated = torch.stack((even * cos + odd * sin, + odd * cos - even * sin), dim=-1) + x = torch.cat((o[..., :nope_dim], + rotated.flatten(-2).to(o.dtype)), dim=-1) + r = heads_per_group * d + w = wo_a.weight.reshape(n_groups, o_lora_rank, r) + x = x.reshape(t, n_groups, r).transpose(0, 1).contiguous() + z = torch.bmm(x.to(torch.bfloat16), w.transpose(1, 2)) + return wo_b(z.transpose(0, 1).contiguous().flatten(1)) + +# PPU cannot use the BF16-input / FP32-output compressor GEMM here. +# Convert before GEMM so an invalid device operation cannot poison the stream. +def ppu_compressor_mm(a, b): + return torch.mm(a.float().contiguous(), b.float().contiguous()) diff --git a/tests/test_apply.py b/tests/test_apply.py new file mode 100644 index 0000000..5f34a86 --- /dev/null +++ b/tests/test_apply.py @@ -0,0 +1,60 @@ +import hashlib +import json +import subprocess +import sys +import tempfile +import unittest +from pathlib import Path + + +ROOT = Path(__file__).resolve().parents[1] +MANIFEST = json.loads((ROOT / "manifest.json").read_text())["files"] +SCRIPT = ROOT / "scripts/apply.py" + + +class ApplyPatchTests(unittest.TestCase): + def setUp(self): + self.temp = tempfile.TemporaryDirectory() + self.site = Path(self.temp.name) + for name in MANIFEST: + target = self.site / name + target.parent.mkdir(parents=True, exist_ok=True) + target.write_bytes((ROOT / "upstream" / name).read_bytes()) + + def tearDown(self): + self.temp.cleanup() + + def command(self, action): + return subprocess.run( + [sys.executable, str(SCRIPT), action, "--site-packages", str(self.site)], + capture_output=True, + text=True, + ) + + def test_apply_is_idempotent_and_rollback_restores_originals(self): + self.assertEqual(self.command("--check").returncode, 0) + for _ in range(2): + result = self.command("--apply") + self.assertEqual(result.returncode, 0, result.stderr) + for name, hashes in MANIFEST.items(): + target = self.site / name + self.assertEqual(hashlib.sha256(target.read_bytes()).hexdigest(), hashes["patched_sha256"]) + backup = target.with_name(target.name + ".ppu-dsv4.orig") + self.assertEqual(hashlib.sha256(backup.read_bytes()).hexdigest(), hashes["original_sha256"]) + result = self.command("--rollback") + self.assertEqual(result.returncode, 0, result.stderr) + for name, hashes in MANIFEST.items(): + self.assertEqual(hashlib.sha256((self.site / name).read_bytes()).hexdigest(), hashes["original_sha256"]) + + def test_unknown_base_refuses_before_any_file_is_changed(self): + changed = self.site / next(iter(MANIFEST)) + changed.write_bytes(b"unexpected version\n") + result = self.command("--apply") + self.assertNotEqual(result.returncode, 0) + self.assertIn("does not match", result.stderr) + for name, hashes in list(MANIFEST.items())[1:]: + self.assertEqual(hashlib.sha256((self.site / name).read_bytes()).hexdigest(), hashes["original_sha256"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/upstream/vllm/model_executor/layers/quantization/fp8.py b/upstream/vllm/model_executor/layers/quantization/fp8.py new file mode 100644 index 0000000..56f43d5 --- /dev/null +++ b/upstream/vllm/model_executor/layers/quantization/fp8.py @@ -0,0 +1,871 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import TYPE_CHECKING, Any + +import torch +from torch.utils._python_dispatch import TorchDispatchMode + +import vllm.envs as envs +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.config import get_current_vllm_config +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.logger import init_logger +from vllm.model_executor.kernels.linear import ( + init_fp8_linear_kernel, +) +from vllm.model_executor.kernels.linear.scaled_mm import ( + CutlassFP8ScaledMMLinearKernel, + MarlinFP8ScaledMMLinearKernel, +) +from vllm.model_executor.layers.attention import Attention +from vllm.model_executor.layers.fused_moe import ( + FusedMoEMethodBase, + FusedMoeWeightScaleSupported, + RoutedExperts, + SharedExperts, + UnquantizedFusedMoEMethod, +) +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEQuantConfig, +) +from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( + Fp8MoeBackend, + convert_to_fp8_moe_kernel_format, + make_fp8_moe_kernel, + make_fp8_moe_quant_config, + select_fp8_moe_backend, +) +from vllm.model_executor.layers.linear import ( + LinearBase, + LinearMethodBase, + UnquantizedLinearMethod, +) +from vllm.model_executor.layers.quantization import QuantizationMethods +from vllm.model_executor.layers.quantization.base_config import ( + QuantizationConfig, + QuantizeMethodBase, +) +from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod +from vllm.model_executor.layers.quantization.utils.fp8_utils import ( + create_fp8_input_scale, + create_fp8_scale_parameter, + create_fp8_weight_parameter, + process_fp8_input_tensor_strategy_moe, + process_fp8_weight_tensor_strategy, + process_fp8_weight_tensor_strategy_moe, + validate_fp8_block_shape, +) +from vllm.model_executor.layers.quantization.utils.marlin_utils import ( + get_marlin_input_dtype, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + GroupShape, + create_fp8_quant_key, + is_layer_skipped, + kFp8Dynamic128Sym, + kFp8DynamicTensorSym, + kFp8DynamicTokenSym, + kFp8Static128BlockSym, + kFp8StaticTensorSym, +) +from vllm.model_executor.layers.quantization.utils.w8a8_utils import ( + cutlass_block_fp8_supported, + cutlass_fp8_supported, + normalize_e4m3fn_to_e4m3fnuz, +) +from vllm.model_executor.parameter import ( + BlockQuantScaleParameter, + PerTensorScaleParameter, +) +from vllm.model_executor.utils import replace_parameter, set_weight_attrs +from vllm.platforms import current_platform + +if current_platform.is_ppu(): + from vllm.utils.ppu_deep_gemm import ( + is_deep_gemm_supported, + ) +else: + from vllm.utils.deep_gemm import ( + is_deep_gemm_supported, + ) + +if TYPE_CHECKING: + from vllm.model_executor.models.utils import WeightsMapper + +ACTIVATION_SCHEMES = ["static", "dynamic"] + +logger = init_logger(__name__) + + +class Fp8Config(QuantizationConfig): + """Config class for FP8.""" + + def __init__( + self, + is_checkpoint_fp8_serialized: bool = False, + activation_scheme: str = "dynamic", + ignored_layers: list[str] | None = None, + weight_block_size: list[int] | None = None, + store_dtype: str | None = None, + ) -> None: + super().__init__() + + self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized + + if activation_scheme not in ACTIVATION_SCHEMES: + raise ValueError(f"Unsupported activation scheme {activation_scheme}") + self.activation_scheme = activation_scheme + self.ignored_layers = ignored_layers or [] + self.store_dtype = store_dtype + if weight_block_size is not None: + if not is_checkpoint_fp8_serialized: + raise ValueError( + "The block-wise quantization only supports fp8-serialized " + "checkpoint for now." + ) + if len(weight_block_size) != 2: + raise ValueError( + "The quantization block size of weight must have 2 " + f"dimensions, but got {len(weight_block_size)} dimensions" + ) + if activation_scheme != "dynamic": + raise ValueError( + "The block-wise quantization only supports " + "dynamic activation scheme for now, but got " + f"{activation_scheme} activation scheme." + ) + self.weight_block_size = weight_block_size + self.use_deep_gemm: bool | None = None + + @classmethod + def get_name(cls) -> QuantizationMethods: + return "fp8" + + @classmethod + def get_supported_act_dtypes(cls) -> list[torch.dtype]: + return [torch.bfloat16, torch.half] + + @classmethod + def get_min_capability(cls) -> int: + return 75 + + @classmethod + def get_config_filenames(cls) -> list[str]: + return [] + + def apply_vllm_mapper(self, hf_to_vllm_mapper: "WeightsMapper"): + if self.ignored_layers is not None: + self.ignored_layers = hf_to_vllm_mapper.apply_list(self.ignored_layers) + + @classmethod + def from_config(cls, config: dict[str, Any]) -> "Fp8Config": + quant_method = cls.get_from_keys(config, ["quant_method"]) + is_checkpoint_fp8_serialized = "fp8" in quant_method + activation_scheme = cls.get_from_keys(config, ["activation_scheme"]) + ignored_layers = cls.get_from_keys_or(config, ["ignored_layers"], None) + weight_block_size = cls.get_from_keys_or(config, ["weight_block_size"], None) + store_dtype = cls.get_from_keys_or(config, ["store_dtype"], None) + if not ignored_layers: + ignored_layers = cls.get_from_keys_or( + config, ["modules_to_not_convert"], None + ) + return cls( + is_checkpoint_fp8_serialized=is_checkpoint_fp8_serialized, + activation_scheme=activation_scheme, + ignored_layers=ignored_layers, + weight_block_size=weight_block_size, + store_dtype=store_dtype, + ) + + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> "QuantizeMethodBase | None": + if isinstance(layer, LinearBase): + if is_layer_skipped( + prefix=prefix, + ignored_layers=self.ignored_layers, + fused_mapping=self.packed_modules_mapping, + ): + return UnquantizedLinearMethod() + if not self.is_checkpoint_fp8_serialized: + from vllm.model_executor.layers.quantization.online.fp8 import ( + Fp8PerTensorOnlineLinearMethod, + ) + + online_method = Fp8PerTensorOnlineLinearMethod() + online_method.marlin_input_dtype = get_marlin_input_dtype(prefix) + return online_method + else: + offline_method = Fp8LinearMethod(self) + offline_method.marlin_input_dtype = get_marlin_input_dtype(prefix) + return offline_method + elif isinstance(layer, RoutedExperts): + if is_layer_skipped( + prefix=prefix, + ignored_layers=self.ignored_layers, + fused_mapping=self.packed_modules_mapping, + ): + return UnquantizedFusedMoEMethod(layer.moe_config) + if self.store_dtype == "mxfp4": + from vllm.model_executor.layers.quantization.mxfp4 import ( + Mxfp4MoEMethod, + ) + + return Mxfp4MoEMethod(layer.moe_config) + if self.is_checkpoint_fp8_serialized: + return Fp8MoEMethod(self, layer) + else: + from vllm.model_executor.layers.quantization.online.fp8 import ( + Fp8PerTensorOnlineMoEMethod, + ) + + return Fp8PerTensorOnlineMoEMethod(layer=layer) + elif isinstance(layer, Attention): + return Fp8KVCacheMethod(self) + return None + + @staticmethod + def get_cache_scale_mapper() -> "WeightsMapper": + """Map compressed-tensors KV-cache scale names to vLLM names.""" + from vllm.model_executor.models.utils import WeightsMapper + + orig_to_new_suffix = { + ".k_proj.output_scale": ".attn.k_scale", + ".v_proj.output_scale": ".attn.v_scale", + ".q_proj.output_scale": ".attn.q_scale", + ".self_attn.prob_output_scale": ".self_attn.attn.prob_scale", + } + cache_scale_mapper = WeightsMapper(orig_to_new_suffix=orig_to_new_suffix) + return cache_scale_mapper | QuantizationConfig.get_cache_scale_mapper() + + +class CopyNumelCounter(TorchDispatchMode): + """ + Tracks total number of elements modified with `copy_`. Useful for keeping + track of weight loading where underlying weights can be arbitrarily + transformed (such as with `narrow`) before calling copy. + """ + + def __init__(self): + super().__init__() + self.copied_numel = 0 + + def __torch_dispatch__(self, func, types, args=(), kwargs=None): + if kwargs is None: + kwargs = {} + out = func(*args, **kwargs) + if func == torch.ops.aten.copy_.default: + self.copied_numel += args[0].numel() + return out + + +def _copy_missing_attrs(old: torch.Tensor, new: torch.Tensor) -> None: + """Copies any attrs present in `old` but not in `new` to `new`""" + new_attrs = set(dir(new)) + attrs_to_set = {} + for attr in dir(old): + if attr not in new_attrs: + attrs_to_set[attr] = getattr(old, attr) + set_weight_attrs(new, attrs_to_set) + + +class Fp8LinearMethod(LinearMethodBase): + """Linear method for FP8. + Supports loading FP8 checkpoints with static weight scale and + dynamic/static activation scale. + + Limitations: + 1. Only support float8_e4m3fn data type due to the limitation of + torch._scaled_mm (https://github.com/pytorch/pytorch/blob/2e48b39603411a41c5025efbe52f89560b827825/aten/src/ATen/native/cuda/Blas.cpp#L854-L856) + + Args: + quant_config: The quantization config. + """ + + def __init__(self, quant_config: Fp8Config): + self.quant_config = quant_config + self.is_scale_e8m0 = getattr(quant_config, "is_scale_e8m0", False) + self.cutlass_block_fp8_supported = cutlass_block_fp8_supported() + self.out_dtype = torch.get_default_dtype() + self.input_dtype = get_current_vllm_config().model_config.dtype + + # For GPUs that lack FP8 hardware support, we can leverage the Marlin + # kernel for fast weight-only FP8 quantization + self.marlin_input_dtype = None + self.use_marlin = False + + if self.quant_config.use_deep_gemm is not None: + self.use_deep_gemm = self.quant_config.use_deep_gemm + else: + self.use_deep_gemm = is_deep_gemm_supported() + + self.weight_block_size = self.quant_config.weight_block_size + self.block_quant = self.weight_block_size is not None + self.act_q_static = self.quant_config.activation_scheme == "static" + + if self.block_quant: + assert not self.act_q_static + assert self.weight_block_size is not None + + self.activation_quant_key = create_fp8_quant_key( + static=self.act_q_static, + group_shape=GroupShape(1, self.weight_block_size[0]), + ) + self.weight_quant_key = create_fp8_quant_key( + static=True, group_shape=GroupShape(*self.weight_block_size) + ) + else: + self.weight_quant_key = kFp8StaticTensorSym + # Use per-token quantization for better perf if dynamic and cutlass + if self.act_q_static: + self.activation_quant_key = kFp8StaticTensorSym + elif cutlass_fp8_supported(): + self.activation_quant_key = kFp8DynamicTokenSym + else: + self.activation_quant_key = kFp8DynamicTensorSym + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: list[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + output_size_per_partition = sum(output_partition_sizes) + weight_loader = extra_weight_attrs.get("weight_loader") + layer.logical_widths = output_partition_sizes + layer.input_size_per_partition = input_size_per_partition + layer.output_size_per_partition = output_size_per_partition + layer.orig_dtype = params_dtype + layer.weight_block_size = None + + if self.block_quant: + assert self.weight_block_size is not None + layer.weight_block_size = self.weight_block_size + validate_fp8_block_shape( + layer, + input_size, + output_size, + input_size_per_partition, + output_partition_sizes, + self.weight_block_size, + ) + + weight = create_fp8_weight_parameter( + output_size_per_partition, input_size_per_partition, weight_loader + ) + layer.register_parameter("weight", weight) + + # WEIGHT SCALE + if not self.block_quant: + scale = create_fp8_scale_parameter( + PerTensorScaleParameter, + output_partition_sizes, + input_size_per_partition, + None, + weight_loader, + ) + layer.register_parameter("weight_scale", scale) + else: + assert not self.act_q_static + assert self.weight_block_size is not None + scale = create_fp8_scale_parameter( + BlockQuantScaleParameter, + output_partition_sizes, + input_size_per_partition, + self.weight_block_size, + weight_loader, + scale_dtype=(torch.float8_e8m0fnu if self.is_scale_e8m0 else None), + ) + # The weight_scale_inv name is intentional for deepseekv3 + layer.register_parameter("weight_scale_inv", scale) + + # INPUT ACTIVATION SCALE + if self.act_q_static: + scale = create_fp8_input_scale(output_partition_sizes, weight_loader) + set_weight_attrs(scale, {"scale_type": "input_scale"}) + layer.register_parameter("input_scale", scale) + + self.fp8_linear = init_fp8_linear_kernel( + activation_quant_key=self.activation_quant_key, + weight_quant_key=self.weight_quant_key, + weight_shape=layer.weight.shape, + input_dtype=self.input_dtype, + out_dtype=self.out_dtype, + module_name=self.__class__.__name__, + ) + + self.use_marlin = isinstance(self.fp8_linear, MarlinFP8ScaledMMLinearKernel) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + if self.use_marlin: + if not self.block_quant: + # Canonicalize to (K, N) for the kernel. + replace_parameter(layer, "weight", layer.weight.t()) + # Only Marlin kernels support `marlin_input_dtype`; guard to avoid + # AttributeError if backend selection changes. + if hasattr(self.fp8_linear, "marlin_input_dtype"): + self.fp8_linear.marlin_input_dtype = self.marlin_input_dtype + self.fp8_linear.process_weights_after_loading(layer) + return + + input_scale = None + # TODO(rob): refactor block quant into separate class. + if self.block_quant: + assert not self.act_q_static + + # If checkpoint not serialized fp8, quantize the weights. + else: + # If checkpoint is fp8 per-tensor, handle that there are N scales for N + # shards in a fused module + weight = layer.weight + weight_scale = layer.weight_scale + + # If using w8a8, torch._scaled_mm needs per tensor, so + # requantize the logical shards as a single weight. + weight, weight_scale, input_scale = process_fp8_weight_tensor_strategy( + weight, + weight_scale, + layer.logical_widths, + getattr(layer, "input_scale", None), + ) + if self.act_q_static: + assert input_scale is not None + input_scale = input_scale.max() + weight = weight.t() + + # Update layer with new values. + replace_parameter(layer, "weight", weight.data) + replace_parameter(layer, "weight_scale", weight_scale.data) + + if input_scale is not None: + replace_parameter(layer, "input_scale", input_scale) + else: + layer.input_scale = None + + self.fp8_linear.process_weights_after_loading(layer) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + # if batch invariant mode is enabled, prefer direct FP8 path + # we will use BF16 dequant when direct FP8 is not supported. + if envs.VLLM_BATCH_INVARIANT: + if self.block_quant: + assert self.weight_block_size is not None + return self.fp8_linear.apply_weights( + layer, + x, + bias, + ) + else: + if isinstance(self.fp8_linear, CutlassFP8ScaledMMLinearKernel): + return self.fp8_linear.apply_weights(layer, x, bias) + + # per-tensor/channel: dequant to BF16 and run GEMM + weight_fp8 = layer.weight.to(torch.bfloat16) + weight_scale = layer.weight_scale.to(torch.bfloat16) + if weight_scale.numel() == 1: + # Per-tensor: simple scalar multiplication + weight_bf16 = weight_fp8 * weight_scale + else: + # Multiple scales (fused modules like QKV) + # Try to infer correct broadcasting + # weight is [K, N], scale could be [num_logical_weights] + # Need to figure out how to broadcast - for now just try + # direct multiplication + if ( + weight_scale.dim() == 1 + and weight_scale.shape[0] == weight_fp8.shape[0] + ): + # Per-row scaling + weight_bf16 = weight_fp8 * weight_scale.unsqueeze(1) + else: + # Fallback + weight_bf16 = weight_fp8 * weight_scale + return torch.nn.functional.linear(x, weight_bf16.t(), bias) + + return self.fp8_linear.apply_weights(layer, x, bias) + + +class Fp8MoEMethod(FusedMoEMethodBase): + """MoE method for FP8. + Supports loading FP8 checkpoints with static weight scale and + dynamic/static activation scale. + + Also supports loading quantized FP16/BF16 model checkpoints with dynamic + activation scaling. The weight scaling factor will be initialized after + the model weights are loaded. + + Args: + quant_config: The quantization config. + """ + + def __init__(self, quant_config: Fp8Config, layer: RoutedExperts): + super().__init__(layer.moe_config) + self.quant_config = quant_config + self.weight_block_size = self.quant_config.weight_block_size + self.block_quant: bool = self.weight_block_size is not None + self.weight_scale_name = ( + "weight_scale_inv" if self.block_quant else "weight_scale" + ) + + # Set weight key and activation key for kernel compatibility + if self.block_quant: + weight_key = kFp8Static128BlockSym + activation_key = kFp8Dynamic128Sym + else: + weight_key = kFp8StaticTensorSym + activation_key = ( + kFp8StaticTensorSym + if self.quant_config.activation_scheme == "static" + else kFp8DynamicTensorSym + ) + + # Select Fp8 MoE backend + self.fp8_backend, self.experts_cls = select_fp8_moe_backend( + config=self.moe, + weight_key=weight_key, + activation_key=activation_key, + allow_vllm_cutlass=False, + ) + + def create_weights( + self, + layer: RoutedExperts, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + layer.num_experts = num_experts + layer.orig_dtype = params_dtype + layer.weight_block_size = None + + assert self.quant_config.is_checkpoint_fp8_serialized + params_dtype = torch.float8_e4m3fn + + if self.block_quant: + assert self.weight_block_size is not None + layer.weight_block_size = self.weight_block_size + tp_size = get_tensor_model_parallel_world_size() + block_n, block_k = ( + self.weight_block_size[0], + self.weight_block_size[1], + ) + # NOTE: To ensure proper alignment of the block-wise quantization + # scales, the output_size of the weights for both the gate and up + # layers must be divisible by block_n. + # Required by column parallel or enabling merged weights + if intermediate_size_per_partition % block_n != 0: + raise ValueError( + f"The output_size of gate's and up's weight = " + f"{intermediate_size_per_partition} is not divisible by " + f"weight quantization block_n = {block_n}." + ) + if tp_size > 1 and intermediate_size_per_partition % block_k != 0: + # Required by row parallel + raise ValueError( + f"The input_size of down's weight = " + f"{intermediate_size_per_partition} is not divisible by " + f"weight quantization block_k = {block_k}." + ) + + # WEIGHTS + w13_weight = torch.nn.Parameter( + torch.empty( + num_experts, + 2 * intermediate_size_per_partition, + hidden_size, + dtype=params_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight", w13_weight) + set_weight_attrs(w13_weight, extra_weight_attrs) + + w2_weight = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + intermediate_size_per_partition, + dtype=params_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight", w2_weight) + set_weight_attrs(w2_weight, extra_weight_attrs) + + # BIASES (for models like GPT-OSS that have biased MoE) + if self.moe.has_bias: + w13_bias = torch.nn.Parameter( + torch.zeros( + num_experts, + 2 * intermediate_size_per_partition, + dtype=layer.orig_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w13_bias", w13_bias) + set_weight_attrs(w13_bias, extra_weight_attrs) + w2_bias = torch.nn.Parameter( + torch.zeros(num_experts, hidden_size, dtype=layer.orig_dtype), + requires_grad=False, + ) + layer.register_parameter("w2_bias", w2_bias) + set_weight_attrs(w2_bias, extra_weight_attrs) + + # WEIGHT_SCALES + if not self.block_quant: + # For per-tensor quant, the scales are per expert and weight. + w13_scale_data = torch.ones(num_experts, 2, dtype=torch.float32) + w2_scale_data = torch.ones(num_experts, dtype=torch.float32) + else: + # For block quant, the scales are per block (typically 128x128). + w13_scale_data = torch.ones( + num_experts, + 2 * ((intermediate_size_per_partition + block_n - 1) // block_n), + (hidden_size + block_k - 1) // block_k, + dtype=torch.float32, + ) + w2_scale_data = torch.ones( + num_experts, + (hidden_size + block_n - 1) // block_n, + (intermediate_size_per_partition + block_k - 1) // block_k, + dtype=torch.float32, + ) + w13_weight_scale = torch.nn.Parameter(w13_scale_data, requires_grad=False) + w2_weight_scale = torch.nn.Parameter(w2_scale_data, requires_grad=False) + # Note: name is weight_scale for tensor, weight_scale_inv for block. + layer.register_parameter(f"w13_{self.weight_scale_name}", w13_weight_scale) + layer.register_parameter(f"w2_{self.weight_scale_name}", w2_weight_scale) + + # Add the quantization method used (per tensor/grouped/channel) + # to ensure the weight scales are loaded in properly + extra_weight_attrs.update( + {"quant_method": FusedMoeWeightScaleSupported.BLOCK.value} + if self.block_quant + else {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value} + ) + set_weight_attrs(w13_weight_scale, extra_weight_attrs) + set_weight_attrs(w2_weight_scale, extra_weight_attrs) + + # INPUT_SCALES + if self.quant_config.activation_scheme == "static": + assert not self.block_quant + w13_input_scale = torch.nn.Parameter( + torch.ones(num_experts, dtype=torch.float32), requires_grad=False + ) + layer.register_parameter("w13_input_scale", w13_input_scale) + set_weight_attrs(w13_input_scale, extra_weight_attrs) + + w2_input_scale = torch.nn.Parameter( + torch.ones(num_experts, dtype=torch.float32), requires_grad=False + ) + layer.register_parameter("w2_input_scale", w2_input_scale) + set_weight_attrs(w2_input_scale, extra_weight_attrs) + + else: + layer.w13_input_scale = None + layer.w2_input_scale = None + + def _setup_kernel( + self, + layer: RoutedExperts, + w13: torch.Tensor, + w2: torch.Tensor, + w13_scale: torch.Tensor, + w2_scale: torch.Tensor, + w13_input_scale: torch.Tensor | None, + w2_input_scale: torch.Tensor | None, + ) -> None: + # Shuffle weights to runtime format. + w13, w2, w13_scale, w2_scale = convert_to_fp8_moe_kernel_format( + fp8_backend=self.fp8_backend, + layer=layer, + w13=w13, + w2=w2, + w13_scale=w13_scale, + w2_scale=w2_scale, + w13_input_scale=w13_input_scale, + w2_input_scale=w2_input_scale, + ) + + # Replace parameters with updated versions. Note that this helper + # function ensures the replacement is compatible with RL weight reloads. + replace_parameter(layer, "w13_weight", w13) + replace_parameter(layer, "w2_weight", w2) + replace_parameter(layer, f"w13_{self.weight_scale_name}", w13_scale) + replace_parameter(layer, f"w2_{self.weight_scale_name}", w2_scale) + + # AITER backend requires weights to be marked as shuffled. + if self.fp8_backend == Fp8MoeBackend.AITER: + layer.w13_weight.is_shuffled = True + layer.w2_weight.is_shuffled = True + + self.moe_quant_config = self.get_fused_moe_quant_config(layer) + assert self.moe_quant_config is not None + assert self.experts_cls is not None + self.moe_kernel = make_fp8_moe_kernel( + moe_quant_config=self.moe_quant_config, + moe_config=self.moe, + fp8_backend=self.fp8_backend, + experts_cls=self.experts_cls, + routing_tables=layer._expert_routing_tables(), + layer=layer, + ) + + def process_weights_after_loading(self, layer: RoutedExperts) -> None: + # Allow for accessing weights and scales in standard way. + w13 = layer.w13_weight + w2 = layer.w2_weight + w13_scale = getattr(layer, f"w13_{self.weight_scale_name}") + w2_scale = getattr(layer, f"w2_{self.weight_scale_name}") + w13_input_scale = layer.w13_input_scale + w2_input_scale = layer.w2_input_scale + + # MI300x and MI325x use FNUZ format for FP8. Convert if needed. + if current_platform.is_fp8_fnuz(): + w13, w13_scale, w13_input_scale = normalize_e4m3fn_to_e4m3fnuz( + w13, + w13_scale, + w13_input_scale, + ) + w2, w2_scale, w2_input_scale = normalize_e4m3fn_to_e4m3fnuz( + w2, + w2_scale, + w2_input_scale, + ) + + # Per tensor kernels require single activation scale. Use the max. + if self.quant_config.activation_scheme == "static": + assert not self.block_quant + assert w13_input_scale is not None and w2_input_scale is not None + w13_input_scale, w2_input_scale = process_fp8_input_tensor_strategy_moe( + w13_input_scale, w2_input_scale + ) + replace_parameter(layer, "w13_input_scale", w13_input_scale) + replace_parameter(layer, "w2_input_scale", w2_input_scale) + + # Per tensor kernels require single weight scale for w13 per expert, but + # on disk there is a scale for w1 and w3. Use the max to requantize. + if not self.block_quant: + shard_size = layer.intermediate_size_per_partition + w13, w13_scale = process_fp8_weight_tensor_strategy_moe( + w13, w13_scale, shard_size, layer.local_num_experts + ) + + # Shuffle weights to runtime format and setup kernel. + self._setup_kernel( + layer, w13, w2, w13_scale, w2_scale, w13_input_scale, w2_input_scale + ) + + def maybe_make_prepare_finalize( + self, + routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, + ) -> mk.FusedMoEPrepareAndFinalizeModular | None: + raise ValueError( + f"{self.__class__.__name__} uses the new modular kernel initialization " + "logic. This function should not be called." + ) + + def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfig: + w1_scale = getattr(layer, f"w13_{self.weight_scale_name}") + w2_scale = getattr(layer, f"w2_{self.weight_scale_name}") + a1_scale = layer.w13_input_scale + a2_scale = layer.w2_input_scale + + quant_config = make_fp8_moe_quant_config( + fp8_backend=self.fp8_backend, + w1_scale=w1_scale, + w2_scale=w2_scale, + a1_scale=a1_scale, + a2_scale=a2_scale, + block_shape=self.weight_block_size, + swiglu_limit=getattr(layer, "swiglu_limit", None), + gemm1_alpha=getattr(layer, "swiglu_alpha", None), + gemm1_beta=getattr(layer, "swiglu_beta", None), + layer=layer, + ) + + # Inject biases into the quant config if the model has them + # (e.g. GPT-OSS biased MoE) + if quant_config is not None and self.moe.has_bias: + w13_bias = getattr(layer, "w13_bias", None) + w2_bias = getattr(layer, "w2_bias", None) + if w13_bias is not None: + quant_config._w1.bias = w13_bias + if w2_bias is not None: + quant_config._w2.bias = w2_bias + + return quant_config + + @property + def supports_eplb(self) -> bool: + return True + + def apply_monolithic( + self, + layer: RoutedExperts, + x: torch.Tensor, + router_logits: torch.Tensor, + input_ids: torch.Tensor | None = None, + ) -> torch.Tensor: + assert self.is_monolithic + assert self.moe_kernel is not None + return self.moe_kernel.apply_monolithic( + x, + layer.w13_weight, + layer.w2_weight, + router_logits, + activation=layer.activation, + global_num_experts=layer.global_num_experts, + expert_map=layer.expert_map, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + num_expert_group=layer.num_expert_group, + topk_group=layer.topk_group, + e_score_correction_bias=layer.e_score_correction_bias, + routed_scaling_factor=layer.routed_scaling_factor, + ) + + def apply( + self, + layer: RoutedExperts, + x: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + shared_experts: SharedExperts | None, + shared_experts_input: torch.Tensor | None, + ) -> torch.Tensor: + assert not self.is_monolithic + assert self.moe_kernel is not None + return self.moe_kernel.apply( + x, + layer.w13_weight, + layer.w2_weight, + topk_weights, + topk_ids, + activation=layer.activation, + global_num_experts=layer.global_num_experts, + expert_map=layer.expert_map, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + shared_experts=shared_experts, + shared_experts_input=shared_experts_input, + ) + + +class Fp8KVCacheMethod(BaseKVCacheMethod): + """ + Supports loading kv-cache scaling factors from FP8 checkpoints. + """ + + def __init__(self, quant_config: Fp8Config): + super().__init__(quant_config) diff --git a/upstream/vllm/models/deepseek_v4/attention.py b/upstream/vllm/models/deepseek_v4/attention.py new file mode 100644 index 0000000..fedea3c --- /dev/null +++ b/upstream/vllm/models/deepseek_v4/attention.py @@ -0,0 +1,824 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +DeepseekV4 MLA Attention Layer +""" + +from abc import ABC, abstractmethod +from collections.abc import Callable +from typing import TYPE_CHECKING, Any, ClassVar, cast + +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers import DeepseekV2Config, DeepseekV3Config + +import vllm.envs as envs +from vllm.compilation.breakable_cudagraph import eager_break_during_capture +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer +from vllm.models.deepseek_v4.common.ops import ( + fused_indexer_q_rope_quant, + fused_q_kv_rmsnorm, +) +from vllm.utils.deep_gemm import fp8_einsum +from vllm.utils.torch_utils import direct_register_custom_op + +if TYPE_CHECKING: + from vllm.v1.attention.backends.mla.sparse_swa import ( + DeepseekSparseSWAMetadata, + ) + +from vllm.config import ( + CacheConfig, + VllmConfig, + get_current_vllm_config, +) +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.forward_context import get_forward_context +from vllm.logger import init_logger +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.models.utils import extract_layer_index +from vllm.models.deepseek_v4.common.rope import build_deepseek_v4_rope +from vllm.models.deepseek_v4.compressor import DeepseekCompressor +from vllm.utils.multi_stream_utils import ( + execute_in_parallel, + maybe_execute_in_parallel, +) +from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata +from vllm.v1.attention.backends.mla.indexer import ( + DeepseekV4IndexerBackend, + get_max_prefill_buffer_size, +) +from vllm.v1.attention.backends.mla.sparse_swa import DeepseekV4SWACache +from vllm.v1.kv_cache_interface import ( + KVCacheSpec, + MLAAttentionSpec, + get_kv_quant_mode, +) + +logger = init_logger(__name__) + + +def _resolve_dsv4_kv_cache_dtype( + use_fp8_ds_mla_layout: bool, + kv_cache_dtype: str, + cache_config: CacheConfig | None, +) -> tuple[str, torch.dtype]: + """Map ``(layout, --kv-cache-dtype)`` to ``(cache_dtype_str, torch_dtype)``. + + Both layouts are paged; they differ in the per-token block format. The + ``fp8_ds_mla`` format is UE8M0 block-scaled fp8 packed as ``uint8`` (the + canonical ``fp8_ds_mla`` string is written back onto ``cache_config`` so the + page-size specs pick the 576B per-token slot). Plain-row backends store each + token's KV row in its element dtype: bf16 or per-tensor FP8 E4M3. + """ + if use_fp8_ds_mla_layout: + # fp8_ds_mla block format: UE8M0 block-scaled fp8 packed as uint8. + assert kv_cache_dtype.startswith("fp8"), ( + f"DeepseekV4 fp8_ds_mla layout only supports fp8 kv-cache, " + f"got {kv_cache_dtype}" + ) + if kv_cache_dtype != "fp8_ds_mla": + if cache_config is not None: + cache_config.cache_dtype = "fp8_ds_mla" + kv_cache_dtype = "fp8_ds_mla" + logger.info_once("Using DeepSeek's fp8_ds_mla KV cache format.") + return kv_cache_dtype, torch.uint8 + + # Plain bf16 / per-tensor fp8 KV row (FlashInfer). + if kv_cache_dtype.startswith("fp8"): + return kv_cache_dtype, torch.float8_e4m3fn + # auto / bfloat16 -> plain bf16 KV row. + return kv_cache_dtype, torch.bfloat16 + + +class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): + """DeepseekV4 MLA attention layer. + + The platform-specific sparse-MLA forward (``forward_mqa`` / + ``get_padded_num_q_heads`` / ``_o_proj`` / ``backend_cls``) is provided by a + subclass — ``DeepseekV4FlashMLAAttention`` / + ``DeepseekV4FlashInferSM120Attention`` / + ``DeepseekV4FlashInferMLAAttention`` (CUDA) or + ``DeepseekV4ROCMAiterMLAAttention`` (ROCm) — selected by the platform-specific + deepseek_v4 model module. The base is never instantiated directly. + """ + + # Provided by the platform subclass. + backend_cls: ClassVar[type[AttentionBackend]] + # KV-cache per-token block format (both layouts are paged). True (default) + # = fp8_ds_mla (UE8M0 block-scaled fp8 packed as uint8); False = plain + # bf16 / per-tensor fp8 KV row. Backends can override the instance hook when + # a single attention class dispatches across arch-specific layouts. + use_fp8_ds_mla_layout: ClassVar[bool] = True + # Prefill is processed in fixed-size chunks; this bounds the bf16 kv-gather + # workspace allocated in _forward_prefill and is also read by the dummy-run + # path to pre-reserve that workspace. + PREFILL_CHUNK_SIZE: ClassVar[int] = 4 + + @classmethod + @abstractmethod + def get_padded_num_q_heads(cls, num_heads: int) -> int: + """Q head count the q/output buffers are allocated at. + + The layer allocates the q/output buffers at + ``[N, get_padded_num_q_heads(n_local_heads), head_dim]``. Must satisfy + ``result >= num_heads``. Backends with no padding constraint return + ``num_heads``. + """ + raise NotImplementedError + + @abstractmethod + def forward_mqa( + self, + q: torch.Tensor, + kv: torch.Tensor, + positions: torch.Tensor, + output: torch.Tensor, + ) -> None: + """Platform-specific sparse MLA forward; writes attention into ``output``.""" + raise NotImplementedError + + @abstractmethod + def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + """Inverse-RoPE + wo_a + wo_b output projection (platform-specific).""" + raise NotImplementedError + + def _uses_fp8_ds_mla_layout(self) -> bool: + """Return whether this instance stores fp8 KV in fp8_ds_mla layout.""" + return self.use_fp8_ds_mla_layout + + def __init__( + self, + vllm_config: VllmConfig, + prefix: str, + topk_indices_buffer: torch.Tensor | None = None, + aux_stream_list: list[torch.cuda.Stream] | None = None, + ) -> None: + super().__init__() + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + cache_config = vllm_config.cache_config + tp_size = get_tensor_model_parallel_world_size() + layer_id = extract_layer_index(prefix) + + self.prefix = prefix # Alias for compatibility with compressor + self.hidden_size = config.hidden_size + self.n_heads = config.num_attention_heads + assert self.n_heads % tp_size == 0 + self.n_local_heads = self.n_heads // tp_size + self.q_lora_rank = config.q_lora_rank + self.o_lora_rank = config.o_lora_rank + self.head_dim = config.head_dim + self.rope_head_dim = config.qk_rope_head_dim + self.nope_head_dim = self.head_dim - self.rope_head_dim + self.n_groups = config.o_groups + self.n_local_groups = self.n_groups // tp_size + self.window_size = config.sliding_window + # NOTE(zyongye) Compress ratio can't be 0 + # we do this for because MTP layer is not included + # in the compress ratio list + if layer_id < config.num_hidden_layers: + self.compress_ratio = max(1, config.compress_ratios[layer_id]) + else: + self.compress_ratio = 1 + self.eps = config.rms_norm_eps + self.scale = self.head_dim**-0.5 + + # Padded Q head count is dictated by the platform subclass. + self.padded_heads = self.get_padded_num_q_heads(self.n_local_heads) + # Sink padded to the same head count, initialized to -inf (no sink + # effect). Weight loading fills the first n_local_heads slots. + self.attn_sink = nn.Parameter( + torch.full((self.padded_heads,), -float("inf"), dtype=torch.float32), + requires_grad=False, + ) + + self.fused_wqa_wkv = MergedColumnParallelLinear( + self.hidden_size, + [self.q_lora_rank, self.head_dim], + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.fused_wqa_wkv", + disable_tp=True, # fused ReplicatedLinear + ) + self.q_norm = RMSNorm(self.q_lora_rank, self.eps) + self.wq_b = ColumnParallelLinear( + self.q_lora_rank, + self.n_heads * self.head_dim, + bias=False, + quant_config=quant_config, + return_bias=False, + prefix=f"{prefix}.wq_b", + ) + + self.kv_norm = RMSNorm(self.head_dim, self.eps) + # wo_a uses the standard quant_config for all quantization schemes + # (FP8 blockwise, FP8 channelwise, INT8 channelwise). The _o_proj + # method in the platform subclass dispatches at runtime based on + # weight dtype and scale attribute: + # - int8 + weight_scale → INT8 channelwise einsum + # - fp8 + weight_scale → FP8 channelwise einsum + # - fp8 + weight_scale_inv → FP8 blockwise einsum + self.wo_a = ColumnParallelLinear( + self.n_heads * self.head_dim // self.n_groups, + self.n_groups * self.o_lora_rank, + bias=False, + quant_config=quant_config, + return_bias=False, + prefix=f"{prefix}.wo_a", + ) + self.wo_a.is_bmm = True + self.wo_a.bmm_batch_size = self.n_local_groups + self.wo_b = RowParallelLinear( + self.n_groups * self.o_lora_rank, + self.hidden_size, + bias=False, + quant_config=quant_config, + return_bias=False, + prefix=f"{prefix}.wo_b", + ) + + # Initialize rotary embedding before the indexer/compressor consume it. + self.rotary_emb = build_deepseek_v4_rope( + config, + head_dim=self.head_dim, + rope_head_dim=self.rope_head_dim, + max_position_embeddings=config.max_position_embeddings, + compress_ratio=self.compress_ratio, + ) + self.indexer_rotary_emb = self.rotary_emb + self.topk_indices_buffer = topk_indices_buffer + + self.indexer = None + if self.compress_ratio == 4: + # Only C4A uses sparse attention and hence has indexer. + # aux_stream_list[2] is free here (outer GEMMs joined) for the inner + # overlap of wq_b+fused_indexer_q_rope_quant vs compressor. None on + # ROCm, where aux_stream_list is None. + indexer_aux_stream = ( + aux_stream_list[2] if aux_stream_list is not None else None + ) + self.indexer = DeepseekV4Indexer( + vllm_config, + config=config, + hidden_size=self.hidden_size, + q_lora_rank=self.q_lora_rank, + quant_config=quant_config, + cache_config=cache_config, + topk_indices_buffer=topk_indices_buffer, + compress_ratio=self.compress_ratio, + prefix=f"{prefix}.indexer", + aux_stream=indexer_aux_stream, + ) + + # Will be None on ROCm for now. + self.aux_stream_list = aux_stream_list + # [0]: GEMM start / post-GEMM event0. [1..3]: GEMM done events; + # [1] doubles as post-GEMM event1. Reuse is safe: GEMM fully joins + # before post-GEMM starts. + self.ln_events = [torch.cuda.Event() for _ in range(4)] + + assert cache_config is not None, "DeepseekV4 attention requires cache_config" + # ---- Attention / KV-cache setup ---- + self.max_num_batched_tokens = ( + vllm_config.scheduler_config.max_num_batched_tokens + ) + self.max_model_len = vllm_config.model_config.max_model_len + + # Resolve the kv-cache dtype from this backend's block format. The same + # resolution drives the SWA cache tensor dtype below. + self.kv_cache_dtype, self.kv_cache_torch_dtype = _resolve_dsv4_kv_cache_dtype( + self._uses_fp8_ds_mla_layout(), cache_config.cache_dtype, cache_config + ) + + self.swa_cache_layer = DeepseekV4SWACache( + head_dim=self.head_dim, + window_size=self.window_size, + dtype=self.kv_cache_torch_dtype, + prefix=f"{prefix}.swa_cache", + cache_config=cache_config, + ) + + # Register with compilation context for metadata lookup. + compilation_config = vllm_config.compilation_config + if prefix and prefix in compilation_config.static_forward_context: + raise ValueError(f"Duplicate layer name: {prefix}") + if prefix: + compilation_config.static_forward_context[prefix] = self + self.kv_cache = torch.tensor([]) + + # Create the compressor for layers with compress_ratio > 1; after the + # attention setup above so its KV-cache prefix (self.prefix) is set. + self.compressor = None + if self.compress_ratio > 1: + self.compressor = DeepseekCompressor( + vllm_config=vllm_config, + compress_ratio=self.compress_ratio, + hidden_size=self.hidden_size, + head_dim=self.head_dim, + rotate=True, + prefix=f"{prefix}.compressor", + k_cache_prefix=self.prefix, + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + llama_4_scaling: torch.Tensor | None = None, + ) -> torch.Tensor: + # Pre-allocate attention output with FlashMLA-padded head count. + # The op writes into `o_padded`; we slice to n_local_heads after. + num_tokens = hidden_states.shape[0] + o_padded = torch.empty( + (num_tokens, self.padded_heads, self.head_dim), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + + # Metadata-independent input GEMMs + RMSNorm stay in the captured + # graph; the metadata-dependent rest (q up-proj + kv-insert, indexer, + # compressor, MLA attention) runs in the eager break. + qr_kv, kv_score, indexer_kv_score, indexer_weights = ( + self.attn_gemm_parallel_execute(hidden_states) + ) + qr, kv = qr_kv.split([self.q_lora_rank, self.head_dim], dim=-1) + qr, kv = fused_q_kv_rmsnorm( + qr, + kv, + self.q_norm.weight.data, + self.kv_norm.weight.data, + self.eps, + ) + + # attention_impl is wrapped with @eager_break_during_capture: this is + # where the breakable cudagraph capture breaks (the attention op runs + # eagerly between captured graph segments). + self.attention_impl( + hidden_states, + qr, + kv, + kv_score, + indexer_kv_score, + indexer_weights, + positions, + o_padded, + ) + o = o_padded[:, : self.n_local_heads, :] + + # Inverse-RoPE + wo_a + wo_b output projection (platform-specific). + return self._o_proj(o, positions) + + def attn_gemm_parallel_execute(self, hidden_states) -> tuple[Any, ...]: + aux_streams = self.aux_stream_list + if aux_streams is not None: + assert len(aux_streams) >= 3 + aux_streams = aux_streams[:3] + + # fused_wqa_wkv (heaviest) on default; the three lighter input GEMMs + # on aux streams 0..2 when their owning module exists. ln_events[0] + # is the fan-out start event; ln_events[1..3] are per-aux done events. + # On ROCm, aux_streams is None and execute_in_parallel runs serially. + aux_fns: list[Callable[[], Any] | None] = [None, None, None] + + if self.compressor is not None: + # Local ref so the closure keeps a non-None type for mypy. + compressor = self.compressor + + def compressor_kv_score() -> torch.Tensor: + return torch.mm( + hidden_states, + compressor.fused_wkv_wgate.weight.T, + out_dtype=torch.float32, + ) + + aux_fns[0] = compressor_kv_score + + if self.indexer is not None: + indexer = self.indexer + + def indexer_weights_proj() -> torch.Tensor: + # ReplicatedLinear returns (output, bias); bias is None. + weights, _ = indexer.weights_proj(hidden_states) + return weights + + def indexer_compressor_kv_score() -> torch.Tensor: + return torch.mm( + hidden_states, + indexer.compressor.fused_wkv_wgate.weight.T, + out_dtype=torch.float32, + ) + + aux_fns[1] = indexer_weights_proj + aux_fns[2] = indexer_compressor_kv_score + + def fused_wqa_wkv() -> torch.Tensor: + # MergedColumnParallelLinear returns (output, bias); bias is None. + qr_kv, _ = self.fused_wqa_wkv(hidden_states) + return qr_kv + + qr_kv, (kv_score, indexer_weights, indexer_kv_score) = execute_in_parallel( + fused_wqa_wkv, + aux_fns, + self.ln_events[0], + self.ln_events[1:4], + aux_streams, + enable=hidden_states.shape[0] + <= envs.VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD, + ) + + return qr_kv, kv_score, indexer_kv_score, indexer_weights + + @eager_break_during_capture + def attention_impl( + self, + hidden_states: torch.Tensor, + qr: torch.Tensor, + kv: torch.Tensor, + kv_score: torch.Tensor, + indexer_kv_score: torch.Tensor, + indexer_weights: torch.Tensor, + positions: torch.Tensor, + out: torch.Tensor, # [num_tokens, padded_heads, head_dim], written in place + ) -> None: + forward_context = get_forward_context() + attn_metadata = forward_context.attn_metadata + + # wq_b + kv_insert (+ MLA compressor when an indexer is present) ride + # on the default stream so q stays on its consumer stream (forward_mqa + # downstream reads q on default). Indexer/compressor go on aux for + # overlap with default's GEMM + cache write. + if self.indexer is not None: + aux_streams = self.aux_stream_list + indexer = self.indexer + # Local ref so the closure keeps a non-None type for mypy. + assert self.compressor is not None + compressor = self.compressor + + def wq_b_kv_insert() -> torch.Tensor: + q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim) + q = self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata) + return q + + # 3-way overlap (matches TRT-LLM PR #14142 Level 1): default runs + # wq_b+kv_insert; slot [0] runs the full indexer; slot [1] runs the + # MLA compressor. Slot [2] is reserved for the indexer's inner + # overlap. ROCm (aux_streams is None) falls back to sequential. + q, _ = execute_in_parallel( + wq_b_kv_insert, + [ + lambda: indexer( + hidden_states, + qr, + indexer_kv_score, + indexer_weights, + positions, + self.indexer_rotary_emb, + ), + lambda: compressor(kv_score, positions, self.rotary_emb), + ], + self.ln_events[0], + [self.ln_events[1], self.ln_events[2]], + [aux_streams[0], aux_streams[1]] if aux_streams is not None else None, + enable=aux_streams is not None, + ) + elif self.compressor is not None: + # wq_b + kv_insert on default, compressor on aux. + aux_stream = ( + self.aux_stream_list[0] if self.aux_stream_list is not None else None + ) + compressor = self.compressor + + def wq_b_kv_insert() -> torch.Tensor: + q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim) + q = self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata) + return q + + q, _ = maybe_execute_in_parallel( + wq_b_kv_insert, + lambda: compressor(kv_score, positions, self.rotary_emb), + self.ln_events[0], + self.ln_events[1], + aux_stream, + ) + else: + # SWA-only layer: no compressor, no overlap. + q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim) + q = self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata) + + # MLA attention writes into the pre-allocated `out` buffer + # ([num_tokens, padded_heads, head_dim]). + self.forward_mqa(q, kv, positions, out) + + def _fused_qnorm_rope_kv_insert( + self, + q: torch.Tensor, + kv: torch.Tensor, + positions: torch.Tensor, + attn_metadata: ( + dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None + ), + ) -> torch.Tensor: + if not isinstance(attn_metadata, dict): + # Profile run: kernel doesn't fire; produce a padded tensor so + # downstream FlashMLA gets the right shape. + if self.n_local_heads < self.padded_heads: + return F.pad( + q, + (0, 0, 0, self.padded_heads - self.n_local_heads), + value=0.0, + ) + return q + + swa_metadata = cast( + "DeepseekSparseSWAMetadata | None", + attn_metadata.get(self.swa_cache_layer.prefix), + ) + assert swa_metadata is not None + + swa_kv_cache = self.swa_cache_layer.kv_cache + # The fused insert ops require int64 position_ids; the runner's positions + # buffer is already int64, so no cast is needed. + assert positions.dtype == torch.int64 + cos_sin_cache = self.rotary_emb.cos_sin_cache + cache_dtype = swa_kv_cache.dtype + + # kv is unchanged; attention reads kv solely via swa_kv_cache. + if cache_dtype == torch.uint8: + # fp8_ds_mla UE8M0 paged path. Horizontally fused: + # Q side: per-head RMSNorm (no weight) + GPT-J RoPE, zero-filling + # the padding head slots; the kernel allocates and returns + # the padded q tensor. + # KV side: GPT-J RoPE + UE8M0 FP8 quant + paged cache insert. + swa_kv_cache_2d = swa_kv_cache.view(swa_kv_cache.shape[0], -1) + return torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + q, + kv, + swa_kv_cache_2d, + swa_metadata.slot_mapping, + positions, + cos_sin_cache, + self.padded_heads, + self.eps, + swa_metadata.block_size, + ) + + # Plain-row path: the [num_blocks, block_size, 512] cache stores the KV + # row in its element dtype (no Q padding). bf16 rewrites q in place; + # per-tensor fp8 writes a separately-allocated fp8 q and quantizes the + # KV row. + block_size = swa_metadata.block_size + swa_kv_cache_3d = swa_kv_cache.view(-1, block_size, self.head_dim) + if cache_dtype == torch.bfloat16: + torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_bf16_insert( + q, + kv, + swa_kv_cache_3d, + swa_metadata.slot_mapping, + positions, + cos_sin_cache, + self.eps, + block_size, + ) + return q + + # per-tensor fp8 (torch.float8_e4m3fn) + q_fp8 = torch.empty_like(q, dtype=torch.float8_e4m3fn) + torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_full_cache_fp8_insert( + q, + kv, + q_fp8, + swa_kv_cache_3d, + swa_metadata.slot_mapping, + positions, + cos_sin_cache, + self._flashinfer_fp8_kv_scale, + self._flashinfer_fp8_q_scale_inv, + self.eps, + block_size, + ) + return q_fp8 + + def get_attn_backend(self) -> type[AttentionBackend]: + return self.backend_cls + + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: + if ( + self.compress_ratio <= 1 + ): # SWA part. Allocated separately as DeepseekV4SWACache. + return None + # fp8_ds_mla is a UE8M0 block-scaled uint8 layout and needs 576B + # alignment; plain bf16 / per-tensor fp8 rows use natural element-size + # pages. + uses_fp8_ds_mla_layout = self.kv_cache_dtype == "fp8_ds_mla" + return MLAAttentionSpec( + block_size=vllm_config.cache_config.block_size, + num_kv_heads=1, + head_size=self.head_dim, + dtype=torch.uint8 if uses_fp8_ds_mla_layout else self.kv_cache_torch_dtype, + compress_ratio=self.compress_ratio, + cache_dtype_str=self.kv_cache_dtype, + alignment=576 if uses_fp8_ds_mla_layout else 512, + model_version="deepseek_v4", + kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype), + ) + + +class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase): + def __init__( + self, + head_dim: int, + dtype: torch.dtype, + prefix: str, + cache_config: CacheConfig, + compress_ratio: int = 1, + n_head: int | None = None, + q_head_dim: int | None = None, + ): + super().__init__() + self.kv_cache = torch.tensor([]) + self.head_dim = head_dim + self.prefix = prefix + self.cache_config = cache_config + self.dtype = dtype + self.compress_ratio = compress_ratio + self.n_head = n_head + self.q_head_dim = q_head_dim + compilation_config = get_current_vllm_config().compilation_config + if prefix in compilation_config.static_forward_context: + raise ValueError(f"Duplicate layer name: {prefix}") + compilation_config.static_forward_context[prefix] = self + + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: + # head_dim already carries the fp8 scale padding + # compress_ratio=1 for V3.2, >1 for DeepseekV4; both use the same cache layout. + uses_fp8_ds_mla_layout = vllm_config.cache_config.cache_dtype == "fp8_ds_mla" + return MLAAttentionSpec( + block_size=self.cache_config.block_size, + num_kv_heads=1, + head_size=self.head_dim, + dtype=self.dtype, + compress_ratio=self.compress_ratio, + # 576B for FlashMLA packing; 512B for FlashInfer sparse (#44577). + alignment=576 if uses_fp8_ds_mla_layout else 512, + indexer_n_head=self.n_head, + indexer_q_head_dim=self.q_head_dim, + ) + + def forward(self): ... + + def get_attn_backend(self) -> type[AttentionBackend]: + return DeepseekV4IndexerBackend + + +class DeepseekV4Indexer(nn.Module): + def __init__( + self, + vllm_config: VllmConfig, + config: DeepseekV2Config | DeepseekV3Config, + hidden_size: int, + q_lora_rank: int, + quant_config: QuantizationConfig | None, + cache_config: CacheConfig | None, + topk_indices_buffer: torch.Tensor | None, + compress_ratio: int = 1, + prefix: str = "", + aux_stream: torch.cuda.Stream | None = None, + ): + super().__init__() + self.vllm_config = vllm_config + self.config = config + self.quant_config = quant_config + # self.indexer_cfg = config.attn_module_list_cfg[0]["attn_index"] + self.topk_tokens = config.index_topk + self.n_head = config.index_n_heads # 64 + self.head_dim = config.index_head_dim # 128 + self.rope_dim = config.qk_rope_head_dim # 64 + self.q_lora_rank = q_lora_rank # 1536 + self.compress_ratio = compress_ratio + self.use_fp4_kv = self.vllm_config.attention_config.use_fp4_indexer_cache + logger.info_once( + "Using %s indexer cache for Lightning Indexer.", + "MXFP4" if self.use_fp4_kv else "FP8", + ) + + # no tensor parallel, just replicated + self.wq_b = ReplicatedLinear( + self.q_lora_rank, + self.head_dim * self.n_head, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.wq_b", + ) + self.weights_proj = ReplicatedLinear( + hidden_size, + self.n_head, + bias=False, + quant_config=None, + prefix=f"{prefix}.weights_proj", + ) + self.softmax_scale = self.head_dim**-0.5 + + self.scale_fmt = "ue8m0" + self.quant_block_size = 128 # TODO: get from config + self.topk_indices_buffer = topk_indices_buffer + + self.max_model_len = ( + vllm_config.model_config.max_model_len // self.compress_ratio + ) + self.prefix = prefix + + self.max_total_seq_len = ( + get_max_prefill_buffer_size(vllm_config) // self.compress_ratio + ) + + assert cache_config is not None, "Deepseek V4 indexer requires cache_config" + # NOTE(yifan): FP8 indxer cache use the same layout as V3.2: + # head_dim bytes = 128 fp8 + 4 fp32 scale = 132. + # For FP4 indexer cache, we still allocate the same amount of memory as FP8, + # but only use the first half of the memory. + k_cache_head_dim = self.head_dim + self.head_dim // self.quant_block_size * 4 + self.k_cache = DeepseekV4IndexerCache( + head_dim=k_cache_head_dim, + dtype=torch.uint8, + prefix=f"{prefix}.k_cache", + cache_config=cache_config, + compress_ratio=self.compress_ratio, + n_head=self.n_head, + q_head_dim=self.head_dim, + ) + self.compressor = DeepseekCompressor( + vllm_config=vllm_config, + compress_ratio=self.compress_ratio, + hidden_size=hidden_size, + head_dim=self.head_dim, + rotate=True, + prefix=f"{prefix}.compressor", + k_cache_prefix=self.k_cache.prefix, + use_fp4_cache=self.use_fp4_kv, + ) + + self.indexer_op = SparseAttnIndexer( + self.k_cache, + self.quant_block_size, + self.scale_fmt, + self.topk_tokens, + self.head_dim, + self.max_model_len, + self.max_total_seq_len, + self.topk_indices_buffer, + skip_k_cache_insert=True, + use_fp4_cache=self.use_fp4_kv, + ) + + # None on ROCm — maybe_execute_in_parallel falls back to sequential. + self.aux_stream = aux_stream + self.ln_events: list[torch.cuda.Event] = [ + torch.cuda.Event(), + torch.cuda.Event(), + ] + + def forward( + self, + hidden_states: torch.Tensor, + qr: torch.Tensor, + compressed_kv_score: torch.Tensor, + indexer_weights: torch.Tensor, + positions: torch.Tensor, + rotary_emb: nn.Module, + ) -> torch.Tensor: + compressor = self.compressor + + def wq_b_and_q_quant(): + # ReplicatedLinear returns (output, bias); bias is None. + q, _ = self.wq_b(qr) + q = q.view(-1, self.n_head, self.head_dim) + return fused_indexer_q_rope_quant( + positions, + q, + rotary_emb.cos_sin_cache, + indexer_weights, + self.softmax_scale, + self.n_head**-0.5, + use_fp4=self.use_fp4_kv, + ) + + # compressor returns None and writes K to the indexer KV cache; the + # join orders that write before indexer_op (skip_k_cache_insert=True). + (q_quant, weights), k = maybe_execute_in_parallel( + wq_b_and_q_quant, + lambda: compressor(compressed_kv_score, positions, rotary_emb), + self.ln_events[0], + self.ln_events[1], + self.aux_stream, + ) + return self.indexer_op(hidden_states, q_quant, k, weights) diff --git a/upstream/vllm/models/deepseek_v4/ppu/flashmla.py b/upstream/vllm/models/deepseek_v4/ppu/flashmla.py new file mode 100644 index 0000000..7e6b944 --- /dev/null +++ b/upstream/vllm/models/deepseek_v4/ppu/flashmla.py @@ -0,0 +1,393 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import TYPE_CHECKING, cast + +import torch + +from vllm.forward_context import get_forward_context +from vllm.models.deepseek_v4.attention import DeepseekV4Attention +from vllm.platforms import current_platform +from vllm.models.deepseek_v4.common.ops import ( + combine_topk_swa_indices, + compute_global_topk_indices_and_lens, + dequantize_and_gather_k_cache, +) +from vllm.models.deepseek_v4.ppu.ops.o_proj import ( + compute_fp8_einsum_recipe, + deep_gemm_fp8_o_proj, + deep_gemm_fp8_channel_o_proj, + deep_gemm_int8_o_proj, +) +from vllm.models.deepseek_v4.sparse_mla import ( + DeepseekV4FlashMLABackend, + DeepseekV4FlashMLAMetadata, +) +from vllm.v1.attention.ops.flashmla import ( + flash_mla_sparse_fwd, + flash_mla_with_kvcache, +) +from vllm.v1.worker.workspace import current_workspace_manager + +if TYPE_CHECKING: + from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata + + +class DeepseekV4FlashMLAAttention(DeepseekV4Attention): + """FlashMLA sparse MLA attention layer for DeepSeek V4 (CUDA).""" + + backend_cls = DeepseekV4FlashMLABackend + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + self._einsum_recipe, self._tma_aligned_scales = compute_fp8_einsum_recipe() + + def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: + if self.wo_a.weight.dtype == torch.int8: + return deep_gemm_int8_o_proj( + o, + positions, + self.rotary_emb.cos_sin_cache, + self.wo_a, + self.wo_b, + n_groups=self.n_local_groups, + heads_per_group=self.n_local_heads // self.n_local_groups, + nope_dim=self.nope_head_dim, + rope_dim=self.rope_head_dim, + o_lora_rank=self.o_lora_rank, + einsum_recipe=self._einsum_recipe, + ) + elif self.wo_a.weight.dtype == torch.float8_e4m3fn and hasattr( + self.wo_a, "weight_scale" + ): + # Channelwise FP8 (PPU): wo_a kept as FP8 + per-output-channel + # weight_scale (no pre-dequant). + return deep_gemm_fp8_channel_o_proj( + o, + positions, + self.rotary_emb.cos_sin_cache, + self.wo_a, + self.wo_b, + n_groups=self.n_local_groups, + heads_per_group=self.n_local_heads // self.n_local_groups, + nope_dim=self.nope_head_dim, + rope_dim=self.rope_head_dim, + o_lora_rank=self.o_lora_rank, + einsum_recipe=self._einsum_recipe, + ) + return deep_gemm_fp8_o_proj( + o, + positions, + self.rotary_emb.cos_sin_cache, + self.wo_a, + self.wo_b, + n_groups=self.n_local_groups, + heads_per_group=self.n_local_heads // self.n_local_groups, + nope_dim=self.nope_head_dim, + rope_dim=self.rope_head_dim, + o_lora_rank=self.o_lora_rank, + einsum_recipe=self._einsum_recipe, + tma_aligned_scales=self._tma_aligned_scales, + ) + + @classmethod + def get_padded_num_q_heads(cls, num_heads: int) -> int: + # FP8 decode kernel only supports h_q = 64 or 128. + if num_heads > 128: + raise ValueError( + f"DeepseekV4 FlashMLA does not support {num_heads} heads " + "(FP8 decode kernel requires h_q in {64, 128})." + ) + if current_platform.is_ppu(): + # PPU FlashMLA supports arbitrary head counts; no padding needed. + return num_heads + return 64 if num_heads <= 64 else 128 + + def forward_mqa( + self, + q: torch.Tensor, + kv: torch.Tensor, + positions: torch.Tensor, + output: torch.Tensor, + ) -> None: + assert output.shape == q.shape, ( + f"output buffer shape {output.shape} must match q shape {q.shape}" + ) + assert output.dtype == q.dtype, ( + f"output buffer dtype {output.dtype} must match q dtype {q.dtype}" + ) + + # Get SWA and indexer metadata from forward context + forward_context = get_forward_context() + attn_metadata = forward_context.attn_metadata + + if attn_metadata is None: + # Warmup dummy run: no real metadata. Reserve the same bf16 + # gather workspace _forward_prefill would; the dequantize / topk + # / sparse_fwd kernels are skipped this step. + swa_only = self.compress_ratio <= 1 + N = ( + 0 + if swa_only + else (self.max_model_len + self.compress_ratio - 1) + // self.compress_ratio + ) + M = N + self.window_size + self.max_num_batched_tokens + current_workspace_manager().get_simultaneous( + ((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16), + ) + output.zero_() + return + + assert isinstance(attn_metadata, dict) + flashmla_metadata = cast( + DeepseekV4FlashMLAMetadata | None, attn_metadata.get(self.prefix) + ) + swa_metadata = cast( + "DeepseekSparseSWAMetadata | None", + attn_metadata.get(self.swa_cache_layer.prefix), + ) + assert swa_metadata is not None + + swa_only = self.compress_ratio <= 1 + # SWA-only layers (compress_ratio <= 1) don't have their own KV cache + # allocation, so self.kv_cache may be empty after profiling cleanup. + self_kv_cache = self.kv_cache if not swa_only else None + swa_kv_cache = self.swa_cache_layer.kv_cache + + # Split prefill and decode + num_decodes = swa_metadata.num_decodes + num_prefills = swa_metadata.num_prefills + num_decode_tokens = swa_metadata.num_decode_tokens + + if num_prefills > 0: + self._forward_prefill( + q=q[num_decode_tokens:], + positions=positions[num_decode_tokens:], + compressed_k_cache=self_kv_cache, + swa_k_cache=swa_kv_cache, + output=output[num_decode_tokens:], + attn_metadata=flashmla_metadata, + swa_metadata=swa_metadata, + ) + if num_decodes > 0: + self._forward_decode( + q=q[:num_decode_tokens], + kv_cache=self_kv_cache, + swa_metadata=swa_metadata, + attn_metadata=flashmla_metadata, + swa_only=swa_only, + output=output[:num_decode_tokens], + ) + + def _forward_decode( + self, + q: torch.Tensor, + kv_cache: torch.Tensor | None, # Only used when compress_ratio > 1 + swa_metadata: "DeepseekSparseSWAMetadata", + attn_metadata: DeepseekV4FlashMLAMetadata | None, + swa_only: bool, + output: torch.Tensor, + ) -> None: + num_decodes = swa_metadata.num_decodes + num_decode_tokens = swa_metadata.num_decode_tokens + + topk_indices = None + topk_lens = None + if not swa_only: + assert attn_metadata is not None + assert swa_metadata.is_valid_token is not None + block_size = attn_metadata.block_size // self.compress_ratio + is_valid = swa_metadata.is_valid_token[:num_decode_tokens] + if self.compress_ratio == 4: + # C4A: local indices differ per layer (filled by Indexer). + assert self.topk_indices_buffer is not None + global_indices, topk_lens = compute_global_topk_indices_and_lens( + self.topk_indices_buffer[:num_decode_tokens], + swa_metadata.token_to_req_indices, + attn_metadata.block_table[:num_decodes], + block_size, + is_valid, + ) + topk_indices = global_indices.view(num_decode_tokens, 1, -1) + else: + # C128A: pre-computed during metadata build. + topk_indices = attn_metadata.c128a_global_decode_topk_indices + topk_lens = attn_metadata.c128a_decode_topk_lens + + swa_indices = swa_metadata.decode_swa_indices + swa_lens = swa_metadata.decode_swa_lens + + # We treat queries in the same seq as different queries + # and later we only attend by generated indices. + # q arrives pre-padded to self.padded_heads by the outer wrapper. + q = q.unsqueeze(1) + + # Prepare SWA cache (num_blocks, swa_block_size, 1, head_bytes) + # Use unsqueeze to preserve strides (handles padded blocks correctly) + swa_cache = self.swa_cache_layer.kv_cache.unsqueeze(-2) + # Reshape KV cache to (num_blocks, block_size, 1, head_bytes) + if kv_cache is not None: + kv_cache = kv_cache.unsqueeze(-2) + + # One FlashMLASchedMeta per layer type, shared across all same-type + # layers within this decode step. The first forward call per type + # triggers the in-kernel planner (allocating tile_scheduler_metadata + # and num_splits via PyTorch's graph-aware allocator so CUDA graph + # capture reuses the same addresses on replay); subsequent same-type + # layers see have_initialized=True and skip the planner. + if self.compress_ratio <= 1: + tile_metadata = swa_metadata.tile_sched_swaonly + elif self.compress_ratio == 4: + tile_metadata = swa_metadata.tile_sched_c4a + elif self.compress_ratio == 128: + tile_metadata = swa_metadata.tile_sched_c128a + else: + raise ValueError( + f"Unsupported compress_ratio={self.compress_ratio}; " + "expected 1, 4, or 128." + ) + assert tile_metadata is not None, ( + "swa_metadata missing tile_sched entry for " + f"compress_ratio={self.compress_ratio}; " + "DeepseekSparseSWAMetadataBuilder.build_tile_scheduler did not " + "allocate one for this layer type." + ) + + out, _ = flash_mla_with_kvcache( + q=q, + k_cache=swa_cache, + block_table=None, + head_dim_v=512, + tile_scheduler_metadata=tile_metadata, + cache_seqlens=None, + is_fp8_kvcache=True, + indices=swa_indices, + topk_length=swa_lens, + softmax_scale=self.scale, + attn_sink=self.attn_sink, + extra_k_cache=kv_cache if not swa_only else None, + extra_indices_in_kvcache=topk_indices, + extra_topk_length=topk_lens, + out=output.unsqueeze(1), + ) + + def _forward_prefill( + self, + q: torch.Tensor, + positions: torch.Tensor, + compressed_k_cache: torch.Tensor | None, # Only used when compress_ratio > 1 + swa_k_cache: torch.Tensor, + output: torch.Tensor, + attn_metadata: DeepseekV4FlashMLAMetadata | None, + swa_metadata: "DeepseekSparseSWAMetadata", + ) -> None: + swa_only = attn_metadata is None + + num_prefills = swa_metadata.num_prefills + num_prefill_tokens = swa_metadata.num_prefill_tokens + num_decodes = swa_metadata.num_decodes + num_decode_tokens = swa_metadata.num_decode_tokens + + # Use pre-computed prefill metadata. + seq_lens = swa_metadata.prefill_seq_lens + gather_lens = swa_metadata.prefill_gather_lens + assert seq_lens is not None + assert gather_lens is not None + + # Derive prefill-local token offsets from the full query_start_loc_cpu. + query_start_loc_cpu = swa_metadata.query_start_loc_cpu + query_start_loc = swa_metadata.query_start_loc + assert query_start_loc_cpu is not None + assert query_start_loc is not None + prefill_token_base = query_start_loc_cpu[num_decodes] + + if not swa_only: + if self.compress_ratio == 4: + assert self.topk_indices_buffer is not None + topk_indices = self.topk_indices_buffer[num_decode_tokens:] + topk_indices = topk_indices[:num_prefill_tokens] + else: + # C128A: pre-computed during metadata build. + assert attn_metadata is not None + topk_indices = attn_metadata.c128a_prefill_topk_indices + top_k = topk_indices.shape[-1] + # Compressed region must fit the full compressed pool (seq_len // + # compress_ratio), not just top_k. top_k bounds how many indices + # the indexer selects, not the pool size it indexes into. + N = (self.max_model_len + self.compress_ratio - 1) // self.compress_ratio + else: + # NOTE(woosuk): topk_indices will not be used for SWA-only layers. + assert self.topk_indices_buffer is not None + topk_indices = self.topk_indices_buffer[num_decode_tokens:] + top_k = 0 + N = 0 + + M = N + self.window_size + self.max_num_batched_tokens + chunk_size_const = self.PREFILL_CHUNK_SIZE + num_chunks = (num_prefills + chunk_size_const - 1) // chunk_size_const + + workspace_manager = current_workspace_manager() + kv = workspace_manager.get_simultaneous( + ((chunk_size_const, M, q.shape[-1]), torch.bfloat16), + )[0] + for chunk_idx in range(num_chunks): + chunk_start = chunk_idx * chunk_size_const + chunk_end = min(chunk_start + chunk_size_const, num_prefills) + chunk_size = chunk_end - chunk_start + if not swa_only: + # Gather compressed KV + assert attn_metadata is not None + block_table = attn_metadata.block_table[num_decodes:] + dequantize_and_gather_k_cache( + kv[:chunk_size], + compressed_k_cache, + seq_lens=seq_lens[chunk_start:chunk_end] // self.compress_ratio, + gather_lens=None, + block_table=block_table[chunk_start:chunk_end], + block_size=attn_metadata.block_size // self.compress_ratio, + offset=0, + ) + + # Gather SWA KV + swa_block_table = swa_metadata.block_table[num_decodes:] + dequantize_and_gather_k_cache( + kv[:chunk_size], + swa_k_cache, + seq_lens=seq_lens[chunk_start:chunk_end], + gather_lens=gather_lens[chunk_start:chunk_end], + block_table=swa_block_table[chunk_start:chunk_end], + block_size=swa_metadata.block_size, + offset=N, + ) + + # Combine the topk indices and SWA indices for gathered KV cache + query_start = ( + query_start_loc_cpu[num_decodes + chunk_start] - prefill_token_base + ) + query_end = ( + query_start_loc_cpu[num_decodes + chunk_end] - prefill_token_base + ) + + combined_indices, combined_lens = combine_topk_swa_indices( + topk_indices[query_start:query_end], + query_start_loc[ + num_decodes + chunk_start : num_decodes + chunk_end + 1 + ], + seq_lens[chunk_start:chunk_end], + gather_lens[chunk_start:chunk_end], + self.window_size, + self.compress_ratio, + top_k, + M, + N, + ) + flash_mla_sparse_fwd( + q=q[query_start:query_end], + kv=kv.view(-1, 1, q.shape[-1]), + indices=combined_indices.unsqueeze(1), + sm_scale=self.scale, + attn_sink=self.attn_sink, + topk_length=combined_lens, + out=output[query_start:query_end], + ) diff --git a/upstream/vllm/models/deepseek_v4/ppu/ops/o_proj.py b/upstream/vllm/models/deepseek_v4/ppu/ops/o_proj.py new file mode 100644 index 0000000..d6872f6 --- /dev/null +++ b/upstream/vllm/models/deepseek_v4/ppu/ops/o_proj.py @@ -0,0 +1,244 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import torch +import torch.nn as nn + +from vllm.models.deepseek_v4.common.ops import ( + fused_inv_rope_fp8_quant, + fused_inv_rope_fp8_quant_channelwise, + fused_inv_rope_int8_quant_channelwise, +) +from vllm.platforms import current_platform +from vllm.utils.ppu_deep_gemm import fp8_einsum +from vllm.utils.ppu_deep_gemm import int8_einsum +from vllm.utils.torch_utils import direct_register_custom_op + + +def compute_fp8_einsum_recipe() -> tuple[tuple[int, int, int], bool]: + """fp8_einsum recipe + scale layout for the current GPU arch. + + SM90: FP32 block scales stay [g, r/128, d/128] → sfb_gran_mn=128. + SM100: INT32 packed scales become [g, r, ...] → sfb_gran_mn=1. + + Returns ``(einsum_recipe, tma_aligned_scales)`` for ``deep_gemm_fp8_o_proj``. + """ + cap = current_platform.get_device_capability() + assert cap is not None, "DeepseekV4 attention requires a CUDA device" + einsum_recipe = (1, 128, 128) if cap.major <= 9 else (1, 1, 128) + tma_aligned_scales = cap.major >= 10 + return einsum_recipe, tma_aligned_scales + + +def deep_gemm_fp8_o_proj( + o: torch.Tensor, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, + wo_a: nn.Module, + wo_b: nn.Module, + *, + n_groups: int, + heads_per_group: int, + nope_dim: int, + rope_dim: int, + o_lora_rank: int, + einsum_recipe: tuple[int, int, int], + tma_aligned_scales: bool, +) -> torch.Tensor: + """O projection: inverse RoPE + FP8 quant + einsum + wo_b. + + Shared by the FlashMLA and FlashInfer CUDA backends. ``einsum_recipe`` / + ``tma_aligned_scales`` come from ``compute_fp8_einsum_recipe``. + """ + o_fp8, o_scale = fused_inv_rope_fp8_quant( + o, + positions, + cos_sin_cache, + n_groups=n_groups, + heads_per_group=heads_per_group, + nope_dim=nope_dim, + rope_dim=rope_dim, + tma_aligned_scales=tma_aligned_scales, + ) + z = torch.empty( + (o.shape[0], n_groups, o_lora_rank), + device=o.device, + dtype=torch.bfloat16, + ) + fp8_einsum( + "bhr,hdr->bhd", + (o_fp8, o_scale), + (wo_a.weight, wo_a.weight_scale_inv), + z, + recipe=einsum_recipe, + ) + return wo_b(z.flatten(1)) + + +def deep_gemm_int8_o_proj( + o: torch.Tensor, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, + wo_a: nn.Module, + wo_b: nn.Module, + *, + n_groups: int, + heads_per_group: int, + nope_dim: int, + rope_dim: int, + o_lora_rank: int, + einsum_recipe: tuple[int, int, int], +): + # Channelwise INT8 GEMM via PPU DeepGEMM int8_einsum: + # - activation: per-(token, group) symmetric INT8 quant. + # o_int8 has shape (b=T, h=G, r=R) so o_scale is (b, h, 1). + # - weight: per-output-channel INT8. compressed-tensors + # W8A8 channel scheme stores wo_a as 2D [G*d, R] / scale as + # [G*d, 1]; reshape to 3D so int8_bmm sees (h=G, d=lora_rank, + # r=R) for `hdr` and a (h, d, 1) scale broadcasting along r. + o_int8, o_scale = fused_inv_rope_int8_quant_channelwise( + o, + positions, + cos_sin_cache, + n_groups, + heads_per_group, + nope_dim, + rope_dim, + ) + wo_a_view = wo_a.weight.reshape( + n_groups, o_lora_rank, -1 + ) + wo_a_scale_view = wo_a.weight_scale.reshape( + n_groups, o_lora_rank, 1 + ) + z = torch.empty( + (o.shape[0], n_groups, o_lora_rank), + device=o.device, + dtype=torch.bfloat16, + ) + torch.ops.vllm.deepseek_v4_int8_einsum( + o_int8, + o_scale, + wo_a_view, + wo_a_scale_view, + z, + "bhr,hdr->bhd", + [1, 1, o_int8.shape[-1]], + ) + return wo_b(z.flatten(1)) + + +def deep_gemm_fp8_channel_o_proj( + o: torch.Tensor, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, + wo_a: nn.Module, + wo_b: nn.Module, + *, + n_groups: int, + heads_per_group: int, + nope_dim: int, + rope_dim: int, + o_lora_rank: int, + einsum_recipe: tuple[int, int, int], +): + # Channelwise FP8 GEMM via PPU DeepGEMM fp8_einsum: + # - activation: per-(token, group) symmetric FP8 quant. + # o_fp8 has shape (b=T, h=G, r=R) so o_scale is (b, h, 1). + # - weight: per-output-channel FP8. Channelwise scheme stores + # wo_a as 2D [G*d, R] / scale as [G*d, 1]; reshape to 3D so + # fp8_bmm sees (h=G, d=lora_rank, r=R) for `hdr` and a + # (h, d, 1) scale broadcasting along r. + o_fp8, o_scale = fused_inv_rope_fp8_quant_channelwise( + o, + positions, + cos_sin_cache, + n_groups, + heads_per_group, + nope_dim, + rope_dim, + ) + wo_a_view = wo_a.weight.reshape( + n_groups, o_lora_rank, -1 + ) + wo_a_scale_view = wo_a.weight_scale.reshape( + n_groups, o_lora_rank, 1 + ) + z = torch.empty( + (o.shape[0], n_groups, o_lora_rank), + device=o.device, + dtype=torch.bfloat16, + ) + torch.ops.vllm.deepseek_v4_fp8_channel_einsum( + o_fp8, + o_scale, + wo_a_view, + wo_a_scale_view, + z, + "bhr,hdr->bhd", + [1, 1, o_fp8.shape[-1]], + ) + return wo_b(z.flatten(1)) + + +def deepseek_v4_int8_einsum( + a: torch.Tensor, + a_scale: torch.Tensor, + b: torch.Tensor, + b_scale: torch.Tensor, + out: torch.Tensor, + equation: str, + recipe: list[int], +) -> None: + int8_einsum(equation, (a, a_scale), (b, b_scale), out, recipe=tuple(recipe)) + + +def deepseek_v4_int8_einsum_fake( + a: torch.Tensor, + a_scale: torch.Tensor, + b: torch.Tensor, + b_scale: torch.Tensor, + out: torch.Tensor, + equation: str, + recipe: list[int], +) -> None: + return None + + +direct_register_custom_op( + op_name="deepseek_v4_int8_einsum", + op_func=deepseek_v4_int8_einsum, + mutates_args=["out"], + fake_impl=deepseek_v4_int8_einsum_fake, +) + + +def deepseek_v4_fp8_channel_einsum( + a: torch.Tensor, + a_scale: torch.Tensor, + b: torch.Tensor, + b_scale: torch.Tensor, + out: torch.Tensor, + equation: str, + recipe: list[int], +) -> None: + fp8_einsum(equation, (a, a_scale), (b, b_scale), out, recipe=tuple(recipe)) + + +def deepseek_v4_fp8_channel_einsum_fake( + a: torch.Tensor, + a_scale: torch.Tensor, + b: torch.Tensor, + b_scale: torch.Tensor, + out: torch.Tensor, + equation: str, + recipe: list[int], +) -> None: + return None + + +direct_register_custom_op( + op_name="deepseek_v4_fp8_channel_einsum", + op_func=deepseek_v4_fp8_channel_einsum, + mutates_args=["out"], + fake_impl=deepseek_v4_fp8_channel_einsum_fake, +) \ No newline at end of file