first commit

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