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 14, 05:30 UTC
[Mamba] Support TP>1 with quantization for mamba2 mixer in case `n_groups % tp_size == 0` (#24593)
Signed-off-by: Tomer Asida <57313761+tomeras91@users.noreply.github.com> Signed-off-by: tomeras91 <57313761+tomeras91@users.noreply.github.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
tomeras91 committed
Sep 16, 2025 at 13:51 UTC
27fcfe7bcfbec033b12a40c2e82342ea2b25bb43
1 file changed
+119
-84
vllm/model_executor/layers/mamba/mamba_mixer2.py
+119
-84
@@ -19,6 +19,7 @@ from vllm.distributed import (divide, get_tensor_model_parallel_rank,
19
from vllm.forward_context import ForwardContext, get_forward_context
20
from vllm.model_executor.custom_op import CustomOp
21
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
22
+ MergedColumnParallelLinear,
23
RowParallelLinear)
24
from vllm.model_executor.layers.mamba.abstract import MambaBase
25
from vllm.model_executor.layers.mamba.mamba2_metadata import (Mamba2Metadata,
@@ -261,12 +262,14 @@ class MambaMixer2(MambaBase, CustomOp):
262
), "Tensor parallel world size must divide num heads."
263
264
assert (n_groups % self.tp_size) == 0 or n_groups == 1, (
264
- "If tensor parallel world size does not divide num_heads, "
265
+ "If tensor parallel world size does not divide num_groups, "
266
"then num_groups must equal 1.")
267
267
- assert (
268
- self.tp_size == 1 or quant_config is None
269
- ), "Tensor parallel currently not supported for quantized models."
268
+ assert (n_groups % self.tp_size == 0) or self.tp_size == 1 or \
269
+ quant_config is None, (
270
+ "Tensor parallel currently supported for quantized models only "
271
+ "if tensor parallel world size divides num groups."
272
+ )
273
274
self.ssm_state_size = ssm_state_size
275
self.conv_kernel_size = conv_kernel_size
@@ -285,94 +288,101 @@ class MambaMixer2(MambaBase, CustomOp):
288
n_groups, self.tp_size)
289
self.n_groups = n_groups + groups
290
288
- self.conv_dim = intermediate_size + 2 * self.n_groups * ssm_state_size
289
- self.conv1d = ColumnParallelLinear(
290
- input_size=conv_kernel_size,
291
- output_size=self.conv_dim,
292
- bias=use_conv_bias,
293
- quant_config=None,
294
- prefix=f"{prefix}.conv1d",
295
- )
296
- # unsqueeze to fit conv1d weights shape into the linear weights shape.
297
- # Can't do this in `weight_loader` since it already exists in
298
- # `ColumnParallelLinear` and `set_weight_attrs`
299
- # doesn't allow to override it
300
- self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1)
291
+ self.groups_ssm_state_size = self.n_groups * self.ssm_state_size
292
+ self.conv_dim = intermediate_size + 2 * self.groups_ssm_state_size
293
+
294
+ if n_groups % self.tp_size == 0:
295
+ self.conv1d = MergedColumnParallelLinear(
296
+ input_size=conv_kernel_size,
297
+ output_sizes=[
298
+ intermediate_size,
299
+ self.groups_ssm_state_size,
300
+ self.groups_ssm_state_size,
301
+ ],
302
+ bias=use_conv_bias,
303
+ quant_config=None,
304
+ prefix=f"{prefix}.conv1d",
305
+ )
306
302
- self.in_proj = ColumnParallelLinear(
303
- input_size=hidden_size,
304
- output_size=intermediate_size + self.conv_dim + self.num_heads,
305
- bias=use_bias,
306
- quant_config=quant_config,
307
- prefix=f"{prefix}.in_proj",
308
- )
307
+ self.in_proj = MergedColumnParallelLinear(
308
+ input_size=hidden_size,
309
+ output_sizes=[
310
+ intermediate_size,
311
+ intermediate_size,
312
+ self.groups_ssm_state_size,
313
+ self.groups_ssm_state_size,
314
+ self.num_heads,
315
+ ],
316
+ bias=use_bias,
317
+ quant_config=quant_config,
318
+ prefix=f"{prefix}.in_proj",
319
+ )
320
+ else:
321
+ # This is the n_groups == 1 case,
322
+ # where we need to duplicate groups if TP>1.
323
+
324
+ self.conv1d = ColumnParallelLinear(
325
+ input_size=conv_kernel_size,
326
+ output_size=self.conv_dim,
327
+ bias=use_conv_bias,
328
+ quant_config=None,
329
+ prefix=f"{prefix}.conv1d",
330
+ )
331
310
- # - because in_proj is a concatenation of 3 weights, we
311
- # need to interleave them before sharding
312
- # - use the custom weight loader mamba_v2_sharded_weight_loader
313
- # for conv1d.bias, covn1d.weight and in_proj.weight
314
- # - need to set these settings, to assign the groups to the head shards
315
- group_shard_settings = (
316
- self.n_groups * self.ssm_state_size, # expected model size
317
- (self.n_groups - n_groups) *
318
- self.ssm_state_size, # extra dims assigned
319
- n_groups == 1, # if there was only one group
320
- )
321
- intermediate_settings = (intermediate_size, 0, False)
322
- head_settings = (self.num_heads, 0, False)
323
-
324
- # - the weight already has a "weight_loader" attribute
325
- # which set_weight_attrs will raise if we do not
326
- # delete before trying to override it
327
- # - ditto for the other two weights below
328
- delattr(self.conv1d.bias, "weight_loader")
329
- set_weight_attrs(
330
- self.conv1d.bias,
331
- {
332
- "weight_loader":
333
- mamba_v2_sharded_weight_loader(
334
- [
335
- intermediate_settings,
336
- group_shard_settings,
337
- group_shard_settings,
338
- ],
339
- self.tp_size,
340
- tp_rank,
341
- )
342
- },
343
- )
332
+ self.in_proj = ColumnParallelLinear(
333
+ input_size=hidden_size,
334
+ output_size=intermediate_size + self.conv_dim + self.num_heads,
335
+ bias=use_bias,
336
+ quant_config=quant_config,
337
+ prefix=f"{prefix}.in_proj",
338
+ )
339
345
- delattr(self.conv1d.weight, "weight_loader")
346
- set_weight_attrs(
347
- self.conv1d.weight,
348
- {
349
- "weight_loader":
350
- mamba_v2_sharded_weight_loader(
351
- [
352
- intermediate_settings,
353
- group_shard_settings,
354
- group_shard_settings,
355
- ],
356
- self.tp_size,
357
- tp_rank,
358
- )
359
- },
360
- )
340
+ # - because in_proj is a concatenation of 3 weights, we
341
+ # need to interleave them before sharding
342
+ # - use the custom weight loader mamba_v2_sharded_weight_loader
343
+ # for conv1d.bias, covn1d.weight and in_proj.weight
344
+ # - need to set these settings, to assign the groups
345
+ # to the head shards
346
+ group_shard_settings = (
347
+ self.groups_ssm_state_size, # expected model size
348
+ (self.n_groups - n_groups) *
349
+ self.ssm_state_size, # extra dims assigned
350
+ n_groups == 1, # if there was only one group
351
+ )
352
+ intermediate_settings = (intermediate_size, 0, False)
353
+ head_settings = (self.num_heads, 0, False)
354
+
355
+ # - the weight already has a "weight_loader" attribute
356
+ # which set_weight_attrs will raise if we do not
357
+ # delete before trying to override it
358
+ # - ditto for the other two weights below
359
+ delattr(self.conv1d.bias, "weight_loader")
360
+ set_weight_attrs(
361
+ self.conv1d.bias,
362
+ {
363
+ "weight_loader":
364
+ mamba_v2_sharded_weight_loader(
365
+ [
366
+ intermediate_settings,
367
+ group_shard_settings,
368
+ group_shard_settings,
369
+ ],
370
+ self.tp_size,
371
+ tp_rank,
372
+ )
373
+ },
374
+ )
375
362
- if quant_config is None:
363
- # - quant layers do not have a weight loader
364
- delattr(self.in_proj.weight, "weight_loader")
376
+ delattr(self.conv1d.weight, "weight_loader")
377
set_weight_attrs(
366
- self.in_proj.weight,
378
+ self.conv1d.weight,
379
{
380
"weight_loader":
381
mamba_v2_sharded_weight_loader(
382
[
371
- intermediate_settings, # for gate
383
intermediate_settings,
384
group_shard_settings,
385
group_shard_settings,
375
- head_settings, # for dt
386
],
387
self.tp_size,
388
tp_rank,
@@ -380,6 +390,33 @@ class MambaMixer2(MambaBase, CustomOp):
390
},
391
)
392
393
+ if quant_config is None:
394
+ # - quant layers do not have a weight loader
395
+ delattr(self.in_proj.weight, "weight_loader")
396
+ set_weight_attrs(
397
+ self.in_proj.weight,
398
+ {
399
+ "weight_loader":
400
+ mamba_v2_sharded_weight_loader(
401
+ [
402
+ intermediate_settings, # for gate
403
+ intermediate_settings,
404
+ group_shard_settings,
405
+ group_shard_settings,
406
+ head_settings, # for dt
407
+ ],
408
+ self.tp_size,
409
+ tp_rank,
410
+ )
411
+ },
412
+ )
413
+
414
+ # unsqueeze to fit conv1d weights shape into the linear weights shape.
415
+ # Can't do this in `weight_loader` since it already exists in
416
+ # `ColumnParallelLinear` and `MergedColumnParallelLinear`,
417
+ # and `set_weight_attrs` doesn't allow to override it
418
+ self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1)
419
+
420
# - these are TPed by heads to reduce the size of the
421
# temporal shape
422
self.A = nn.Parameter(
@@ -498,8 +535,6 @@ class MambaMixer2(MambaBase, CustomOp):
535
chunk_indices_p = mamba2_metadata.chunk_indices
536
chunk_offsets_p = mamba2_metadata.chunk_offsets
537
501
- groups_time_state_size = self.n_groups * self.ssm_state_size
502
-
538
# 1. Gated MLP's linear projection
539
projected_states, _ = self.in_proj(hidden_states)
540
@@ -524,8 +559,8 @@ class MambaMixer2(MambaBase, CustomOp):
559
hidden_states_B_C,
560
[
561
self.intermediate_size // self.tp_size,
527
- groups_time_state_size // self.tp_size,
528
- groups_time_state_size // self.tp_size,
562
+ self.groups_ssm_state_size // self.tp_size,
563
+ self.groups_ssm_state_size // self.tp_size,
564
],
565
dim=-1,
566
)