@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 158 lines 5.72 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from collections.abc import Callable
4
5 import torch
6 import torch.nn as nn
7
8 from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation
9 from vllm.model_executor.layers.quantization.utils.quant_utils import (
10 kFp8Dynamic128Sym,
11 )
12 from vllm.models.deepseek_v4.common.ops.fused_inv_rope_fp8_quant import (
13 fused_inv_rope_fp8_quant,
14 )
15 from vllm.platforms import current_platform
16 from vllm.utils.deep_gemm import fp8_einsum, get_tma_aligned_size
17
18
19 def compute_fp8_einsum_recipe(
20 block_size: int = 128,
21 ) -> tuple[tuple[int, int, int], bool]:
22 """fp8_einsum recipe + scale layout for the current GPU arch.
23
24 SM90 keeps block-row FP32 scales. SM100 uses packed per-row E8M0 scales.
25
26 Returns ``(einsum_recipe, tma_aligned_scales)`` for ``deep_gemm_fp8_o_proj``.
27 """
28 cap = current_platform.get_device_capability()
29 assert cap is not None, "DeepseekV4 attention requires a CUDA device"
30 einsum_recipe = (1, 128, 128) if cap.major <= 9 else (1, 1, block_size)
31 tma_aligned_scales = cap.major >= 10
32 return einsum_recipe, tma_aligned_scales
33
34
35 def deep_gemm_fp8_o_proj(
36 o: torch.Tensor | QuantizedActivation,
37 positions: torch.Tensor,
38 cos_sin_cache: torch.Tensor,
39 wo_a: nn.Module,
40 wo_b: Callable[[torch.Tensor], torch.Tensor],
41 *,
42 n_groups: int,
43 heads_per_group: int,
44 nope_dim: int,
45 rope_dim: int,
46 o_lora_rank: int,
47 einsum_recipe: tuple[int, int, int],
48 tma_aligned_scales: bool,
49 ) -> torch.Tensor:
50 """O projection: inverse RoPE + grouped wo_a + wo_b.
51
52 Shared by the FlashMLA and FlashInfer CUDA backends. The attention
53 layer selects the recipe at initialization. ``wo_b`` is any callable over
54 the flattened ``z``: the projection module itself, or a wrapper that also
55 reduce-scatters its output (DeepSeek-V4.1 GEMM-RS).
56
57 A QuantizedActivation ``o`` was already rotated and cast by the attention
58 kernel, one ``wo_a`` group per slot with the live groups first, so only
59 the padding groups are sliced off.
60 """
61 use_fp8 = wo_a.weight.dtype == torch.float8_e4m3fn
62 if isinstance(o, QuantizedActivation):
63 assert use_fp8
64 o_proj_input, o_scale = o.data[:, :n_groups], o.scale[:, :n_groups]
65 else:
66 o_proj_input, o_scale = fused_inv_rope_fp8_quant(
67 o,
68 positions,
69 cos_sin_cache,
70 n_groups=n_groups,
71 heads_per_group=heads_per_group,
72 nope_dim=nope_dim,
73 rope_dim=rope_dim,
74 quant_group_size=einsum_recipe[2],
75 tma_aligned_scales=tma_aligned_scales,
76 quantize=use_fp8,
77 )
78 z = torch.empty(
79 (o_proj_input.shape[0], n_groups, o_lora_rank),
80 device=o_proj_input.device,
81 dtype=torch.bfloat16,
82 )
83 if use_fp8:
84 weight_scale = (
85 wo_a.weight_scale
86 if hasattr(wo_a, "weight_scale")
87 else wo_a.weight_scale_inv
88 )
89 fp8_einsum(
90 "bhr,hdr->bhd",
91 (o_proj_input, o_scale),
92 (wo_a.weight, weight_scale),
93 z,
94 recipe=einsum_recipe,
95 )
96 else:
97 grouped_weight = wo_a.weight.view(n_groups, o_lora_rank, -1)
98 torch.bmm(
99 o_proj_input.transpose(0, 1),
100 grouped_weight.transpose(1, 2),
101 out=z.transpose(0, 1),
102 )
103 return wo_b(z.flatten(1))
104
105
106 # TRTLLM-gen's DSv4 RopeQuant epilogue (flashinfer-ai/flashinfer#4918) has one
107 # fixed schedule: 128 query heads in groups of 8 512-wide heads, and FP8 with
108 # one packed UE8M0 scale per 128 elements.
109 _ROPE_QUANT_NUM_HEADS = 128
110 _ROPE_QUANT_HEADS_PER_GROUP = 8
111 _ROPE_QUANT_HEAD_DIM = 512
112 _ROPE_QUANT_BLOCK = 128
113
114
115 def rope_quant_unsupported_reason(layer: torch.nn.Module) -> str | None:
116 """Why ``layer`` cannot take FlashInfer's fused inverse RoPE + FP8 output.
117
118 DSv4 always has 8 heads per ``wo_a`` group at any TP size, and the rest of
119 the kernel's fixed shape (head/RoPE dims, per-128 einsum scales, FP32 RoPE
120 cache) is what the SM100 layer uses anyway. What varies is the local head
121 count and the configs the layer also serves unquantized.
122 """
123 if layer.n_local_heads != _ROPE_QUANT_NUM_HEADS:
124 # Padding fewer heads up to the kernel's fixed 128 costs more
125 # attention than the fused epilogue saves.
126 return f"RopeQuant needs {_ROPE_QUANT_NUM_HEADS} local query heads"
127 if layer.kv_cache_torch_dtype != torch.float8_e4m3fn:
128 return "RopeQuant needs the per-tensor FP8 KV cache"
129 if layer.wo_a.weight.dtype != torch.float8_e4m3fn:
130 return "RopeQuant needs an FP8 wo_a"
131 return None
132
133
134 def rope_quant_attn_out(num_tokens: int, device: torch.device) -> QuantizedActivation:
135 """Allocate the output pair the RopeQuant kernel writes.
136
137 Values are group-major, ``[num_tokens, 16, 8 * 512]`` over physical
138 ``[group, token, K]``. Scales are MN-major packed UE8M0, one int32 per
139 head over its four 128-wide blocks. Both are what ``fp8_einsum`` consumes.
140 """
141 n_groups = _ROPE_QUANT_NUM_HEADS // _ROPE_QUANT_HEADS_PER_GROUP
142 group_width = _ROPE_QUANT_HEADS_PER_GROUP * _ROPE_QUANT_HEAD_DIM
143 data = torch.empty(
144 (n_groups, num_tokens, group_width), dtype=torch.float8_e4m3fn, device=device
145 ).transpose(0, 1)
146 aligned = get_tma_aligned_size(num_tokens, torch.int32.itemsize)
147 scale = torch.empty(
148 (n_groups, group_width // (_ROPE_QUANT_BLOCK * 4), aligned),
149 dtype=torch.int32,
150 device=device,
151 ).permute(2, 0, 1)[:num_tokens]
152 return QuantizedActivation(
153 data=data,
154 scale=scale,
155 orig_dtype=torch.bfloat16,
156 orig_shape=data.shape,
157 quant_key=kFp8Dynamic128Sym,
158 )