@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 4fb56914c5

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 01:00 UTC

[perf] Add fused MLA QKV + strided layernorm (#21116)

Signed-off-by: Mickael Seznec <mickael@mistral.ai> Co-authored-by: mgoin <mgoin64@gmail.com>

Mickaël Seznec committed Jul 22, 2025 at 16:07 UTC 4fb56914c5f27ef062e10d44a0f79c6ceab382f9

7 files changed +214 -66

csrc/layernorm_kernels.cu

+40 -23

@@ -15,15 +15,16 @@ namespace vllm {

15 // TODO(woosuk): Further optimize this kernel.

16 template <typename scalar_t>

17 __global__ void rms_norm_kernel(

18 - scalar_t* __restrict__ out, // [..., hidden_size]

19 - const scalar_t* __restrict__ input, // [..., hidden_size]

18 + scalar_t* __restrict__ out, // [..., hidden_size]

19 + const scalar_t* __restrict__ input, // [..., hidden_size]

20 + const int64_t input_stride,

21 const scalar_t* __restrict__ weight, // [hidden_size]

22 const float epsilon, const int num_tokens, const int hidden_size) {

23 __shared__ float s_variance;

24 float variance = 0.0f;

26 for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {

26 - const float x = (float)input[blockIdx.x * hidden_size + idx];

27 + const float x = (float)input[blockIdx.x * input_stride + idx];

28 variance += x * x;

29 }

@@ -37,7 +38,7 @@ __global__ void rms_norm_kernel(

38 __syncthreads();

40 for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {

40 - float x = (float)input[blockIdx.x * hidden_size + idx];

41 + float x = (float)input[blockIdx.x * input_stride + idx];

42 out[blockIdx.x * hidden_size + idx] =

43 ((scalar_t)(x * s_variance)) * weight[idx];

44 }

@@ -50,7 +51,8 @@ __global__ void rms_norm_kernel(

51 template <typename scalar_t, int width>

52 __global__ std::enable_if_t<(width > 0) && _typeConvert<scalar_t>::exists>

53 fused_add_rms_norm_kernel(

53 - scalar_t* __restrict__ input, // [..., hidden_size]

54 + scalar_t* __restrict__ input, // [..., hidden_size]

55 + const int64_t input_stride,

56 scalar_t* __restrict__ residual, // [..., hidden_size]

57 const scalar_t* __restrict__ weight, // [hidden_size]

58 const float epsilon, const int num_tokens, const int hidden_size) {

@@ -59,6 +61,7 @@ fused_add_rms_norm_kernel(

61 static_assert(sizeof(_f16Vec<scalar_t, width>) == sizeof(scalar_t) * width);

63 const int vec_hidden_size = hidden_size / width;

64 + const int64_t vec_input_stride = input_stride / width;

65 __shared__ float s_variance;

66 float variance = 0.0f;

67 /* These and the argument pointers are all declared `restrict` as they are

@@ -73,7 +76,8 @@ fused_add_rms_norm_kernel(

77 for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) {

78 int id = blockIdx.x * vec_hidden_size + idx;

76 - _f16Vec<scalar_t, width> temp = input_v[id];

79 + int64_t strided_id = blockIdx.x * vec_input_stride + idx;

80 + _f16Vec<scalar_t, width> temp = input_v[strided_id];

81 temp += residual_v[id];

82 variance += temp.sum_squares();

83 residual_v[id] = temp;

@@ -90,10 +94,11 @@ fused_add_rms_norm_kernel(

95 for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) {

96 int id = blockIdx.x * vec_hidden_size + idx;

97 + int64_t strided_id = blockIdx.x * vec_input_stride + idx;

98 _f16Vec<scalar_t, width> temp = residual_v[id];

99 temp *= s_variance;

100 temp *= weight_v[idx];

96 - input_v[id] = temp;

101 + input_v[strided_id] = temp;

102 }

103 }

104

@@ -103,7 +108,8 @@ fused_add_rms_norm_kernel(

108 template <typename scalar_t, int width>

109 __global__ std::enable_if_t<(width == 0) || !_typeConvert<scalar_t>::exists>

110 fused_add_rms_norm_kernel(

106 - scalar_t* __restrict__ input, // [..., hidden_size]

111 + scalar_t* __restrict__ input, // [..., hidden_size]

112 + const int64_t input_stride,

113 scalar_t* __restrict__ residual, // [..., hidden_size]

114 const scalar_t* __restrict__ weight, // [hidden_size]

115 const float epsilon, const int num_tokens, const int hidden_size) {

@@ -111,7 +117,7 @@ fused_add_rms_norm_kernel(

117 float variance = 0.0f;

118

119 for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {

114 - scalar_t z = input[blockIdx.x * hidden_size + idx];

120 + scalar_t z = input[blockIdx.x * input_stride + idx];

121 z += residual[blockIdx.x * hidden_size + idx];

122 float x = (float)z;

123 variance += x * x;

@@ -129,7 +135,7 @@ fused_add_rms_norm_kernel(

135

136 for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {

137 float x = (float)residual[blockIdx.x * hidden_size + idx];

132 - input[blockIdx.x * hidden_size + idx] =

138 + input[blockIdx.x * input_stride + idx] =

139 ((scalar_t)(x * s_variance)) * weight[idx];

140 }

141 }

@@ -141,11 +147,12 @@ void rms_norm(torch::Tensor& out, // [..., hidden_size]

147 torch::Tensor& weight, // [hidden_size]

148 double epsilon) {

149 TORCH_CHECK(out.is_contiguous());

144 - TORCH_CHECK(input.is_contiguous());

150 + TORCH_CHECK(input.stride(-1) == 1);

151 TORCH_CHECK(weight.is_contiguous());

152

153 int hidden_size = input.size(-1);

154 int num_tokens = input.numel() / hidden_size;

155 + int64_t input_stride = input.stride(-2);

156

157 dim3 grid(num_tokens);

158 dim3 block(std::min(hidden_size, 1024));

@@ -153,26 +160,29 @@ void rms_norm(torch::Tensor& out, // [..., hidden_size]

160 const cudaStream_t stream = at::cuda::getCurrentCUDAStream();

161 VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "rms_norm_kernel", [&] {

162 vllm::rms_norm_kernel<scalar_t><<<grid, block, 0, stream>>>(

156 - out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(),

163 + out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), input_stride,

164 weight.data_ptr<scalar_t>(), epsilon, num_tokens, hidden_size);

165 });

166 }

167

161 -#define LAUNCH_FUSED_ADD_RMS_NORM(width) \

162 - VLLM_DISPATCH_FLOATING_TYPES( \

163 - input.scalar_type(), "fused_add_rms_norm_kernel", [&] { \

164 - vllm::fused_add_rms_norm_kernel<scalar_t, width> \

165 - <<<grid, block, 0, stream>>>(input.data_ptr<scalar_t>(), \

166 - residual.data_ptr<scalar_t>(), \

167 - weight.data_ptr<scalar_t>(), epsilon, \

168 - num_tokens, hidden_size); \

168 +#define LAUNCH_FUSED_ADD_RMS_NORM(width) \

169 + VLLM_DISPATCH_FLOATING_TYPES( \

170 + input.scalar_type(), "fused_add_rms_norm_kernel", [&] { \

171 + vllm::fused_add_rms_norm_kernel<scalar_t, width> \

172 + <<<grid, block, 0, stream>>>( \

173 + input.data_ptr<scalar_t>(), input_stride, \

174 + residual.data_ptr<scalar_t>(), weight.data_ptr<scalar_t>(), \

175 + epsilon, num_tokens, hidden_size); \

176 });

177

178 void fused_add_rms_norm(torch::Tensor& input, // [..., hidden_size]

179 torch::Tensor& residual, // [..., hidden_size]

180 torch::Tensor& weight, // [hidden_size]

181 double epsilon) {

182 + TORCH_CHECK(residual.is_contiguous());

183 + TORCH_CHECK(weight.is_contiguous());

184 int hidden_size = input.size(-1);

185 + int64_t input_stride = input.stride(-2);

186 int num_tokens = input.numel() / hidden_size;

187

188 dim3 grid(num_tokens);

@@ -194,9 +204,16 @@ void fused_add_rms_norm(torch::Tensor& input, // [..., hidden_size]

204 auto inp_ptr = reinterpret_cast<std::uintptr_t>(input.data_ptr());

205 auto res_ptr = reinterpret_cast<std::uintptr_t>(residual.data_ptr());

206 auto wt_ptr = reinterpret_cast<std::uintptr_t>(weight.data_ptr());

197 - bool ptrs_are_aligned =

198 - inp_ptr % 16 == 0 && res_ptr % 16 == 0 && wt_ptr % 16 == 0;

199 - if (ptrs_are_aligned && hidden_size % 8 == 0) {

207 + constexpr int vector_width = 8;

208 + constexpr int req_alignment_bytes =

209 + vector_width * 2; // vector_width * sizeof(bfloat16 or float16) (float32

210 + // falls back to non-vectorized version anyway)

211 + bool ptrs_are_aligned = inp_ptr % req_alignment_bytes == 0 &&

212 + res_ptr % req_alignment_bytes == 0 &&

213 + wt_ptr % req_alignment_bytes == 0;

214 + bool offsets_are_multiple_of_vector_width =

215 + hidden_size % vector_width == 0 && input_stride % vector_width == 0;

216 + if (ptrs_are_aligned && offsets_are_multiple_of_vector_width) {

217 LAUNCH_FUSED_ADD_RMS_NORM(8);

218 } else {

219 LAUNCH_FUSED_ADD_RMS_NORM(0);

csrc/layernorm_quant_kernels.cu

+25 -14

@@ -23,8 +23,9 @@ namespace vllm {

23 // TODO(woosuk): Further optimize this kernel.

24 template <typename scalar_t, typename fp8_type>

25 __global__ void rms_norm_static_fp8_quant_kernel(

26 - fp8_type* __restrict__ out, // [..., hidden_size]

27 - const scalar_t* __restrict__ input, // [..., hidden_size]

26 + fp8_type* __restrict__ out, // [..., hidden_size]

27 + const scalar_t* __restrict__ input, // [..., hidden_size]

28 + const int input_stride,

29 const scalar_t* __restrict__ weight, // [hidden_size]

30 const float* __restrict__ scale, // [1]

31 const float epsilon, const int num_tokens, const int hidden_size) {

@@ -32,7 +33,7 @@ __global__ void rms_norm_static_fp8_quant_kernel(

33 float variance = 0.0f;

35 for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {

35 - const float x = (float)input[blockIdx.x * hidden_size + idx];

36 + const float x = (float)input[blockIdx.x * input_stride + idx];

37 variance += x * x;

38 }

@@ -49,7 +50,7 @@ __global__ void rms_norm_static_fp8_quant_kernel(

50 float const scale_inv = 1.0f / *scale;

52 for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {

52 - float x = (float)input[blockIdx.x * hidden_size + idx];

53 + float x = (float)input[blockIdx.x * input_stride + idx];

54 float const out_norm = ((scalar_t)(x * s_variance)) * weight[idx];

55 out[blockIdx.x * hidden_size + idx] =

56 scaled_fp8_conversion<true, fp8_type>(out_norm, scale_inv);

@@ -63,8 +64,9 @@ __global__ void rms_norm_static_fp8_quant_kernel(

64 template <typename scalar_t, int width, typename fp8_type>

65 __global__ std::enable_if_t<(width > 0) && _typeConvert<scalar_t>::exists>

66 fused_add_rms_norm_static_fp8_quant_kernel(

66 - fp8_type* __restrict__ out, // [..., hidden_size]

67 - scalar_t* __restrict__ input, // [..., hidden_size]

67 + fp8_type* __restrict__ out, // [..., hidden_size]

68 + scalar_t* __restrict__ input, // [..., hidden_size]

69 + const int input_stride,

70 scalar_t* __restrict__ residual, // [..., hidden_size]

71 const scalar_t* __restrict__ weight, // [hidden_size]

72 const float* __restrict__ scale, // [1]

@@ -74,6 +76,7 @@ fused_add_rms_norm_static_fp8_quant_kernel(

76 static_assert(sizeof(_f16Vec<scalar_t, width>) == sizeof(scalar_t) * width);

78 const int vec_hidden_size = hidden_size / width;

79 + const int vec_input_stride = input_stride / width;

80 __shared__ float s_variance;

81 float variance = 0.0f;

82 /* These and the argument pointers are all declared `restrict` as they are

@@ -87,8 +90,9 @@ fused_add_rms_norm_static_fp8_quant_kernel(

90 reinterpret_cast<const _f16Vec<scalar_t, width>*>(weight);

92 for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) {

93 + int stride_id = blockIdx.x * vec_input_stride + idx;

94 int id = blockIdx.x * vec_hidden_size + idx;

91 - _f16Vec<scalar_t, width> temp = input_v[id];

95 + _f16Vec<scalar_t, width> temp = input_v[stride_id];

96 temp += residual_v[id];

97 variance += temp.sum_squares();

98 residual_v[id] = temp;

@@ -125,8 +129,9 @@ fused_add_rms_norm_static_fp8_quant_kernel(

129 template <typename scalar_t, int width, typename fp8_type>

130 __global__ std::enable_if_t<(width == 0) || !_typeConvert<scalar_t>::exists>

131 fused_add_rms_norm_static_fp8_quant_kernel(

128 - fp8_type* __restrict__ out, // [..., hidden_size]

129 - scalar_t* __restrict__ input, // [..., hidden_size]

132 + fp8_type* __restrict__ out, // [..., hidden_size]

133 + scalar_t* __restrict__ input, // [..., hidden_size]

134 + const int input_stride,

135 scalar_t* __restrict__ residual, // [..., hidden_size]

136 const scalar_t* __restrict__ weight, // [hidden_size]

137 const float* __restrict__ scale, // [1]

@@ -135,7 +140,7 @@ fused_add_rms_norm_static_fp8_quant_kernel(

140 float variance = 0.0f;

141

142 for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {

138 - scalar_t z = input[blockIdx.x * hidden_size + idx];

143 + scalar_t z = input[blockIdx.x * input_stride + idx];

144 z += residual[blockIdx.x * hidden_size + idx];

145 float x = (float)z;

146 variance += x * x;

@@ -169,7 +174,9 @@ void rms_norm_static_fp8_quant(torch::Tensor& out, // [..., hidden_size]

174 torch::Tensor& weight, // [hidden_size]

175 torch::Tensor& scale, // [1]

176 double epsilon) {

177 + TORCH_CHECK(out.is_contiguous());

178 int hidden_size = input.size(-1);

179 + int input_stride = input.stride(-2);

180 int num_tokens = input.numel() / hidden_size;

181

182 dim3 grid(num_tokens);

@@ -183,8 +190,9 @@ void rms_norm_static_fp8_quant(torch::Tensor& out, // [..., hidden_size]

190 vllm::rms_norm_static_fp8_quant_kernel<scalar_t, fp8_t>

191 <<<grid, block, 0, stream>>>(

192 out.data_ptr<fp8_t>(), input.data_ptr<scalar_t>(),

186 - weight.data_ptr<scalar_t>(), scale.data_ptr<float>(),

187 - epsilon, num_tokens, hidden_size);

193 + input_stride, weight.data_ptr<scalar_t>(),

194 + scale.data_ptr<float>(), epsilon, num_tokens,

195 + hidden_size);

196 });

197 });

198 }

@@ -198,7 +206,7 @@ void rms_norm_static_fp8_quant(torch::Tensor& out, // [..., hidden_size]

206 width, fp8_t> \

207 <<<grid, block, 0, stream>>>( \

208 out.data_ptr<fp8_t>(), input.data_ptr<scalar_t>(), \

201 - residual.data_ptr<scalar_t>(), \

209 + input_stride, residual.data_ptr<scalar_t>(), \

210 weight.data_ptr<scalar_t>(), scale.data_ptr<float>(), \

211 epsilon, num_tokens, hidden_size); \

212 }); \

@@ -210,7 +218,10 @@ void fused_add_rms_norm_static_fp8_quant(

218 torch::Tensor& weight, // [hidden_size]

219 torch::Tensor& scale, // [1]

220 double epsilon) {

221 + TORCH_CHECK(out.is_contiguous());

222 + TORCH_CHECK(residual.is_contiguous());

223 int hidden_size = input.size(-1);

224 + int input_stride = input.stride(-2);

225 int num_tokens = input.numel() / hidden_size;

226

227 dim3 grid(num_tokens);

@@ -234,7 +245,7 @@ void fused_add_rms_norm_static_fp8_quant(

245 auto wt_ptr = reinterpret_cast<std::uintptr_t>(weight.data_ptr());

246 bool ptrs_are_aligned =

247 inp_ptr % 16 == 0 && res_ptr % 16 == 0 && wt_ptr % 16 == 0;

237 - if (ptrs_are_aligned && hidden_size % 8 == 0) {

248 + if (ptrs_are_aligned && hidden_size % 8 == 0 && input_stride % 8 == 0) {

249 LAUNCH_FUSED_ADD_RMS_NORM(8);

250 } else {

251 LAUNCH_FUSED_ADD_RMS_NORM(0);

csrc/quantization/fp8/common.cu

@@ -88,6 +88,8 @@ void static_scaled_fp8_quant(torch::Tensor& out, // [..., d]

88 torch::Tensor const& input, // [..., d]

89 torch::Tensor const& scale) // [1]

90 {

91 + TORCH_CHECK(input.is_contiguous());

92 + TORCH_CHECK(out.is_contiguous());

93 int const block_size = 256;

94 int const num_tokens = input.numel() / input.size(-1);

95 int const num_elems = input.numel();

@@ -111,6 +113,8 @@ void dynamic_scaled_fp8_quant(torch::Tensor& out, // [..., d]

113 torch::Tensor const& input, // [..., d]

114 torch::Tensor& scale) // [1]

115 {

116 + TORCH_CHECK(input.is_contiguous());

117 + TORCH_CHECK(out.is_contiguous());

118 int const block_size = 256;

119 int const num_tokens = input.numel() / input.size(-1);

120 int const num_elems = input.numel();

tests/kernels/core/test_layernorm.py

+19 -7

@@ -26,6 +26,7 @@ CUDA_DEVICES = [

26 @pytest.mark.parametrize("dtype", DTYPES)

27 @pytest.mark.parametrize("seed", SEEDS)

28 @pytest.mark.parametrize("device", CUDA_DEVICES)

29 +@pytest.mark.parametrize("strided_input", [False, True])

30 @torch.inference_mode()

31 def test_rms_norm(

32 num_tokens: int,

@@ -34,13 +35,17 @@ def test_rms_norm(

35 dtype: torch.dtype,

36 seed: int,

37 device: str,

38 + strided_input: bool,

39 ) -> None:

40 current_platform.seed_everything(seed)

41 torch.set_default_device(device)

42 layer = RMSNorm(hidden_size).to(dtype=dtype)

43 layer.weight.data.normal_(mean=1.0, std=0.1)

44 scale = 1 / (2 * hidden_size)

43 - x = torch.randn(num_tokens, hidden_size, dtype=dtype)

45 + last_dim = 2 * hidden_size if strided_input else hidden_size

46 + x = torch.randn(num_tokens, last_dim, dtype=dtype)

47 + x = x[..., :hidden_size]

48 + assert x.is_contiguous() != strided_input

49 x *= scale

50 residual = torch.randn_like(x) * scale if add_residual else None

@@ -72,6 +77,7 @@ def test_rms_norm(

77 @pytest.mark.parametrize("quant_scale", [1.0, 0.01, 10.0])

78 @pytest.mark.parametrize("seed", SEEDS)

79 @pytest.mark.parametrize("device", CUDA_DEVICES)

80 +@pytest.mark.parametrize("strided_input", [False, True])

81 def test_fused_rms_norm_quant(

82 num_tokens: int,

83 hidden_size: int,

@@ -80,13 +86,18 @@ def test_fused_rms_norm_quant(

86 quant_scale: float,

87 seed: int,

88 device: str,

89 + strided_input: bool,

90 ) -> None:

91 current_platform.seed_everything(seed)

92 torch.set_default_device(device)

94 weight = torch.empty(hidden_size, dtype=dtype).normal_(mean=1.0, std=0.1)

95 scale = 1 / (2 * hidden_size)

89 - x = torch.randn(num_tokens, hidden_size, dtype=dtype)

96 + last_dim = 2 * hidden_size if strided_input else hidden_size

97 + x_base = torch.randn(num_tokens, last_dim, dtype=dtype)

98 + x = x_base[..., :hidden_size]

99 + assert x.is_contiguous() != strided_input

100 +

101 x *= scale

102 if add_residual:

103 residual = torch.randn_like(x) * scale

@@ -106,9 +117,11 @@ def test_fused_rms_norm_quant(

117

118 # Unfused kernel is in-place so it goes second

119 # Also use a separate clone of x to avoid modifying the input

109 - x_unfused = x.clone()

120 + x_unfused_base = x_base.clone()

121 + x_unfused = x_unfused_base[..., :hidden_size]

122 + assert x_unfused.is_contiguous() != strided_input

123 torch.ops._C.fused_add_rms_norm(x_unfused, residual, weight, 1e-6)

111 - torch.ops._C.static_scaled_fp8_quant(out_quant, x_unfused,

124 + torch.ops._C.static_scaled_fp8_quant(out_quant, x_unfused.contiguous(),

125 quant_scale_t)

126

127 torch.cuda.synchronize()

@@ -116,7 +129,6 @@ def test_fused_rms_norm_quant(

129 residual,

130 atol=1e-2,

131 rtol=1e-2)

119 -

132 opcheck(

133 torch.ops._C.fused_add_rms_norm_static_fp8_quant,

134 (out_quant_fused, x, residual_fused, weight, quant_scale_t, 1e-6))

@@ -131,7 +143,7 @@ def test_fused_rms_norm_quant(

143 opcheck(torch.ops._C.rms_norm_static_fp8_quant,

144 (out_quant_fused, x, weight, quant_scale_t, 1e-6))

145

134 - torch.testing.assert_close(out_quant_fused.to(dtype=torch.float32),

135 - out_quant.to(dtype=torch.float32),

146 + torch.testing.assert_close(out_quant.to(dtype=torch.float32),

147 + out_quant_fused.to(dtype=torch.float32),

148 atol=1e-3,

149 rtol=1e-3)

vllm/model_executor/layers/linear.py

+77 -1

@@ -259,6 +259,8 @@ class LinearBase(torch.nn.Module):

259 if params_dtype is None:

260 params_dtype = torch.get_default_dtype()

261 self.params_dtype = params_dtype

262 + self.quant_config = quant_config

263 + self.prefix = prefix

264 if quant_config is None:

265 self.quant_method: Optional[

266 QuantizeMethodBase] = UnquantizedLinearMethod()

@@ -300,6 +302,12 @@ class ReplicatedLinear(LinearBase):

302 *,

303 return_bias: bool = True,

304 ):

305 + # If MergedReplicatedLinear, use output size of each partition.

306 + if hasattr(self, "output_sizes"):

307 + self.output_partition_sizes = self.output_sizes

308 + else:

309 + self.output_partition_sizes = [output_size]

310 +

311 super().__init__(input_size,

312 output_size,

313 skip_bias_add,

@@ -311,7 +319,8 @@ class ReplicatedLinear(LinearBase):

319 # All the linear layer supports quant method.

320 assert self.quant_method is not None

321 self.quant_method.create_weights(self,

314 - self.input_size, [self.output_size],

322 + self.input_size,

323 + self.output_partition_sizes,

324 self.input_size,

325 self.output_size,

326 self.params_dtype,

@@ -367,6 +376,73 @@ class ReplicatedLinear(LinearBase):

376 return s

377

378

379 +class MergedReplicatedLinear(ReplicatedLinear):

380 + """Replicated linear layer.

381 +

382 + Args:

383 + input_size: input dimension of the linear layer.

384 + output_size: output dimension of the linear layer.

385 + bias: If true, add bias.

386 + skip_bias_add: If true, skip adding bias but instead return it.

387 + params_dtype: Data type for the parameters.

388 + quant_config: Quantization configure.

389 + prefix: The name of the layer in the state dict, including all parents

390 + (e.g. model.layers.0.qkv_proj)

391 + """

392 +

393 + def __init__(

394 + self,

395 + input_size: int,

396 + output_sizes: list[int],

397 + bias: bool = True,

398 + skip_bias_add: bool = False,

399 + params_dtype: Optional[torch.dtype] = None,

400 + quant_config: Optional[QuantizationConfig] = None,

401 + prefix: str = "",

402 + *,

403 + return_bias: bool = True,

404 + ):

405 + self.output_sizes = output_sizes

406 + super().__init__(input_size,

407 + sum(output_sizes),

408 + bias,

409 + skip_bias_add,

410 + params_dtype,

411 + quant_config,

412 + prefix=prefix,

413 + return_bias=return_bias)

414 +

415 + def weight_loader(self,

416 + param: Union[Parameter, BasevLLMParameter],

417 + loaded_weight: torch.Tensor,

418 + loaded_shard_id: Optional[int] = None):

419 + assert loaded_shard_id is not None

420 + assert loaded_shard_id < len(self.output_sizes)

421 +

422 + if isinstance(param, BlockQuantScaleParameter):

423 + from vllm.model_executor.layers.quantization.fp8 import (

424 + Fp8LinearMethod, Fp8MoEMethod)

425 + assert self.quant_method is not None

426 + assert isinstance(self.quant_method,

427 + (Fp8LinearMethod, Fp8MoEMethod))

428 + weight_block_size = self.quant_method.quant_config.weight_block_size

429 + assert weight_block_size is not None

430 + block_n, _ = weight_block_size[0], weight_block_size[1]

431 + shard_offset = (

432 + (sum(self.output_sizes[:loaded_shard_id]) + block_n - 1) //

433 + block_n)

434 + shard_size = ((self.output_sizes[loaded_shard_id] + block_n - 1) //

435 + block_n)

436 + elif isinstance(param, PerTensorScaleParameter):

437 + shard_offset = loaded_shard_id

438 + shard_size = 1

439 + else:

440 + shard_offset = sum(self.output_sizes[:loaded_shard_id])

441 + shard_size = self.output_sizes[loaded_shard_id]

442 +

443 + param[shard_offset:shard_offset + shard_size] = loaded_weight

444 +

445 +

446 class ColumnParallelLinear(LinearBase):

447 """Linear layer with column parallelism.

448

vllm/model_executor/layers/quantization/fp8.py

+10 -3

@@ -257,9 +257,16 @@ class Fp8LinearMethod(LinearMethodBase):

257 f"{input_size_per_partition} is not divisible by "

258 f"weight quantization block_k = {block_k}.")

259 # Required by column parallel or enabling merged weights

260 - if (tp_size > 1 and output_size // output_size_per_partition

261 - == tp_size) or len(output_partition_sizes) > 1:

262 - for output_partition_size in output_partition_sizes:

260 + is_tp_split = (tp_size > 1 and

261 + output_size // output_size_per_partition == tp_size)

262 + is_merged_gemm = len(output_partition_sizes) > 1

263 + if is_tp_split or is_merged_gemm:

264 + sizes_to_check = output_partition_sizes

265 + if not is_tp_split and is_merged_gemm:

266 + # In case of merged matrices, we allow the last

267 + # matrix to not be a multiple of block size

268 + sizes_to_check = output_partition_sizes[:-1]

269 + for output_partition_size in sizes_to_check:

270 if output_partition_size % block_n != 0:

271 raise ValueError(

272 f"Weight output_partition_size = "

vllm/model_executor/models/deepseek_v2.py

+39 -18

@@ -42,6 +42,7 @@ from vllm.model_executor.layers.fused_moe import FusedMoE

42 from vllm.model_executor.layers.layernorm import RMSNorm

43 from vllm.model_executor.layers.linear import (ColumnParallelLinear,

44 MergedColumnParallelLinear,

45 + MergedReplicatedLinear,

46 ReplicatedLinear,

47 RowParallelLinear)

48 from vllm.model_executor.layers.logits_processor import LogitsProcessor

@@ -336,7 +337,7 @@ class DeepseekV2Attention(nn.Module):

337 kv_a, _ = latent_cache.split(

338 [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)

339 latent_cache = latent_cache.unsqueeze(1)

339 - kv_a = self.kv_a_layernorm(kv_a.contiguous())

340 + kv_a = self.kv_a_layernorm(kv_a)

341 kv = self.kv_b_proj(kv_a)[0]

342 kv = kv.view(-1, self.num_local_heads,

343 self.qk_nope_head_dim + self.v_head_dim)

@@ -407,14 +408,24 @@ class DeepseekV2MLAAttention(nn.Module):

408 self.max_position_embeddings = max_position_embeddings

409

410 if self.q_lora_rank is not None:

410 - self.q_a_proj = ReplicatedLinear(self.hidden_size,

411 - self.q_lora_rank,

412 - bias=False,

413 - quant_config=quant_config,

414 - prefix=f"{prefix}.q_a_proj")

411 + self.fused_qkv_a_proj = MergedReplicatedLinear(

412 + self.hidden_size,

413 + [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],

414 + bias=False,

415 + quant_config=quant_config,

416 + prefix=f"{prefix}.fused_qkv_a_proj")

417 + else:

418 + self.kv_a_proj_with_mqa = ReplicatedLinear(

419 + self.hidden_size,

420 + self.kv_lora_rank + self.qk_rope_head_dim,

421 + bias=False,

422 + quant_config=quant_config,

423 + prefix=f"{prefix}.kv_a_proj_with_mqa")

424 +

425 + if self.q_lora_rank is not None:

426 self.q_a_layernorm = RMSNorm(self.q_lora_rank,

427 eps=config.rms_norm_eps)

417 - self.q_b_proj = ColumnParallelLinear(q_lora_rank,

428 + self.q_b_proj = ColumnParallelLinear(self.q_lora_rank,

429 self.num_heads *

430 self.qk_head_dim,

431 bias=False,

@@ -427,13 +438,6 @@ class DeepseekV2MLAAttention(nn.Module):

438 bias=False,

439 quant_config=quant_config,

440 prefix=f"{prefix}.q_proj")

430 -

431 - self.kv_a_proj_with_mqa = ReplicatedLinear(

432 - self.hidden_size,

433 - self.kv_lora_rank + self.qk_rope_head_dim,

434 - bias=False,

435 - quant_config=quant_config,

436 - prefix=f"{prefix}.kv_a_proj_with_mqa")

441 self.kv_a_layernorm = RMSNorm(self.kv_lora_rank,

442 eps=config.rms_norm_eps)

443 self.kv_b_proj = ColumnParallelLinear(

@@ -495,15 +499,24 @@ class DeepseekV2MLAAttention(nn.Module):

499 positions: torch.Tensor,

500 hidden_states: torch.Tensor,

501 ) -> torch.Tensor:

502 + q_c = None

503 + kv_lora = None

504 +

505 if self.q_lora_rank is not None:

499 - q_c = self.q_a_proj(hidden_states)[0]

506 + qkv_lora = self.fused_qkv_a_proj(hidden_states)[0]

507 + q_c, kv_lora = qkv_lora.split(

508 + [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],

509 + dim=-1,

510 + )

511 q_c = self.q_a_layernorm(q_c)

512 q = self.q_b_proj(q_c)[0]

513 else:

514 + kv_lora = self.kv_a_proj_with_mqa(hidden_states)[0]

515 q = self.q_proj(hidden_states)[0]

504 - kv_c, k_pe = self.kv_a_proj_with_mqa(hidden_states)[0].split(

505 - [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)

506 - kv_c_normed = self.kv_a_layernorm(kv_c.contiguous())

516 +

517 + kv_c, k_pe = kv_lora.split([self.kv_lora_rank, self.qk_rope_head_dim],

518 + dim=-1)

519 + kv_c_normed = self.kv_a_layernorm(kv_c)

520

521 q = q.view(-1, self.num_local_heads, self.qk_head_dim)

522 # Add head dim of 1 to k_pe

@@ -837,6 +850,8 @@ class DeepseekV2ForCausalLM(nn.Module, SupportsPP, MixtureOfExperts):

850 # (param_name, shard_name, shard_id)

851 ("gate_up_proj", "gate_proj", 0),

852 ("gate_up_proj", "up_proj", 1),

853 + ("fused_qkv_a_proj", "q_a_proj", 0),

854 + ("fused_qkv_a_proj", "kv_a_proj_with_mqa", 1),

855 ]

856

857 # Params for weights, fp8 weight scales, fp8 activation scales

@@ -871,6 +886,12 @@ class DeepseekV2ForCausalLM(nn.Module, SupportsPP, MixtureOfExperts):

886 if (("mlp.experts." in name) and name not in params_dict):

887 continue

888 name = name.replace(weight_name, param_name)

889 +

890 + # QKV fusion is optional, fall back to normal

891 + # weight loading if it's not enabled

892 + if ((param_name == "fused_qkv_a_proj")

893 + and name not in params_dict):

894 + continue

895 # Skip loading extra bias for GPTQ models.

896 if name.endswith(".bias") and name not in params_dict:

897 continue