ppu_deepseek_v4_pro_patch/patches/vllm.patch

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