first commit

.dockerignore

@@ -0,0 +1,6 @@
+.git
+__pycache__
+*.pyc
+upstream_archive/
+patches/
+README.md

.gitignore

@@ -0,0 +1,3 @@
+__pycache__/
+*.pyc
+upstream_archive/

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

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/`。

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"
+  }
+}

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)

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"
+    }
+  }
+}

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())

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)

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)

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)

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],
+            )

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())

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()

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)

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)

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],
+            )

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,
+)