| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """DSV4.1 output projection with a small-batch SM100/SM103 fusion.""" |
| 4 | |
| 5 | import torch |
| 6 | from torch import nn |
| 7 | |
| 8 | from vllm.model_executor.kernels.linear.mxfp8.flashinfer import ( |
| 9 | FlashInferCutedslMxfp8LinearKernel, |
| 10 | ) |
| 11 | from vllm.model_executor.layers.fusion.quant_activation import ( |
| 12 | QuantizedActivation, |
| 13 | get_input_quant_key, |
| 14 | ) |
| 15 | from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Dynamic |
| 16 | from vllm.models.deepseek_v4.nvidia.ops.o_proj import deep_gemm_fp8_o_proj |
| 17 | from vllm.platforms import current_platform |
| 18 | |
| 19 | |
| 20 | def _fused_wo_a_max_tokens(layer: nn.Module) -> int: |
| 21 | max_tokens = getattr(layer, "_fused_wo_a_token_limit", None) |
| 22 | if max_tokens is None: |
| 23 | # Additional token tiles only pay off for smaller grids. |
| 24 | ctas_per_tile = layer.n_local_heads * (layer.o_lora_rank // 128) |
| 25 | max_tokens = 96 if ctas_per_tile <= 128 else 32 |
| 26 | layer._fused_wo_a_token_limit = max_tokens |
| 27 | return max_tokens |
| 28 | |
| 29 | |
| 30 | def _can_fuse_wo_a(layer: nn.Module) -> bool: |
| 31 | """Whether the attention layer matches the fused WO-A kernel's layout.""" |
| 32 | method = getattr(layer.wo_b, "scheme", layer.wo_b.quant_method) |
| 33 | return ( |
| 34 | current_platform.is_device_capability_family(100) |
| 35 | # WO-A is FP8 with per-32 scales, and WO-B takes MXFP8 input directly. |
| 36 | and layer._einsum_recipe == (1, 1, 32) |
| 37 | and get_input_quant_key(layer.wo_b) == kMxfp8Dynamic |
| 38 | and isinstance( |
| 39 | getattr(method, "kernel", None), FlashInferCutedslMxfp8LinearKernel |
| 40 | ) |
| 41 | # A group's heads form one portable (<= 8 CTA) cluster. |
| 42 | and layer.n_local_heads // layer.n_local_groups <= 8 |
| 43 | and (layer.nope_head_dim, layer.rope_head_dim) == (448, 64) |
| 44 | and layer.o_lora_rank % 128 == 0 |
| 45 | ) |
| 46 | |
| 47 | |
| 48 | def register_dsv41_o_proj_warmup(layer: nn.Module) -> None: |
| 49 | """Warm the fused WO-A token counts ``dsv41_o_proj`` can dispatch.""" |
| 50 | # Draft models load outside the warmup registry; their layers share the |
| 51 | # target's attention shapes, so the target's registration covers them too. |
| 52 | if _can_fuse_wo_a(layer): |
| 53 | from .fused_wo_a import _FUSED_WO_A_KERNEL |
| 54 | |
| 55 | _FUSED_WO_A_KERNEL.register_warmup( |
| 56 | max_tokens=_fused_wo_a_max_tokens(layer), |
| 57 | n_groups=layer.n_local_groups, |
| 58 | heads_per_group=layer.n_local_heads // layer.n_local_groups, |
| 59 | o_lora_rank=layer.o_lora_rank, |
| 60 | ) |
| 61 | |
| 62 | |
| 63 | def dsv41_o_proj( |
| 64 | layer: nn.Module, attn_out: torch.Tensor, positions: torch.Tensor |
| 65 | ) -> torch.Tensor: |
| 66 | """``deep_gemm_fp8_o_proj`` with WO-A fused for small SM100/SM103 batches.""" |
| 67 | o = attn_out[:, : layer.n_local_heads, :] |
| 68 | cos_sin_cache = layer.rotary_emb.cos_sin_cache |
| 69 | if 1 <= o.shape[0] <= _fused_wo_a_max_tokens(layer) and _can_fuse_wo_a(layer): |
| 70 | from .fused_wo_a import _FUSED_WO_A_KERNEL |
| 71 | |
| 72 | q, scales = _FUSED_WO_A_KERNEL( |
| 73 | x=o, |
| 74 | positions=positions, |
| 75 | rope=cos_sin_cache, |
| 76 | weight=layer.wo_a.weight, |
| 77 | weight_scale=layer.wo_a.weight_scale, |
| 78 | ) |
| 79 | return layer._wo_b_proj( |
| 80 | QuantizedActivation(q, scales, o.dtype, q.shape, kMxfp8Dynamic) |
| 81 | ) |
| 82 | return deep_gemm_fp8_o_proj( |
| 83 | o, |
| 84 | positions, |
| 85 | cos_sin_cache, |
| 86 | layer.wo_a, |
| 87 | layer._wo_b_proj, |
| 88 | n_groups=layer.n_local_groups, |
| 89 | heads_per_group=layer.n_local_heads // layer.n_local_groups, |
| 90 | nope_dim=layer.nope_head_dim, |
| 91 | rope_dim=layer.rope_head_dim, |
| 92 | o_lora_rank=layer.o_lora_rank, |
| 93 | einsum_recipe=layer._einsum_recipe, |
| 94 | tma_aligned_scales=layer._tma_aligned_scales, |
| 95 | ) |