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