A high-throughput and memory-efficient inference and serving engine for LLMs
[AMD][Bugfix][Quantization] Honor fused-name match in is_layer_skipped (#43981)
RoyWang committed
Jun 16, 2026 at 00:37 UTC
a3195fab7b1227e75403fef83891e961e69228a6
2 files changed
+97
-1
tests/quantization/test_quark.py
+88
@@ -25,6 +25,9 @@ from vllm.model_executor.layers.quantization.quark.quark import ( # noqa: E501
25
from vllm.model_executor.layers.quantization.quark.quark_moe import ( # noqa: E501
26
QuarkW8A8Int8MoEMethod,
27
)
28
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
29
+ is_layer_skipped,
30
+)
31
from vllm.platforms import current_platform
32
33
from .reference_mxfp4 import dq_mxfp4_torch, qdq_mxfp4_torch
@@ -437,3 +440,88 @@ def test_mxfp4_dequant_kernel_match_quark(
440
out_torch = dq_mxfp4_torch(w_mxfp4, scale, float_dtype)
441
442
assert torch.equal(out_hip, out_torch)
443
+
444
+
445
+# Unit tests for ``is_layer_skipped`` fused-name handling.
446
+
447
+FUSED_MAPPING = {
448
+ "qkv_proj": ["q_proj", "k_proj", "v_proj"],
449
+ "gate_up_proj": ["gate_proj", "up_proj"],
450
+}
451
+
452
+
453
+def test_fused_name_listed_directly_is_skipped():
454
+ # Regression for Step-3.5-Flash-FP8: the checkpoint lists the fused
455
+ # name (``qkv_proj``) directly in ``modules_to_not_convert``. When a
456
+ # ``packed_modules_mapping`` is registered on the model, the fused
457
+ # match must still win over per-shard expansion.
458
+ ignored = ["model.layers.0.self_attn.qkv_proj"]
459
+ assert is_layer_skipped(
460
+ prefix="model.layers.0.self_attn.qkv_proj",
461
+ ignored_layers=ignored,
462
+ fused_mapping=FUSED_MAPPING,
463
+ )
464
+ assert is_layer_skipped(
465
+ prefix="model.layers.0.mlp.gate_up_proj",
466
+ ignored_layers=["model.layers.0.mlp.gate_up_proj"],
467
+ fused_mapping=FUSED_MAPPING,
468
+ )
469
+
470
+
471
+def test_unfused_shards_listed_is_skipped():
472
+ # Quark INT8 style: per-shard names listed; all shards present means
473
+ # the fused layer is skipped via expansion.
474
+ ignored = [
475
+ "model.layers.0.self_attn.q_proj",
476
+ "model.layers.0.self_attn.k_proj",
477
+ "model.layers.0.self_attn.v_proj",
478
+ ]
479
+ assert is_layer_skipped(
480
+ prefix="model.layers.0.self_attn.qkv_proj",
481
+ ignored_layers=ignored,
482
+ fused_mapping=FUSED_MAPPING,
483
+ )
484
+
485
+
486
+def test_partial_shards_raises():
487
+ # Only some shards listed -> ambiguous, must raise. Fused name is
488
+ # not in ignored_layers, so we fall through to per-shard expansion.
489
+ ignored = ["model.layers.0.self_attn.q_proj"]
490
+ with pytest.raises(ValueError):
491
+ is_layer_skipped(
492
+ prefix="model.layers.0.self_attn.qkv_proj",
493
+ ignored_layers=ignored,
494
+ fused_mapping=FUSED_MAPPING,
495
+ )
496
+
497
+
498
+def test_not_skipped_when_nothing_listed():
499
+ assert not is_layer_skipped(
500
+ prefix="model.layers.0.self_attn.qkv_proj",
501
+ ignored_layers=["model.layers.0.mlp.gate_up_proj"],
502
+ fused_mapping=FUSED_MAPPING,
503
+ )
504
+
505
+
506
+def test_non_fused_layer_unaffected():
507
+ assert is_layer_skipped(
508
+ prefix="model.layers.0.self_attn.o_proj",
509
+ ignored_layers=["model.layers.0.self_attn.o_proj"],
510
+ fused_mapping=FUSED_MAPPING,
511
+ )
512
+ assert not is_layer_skipped(
513
+ prefix="model.layers.0.self_attn.o_proj",
514
+ ignored_layers=["model.layers.1.self_attn.o_proj"],
515
+ fused_mapping=FUSED_MAPPING,
516
+ )
517
+
518
+
519
+def test_substr_match_on_fused_name():
520
+ # skip_with_substr=True path: fused-name substring match should also
521
+ # short-circuit before shard expansion.
522
+ assert is_layer_skipped(
523
+ prefix="model.layers.0.self_attn.qkv_proj",
524
+ ignored_layers=["self_attn.qkv_proj"],
525
+ fused_mapping=FUSED_MAPPING,
526
+ skip_with_substr=True,
527
+ )
vllm/model_executor/layers/quantization/utils/quant_utils.py
+9
-1
@@ -520,7 +520,15 @@ def is_layer_skipped(
520
# in the safetensors checkpoint. So, we convert the name
521
# from the fused version to unfused + check to make sure that
522
# each shard of the fused layer has the same scheme.
523
- if proj_name in fused_mapping:
523
+ #
524
+ # Some checkpoints (e.g. block-FP8 Step-3.5-Flash) already list the
525
+ # fused name (e.g. ``self_attn.qkv_proj``) directly in
526
+ # ``modules_to_not_convert``. Honor that fused-name match first so
527
+ # those layers are still correctly skipped even when a
528
+ # ``packed_modules_mapping`` is registered on the model.
529
+ if proj_name in fused_mapping and match_func(prefix, ignored_layers):
530
+ is_skipped = True
531
+ elif proj_name in fused_mapping:
532
shard_prefixes = [
533
prefix.replace(proj_name, shard_proj_name)
534
for shard_proj_name in fused_mapping[proj_name]