diff --git a/python/aisimulate/src/aisimulate_core/sdk/models/deepseek_v32.py b/python/aisimulate/src/aisimulate_core/sdk/models/deepseek_v32.py index 6e0064ed5c..69399d7262 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/models/deepseek_v32.py +++ b/python/aisimulate/src/aisimulate_core/sdk/models/deepseek_v32.py @@ -279,10 +279,7 @@ def create(cls, model_info: dict, model_config, backend_name: str) -> BaseModel: # -> is_layer_skipped, modelopt.py:150-161 @0.24.0). Exclusions are # PER-PROJECTION: DeepSeek-V3.2-NVFP4 keeps q/kv/indexer in BF16 but # quantizes o_proj; GLM-5 NVFP4 excludes the whole self_attn block. - extra_params.setdefault( - "dsa_attn_quant_exclusions", - attention_projection_exclusions(raw_config, precise_module_paths=backend_name == "vllm"), - ) + extra_params.setdefault("dsa_attn_quant_exclusions", attention_projection_exclusions(raw_config)) if _shared_experts_excluded_from_quant(raw_config): extra_params.setdefault("dsa_shared_expert_quant_mode", common.GEMMQuantMode.bfloat16) # GLM-5.2 shares one DSA topk index across ``index_topk_freq`` layers diff --git a/python/aisimulate/src/aisimulate_core/sdk/models/helpers.py b/python/aisimulate/src/aisimulate_core/sdk/models/helpers.py index f800a8377c..062d4dd1af 100644 --- a/python/aisimulate/src/aisimulate_core/sdk/models/helpers.py +++ b/python/aisimulate/src/aisimulate_core/sdk/models/helpers.py @@ -86,34 +86,32 @@ def quant_exclude_patterns(raw_config: dict) -> list: "q": ("q_a_proj", "q_b_proj", "q_proj"), "kv": ("kv_a_proj", "kv_b_proj", "kv_proj", "k_proj", "v_proj"), "o": ("o_proj",), - "indexer": ("indexer",), -} -_PRECISE_PROJECTION_GROUP_MARKERS = { - **_PROJECTION_GROUP_MARKERS, "indexer": ("indexer.wq_b", "indexer.wk", "indexer.weights_proj"), } _PROJECTION_MODULE_PATHS = { - **_PRECISE_PROJECTION_GROUP_MARKERS, + **_PROJECTION_GROUP_MARKERS, "kv": (*_PROJECTION_GROUP_MARKERS["kv"], "kv_a_proj_with_mqa"), - "indexer": (*_PRECISE_PROJECTION_GROUP_MARKERS["indexer"], "indexer.wk_weights_proj"), + "indexer": (*_PROJECTION_GROUP_MARKERS["indexer"], "indexer.wk_weights_proj"), } -def attention_projection_exclusions(raw_config: dict, *, precise_module_paths: bool = False) -> frozenset: +def attention_projection_exclusions(raw_config: dict) -> frozenset: """Which attention projection groups the checkpoint keeps unquantized. Returns a subset of :data:`ATTENTION_PROJECTION_GROUPS`. A pattern naming the whole block (``self_attn*`` / ``re:.*self_attn.*``) covers every - group; otherwise groups are matched per projection name. + group; otherwise groups are matched per projection module path. - vLLM DSA opts into precise module paths: norm and auxiliary exclusions - do not exclude projection GEMMs. Other consumers retain their existing - interpretation until their precision/data correction is qualified. + Norm and auxiliary exclusions do not exclude projection GEMMs. Native GLM + FP8 checkpoints list ``q_a_layernorm``, ``kv_a_layernorm``, + ``indexer.k_norm`` and ``indexers_proj`` while shipping their attention + projections in block FP8; vLLM, SGLang and TRT-LLM skip only the listed + module (or its descendants), so every backend runs those GEMMs quantized. A targets-only ModelOpt MIXED_PRECISION export (see :func:`modelopt_mixed_precision_targets`) additionally keeps every group - that no ``quantized_layers`` entry names unquantized, on both paths. As - on vLLM and SGLang, only a key naming a projection (or a fused module + that no ``quantized_layers`` entry names unquantized. As on vLLM and + SGLang, only a key naming a projection (or a fused module packing it) exactly, or a descendant of one, lists it; an ancestor key (``...self_attn``, ``...self_attn.indexer``) quantizes nothing. Keys the exclusion list covers and keys at layer index >= ``num_hidden_layers`` @@ -121,20 +119,19 @@ def attention_projection_exclusions(raw_config: dict, *, precise_module_paths: b listed in any decoder layer counts as quantized (the DSA indexer exists only in full layers). Exclusion globs still apply and still win. """ - excluded = _glob_attention_projection_exclusions(raw_config, precise_module_paths=precise_module_paths) + excluded = _glob_attention_projection_exclusions(raw_config) targets = modelopt_mixed_precision_targets(raw_config) if targets is not None: excluded |= frozenset(ATTENTION_PROJECTION_GROUPS) - _modelopt_listed_attention_groups(raw_config, targets) return excluded -def _glob_attention_projection_exclusions(raw_config: dict, *, precise_module_paths: bool) -> frozenset: +def _glob_attention_projection_exclusions(raw_config: dict) -> frozenset: """Attention projection groups named by the checkpoint's exclusion patterns.""" - markers_by_group = _PRECISE_PROJECTION_GROUP_MARKERS if precise_module_paths else _PROJECTION_GROUP_MARKERS excluded: set = set() for pattern in quant_exclude_patterns(raw_config): p = str(pattern) - if precise_module_paths and p.startswith("re:"): + if p.startswith("re:"): # Compressed-tensors regexes match native module paths, including # anchors and escaped dots; projection-name substrings are not # sufficient (e.g. a norm-only regex must not exclude a GEMM). @@ -150,13 +147,6 @@ def _glob_attention_projection_exclusions(raw_config: dict, *, precise_module_pa ): excluded.add(group) continue - if not precise_module_paths: - if "self_attn" in p and not any(m in p for markers in markers_by_group.values() for m in markers): - return frozenset(ATTENTION_PROJECTION_GROUPS) - for group, markers in markers_by_group.items(): - if any(m in p for m in markers): - excluded.add(group) - continue # A named child such as q_a_layernorm is not an exclusion of the # entire attention block. Native GLM FP8 configs exclude many such # norms while retaining FP8 attention projections. @@ -167,7 +157,7 @@ def _glob_attention_projection_exclusions(raw_config: dict, *, precise_module_pa # the separate indexers_proj module. Neither excludes wq_b/wk GEMMs. if "indexer" in p and p.rsplit("indexer", 1)[1] in block_suffixes: excluded.add("indexer") - for group, markers in markers_by_group.items(): + for group, markers in _PROJECTION_GROUP_MARKERS.items(): if any(m in p for m in markers): excluded.add(group) return frozenset(excluded) diff --git a/python/aisimulate/tests/unit/sdk/models/test_model_config.py b/python/aisimulate/tests/unit/sdk/models/test_model_config.py index 7698a27c2a..52c18ad554 100644 --- a/python/aisimulate/tests/unit/sdk/models/test_model_config.py +++ b/python/aisimulate/tests/unit/sdk/models/test_model_config.py @@ -1915,23 +1915,10 @@ class TestAttentionProjectionExclusions: """Per-projection exclusion parsing (V3.1/V3.2 exclude q/kv but not o_proj).""" @staticmethod - def _excl(patterns, *, precise_module_paths=True): + def _excl(patterns): from aisimulate.sdk.models.helpers import attention_projection_exclusions - return attention_projection_exclusions( - {"quantization_config": {"ignore": patterns}}, precise_module_paths=precise_module_paths - ) - - def test_existing_consumers_keep_their_projection_exclusion_interpretation(self): - from aisimulate.sdk.models.helpers import attention_projection_exclusions - - # SGLang, TRT-LLM and ordinary MLA retain the default exclusion matching. - # vLLM DSA opts into precise projection paths. - raw = {"quantization_config": {"ignore": ["model.layers.0.self_attn.q_a_layernorm"]}} - assert attention_projection_exclusions(raw) == frozenset({"q", "kv", "o", "indexer"}) - assert self._excl(["model.layers.0.self_attn.indexers_proj"], precise_module_paths=False) == frozenset( - {"indexer"} - ) + return attention_projection_exclusions({"quantization_config": {"ignore": patterns}}) def test_whole_block_glob_covers_all_groups(self): assert self._excl(["model.layers.3.self_attn*"]) == frozenset({"q", "kv", "o", "indexer"}) @@ -1999,25 +1986,44 @@ def test_norm_and_auxiliary_exclusions_do_not_reclassify_projection_groups(self) def test_named_indexer_projection_exclusion(self, projection): assert self._excl([f"model.layers.0.self_attn.indexer.{projection}"]) == frozenset({"indexer"}) - @pytest.mark.parametrize("backend", ["vllm", "sglang", "trtllm"]) - @pytest.mark.parametrize("phase", ["context", "generation"]) - @pytest.mark.parametrize("quant", ["fp8", "fp8_block", "nvfp4"]) - def test_native_glm52_precision_correction_is_scoped_to_vllm_dsa(self, backend, phase, quant): - # The real checkpoint excludes layernorms and indexers_proj, not the - # q/kv/o projection GEMMs. NVFP4 mixed exclusions have separate tests - # in TestDSV32NVFP4AttentionExclusion above. + @staticmethod + def _dsa_gemm_mode(model_path, backend, phase, **model_config): from aisimulate.sdk import engine model = models.get_model( - "zai-org/GLM-5.2-FP8", - config.ModelConfig(tp_size=8, moe_tp_size=1, moe_ep_size=8, gemm_quant_mode=common.GEMMQuantMode[quant]), + model_path, + config.ModelConfig(tp_size=8, moe_tp_size=1, moe_ep_size=8, **model_config), backend_name=backend, ) - expected_exclusions = frozenset() if backend == "vllm" else frozenset({"q", "kv", "o", "indexer"}) - assert model.extra_params["dsa_attn_quant_exclusions"] == expected_exclusions specs = json.loads(engine._ops_json(getattr(model, f"{phase}_ops"))) attention = next(fields for spec in specs for tag, fields in spec.items() if tag == f"Dsa{phase.capitalize()}") - assert attention["gemm_quant_mode"] == (quant if backend == "vllm" else "bfloat16") + return model.extra_params["dsa_attn_quant_exclusions"], attention["gemm_quant_mode"] + + @pytest.mark.parametrize("backend", ["vllm", "sglang", "trtllm"]) + @pytest.mark.parametrize("phase", ["context", "generation"]) + @pytest.mark.parametrize("quant", ["fp8", "fp8_block", "nvfp4"]) + def test_native_glm52_projections_follow_the_gemm_mode_on_every_backend(self, backend, phase, quant): + # The real checkpoint excludes layernorms and indexers_proj, not the + # q/kv/o projection GEMMs, and every backend loads it that way. NVFP4 + # mixed exclusions have separate tests in + # TestDSV32NVFP4AttentionExclusion above. + exclusions, gemm = self._dsa_gemm_mode( + "zai-org/GLM-5.2-FP8", backend, phase, gemm_quant_mode=common.GEMMQuantMode[quant] + ) + assert exclusions == frozenset() + assert gemm == quant + + @pytest.mark.parametrize("backend", ["vllm", "sglang", "trtllm"]) + @pytest.mark.parametrize("phase", ["context", "generation"]) + @pytest.mark.parametrize( + "model_path", ["zai-org/GLM-5-FP8", "zai-org/GLM-5.1-FP8", "zai-org/GLM-5.2-FP8", "zai-org/GLM-5.3"] + ) + def test_native_glm_fp8_checkpoint_prices_block_fp8_projections(self, model_path, backend, phase): + # quant_method fp8 with weight_block_size [128, 128]: without an + # explicit override, every backend's DSA module key is fp8_block. + exclusions, gemm = self._dsa_gemm_mode(model_path, backend, phase) + assert exclusions == frozenset() + assert gemm == "fp8_block" @pytest.mark.parametrize("backend", ["vllm", "sglang", "trtllm"]) @pytest.mark.parametrize( diff --git a/python/aisimulate/tests/unit/sdk/models/test_modelopt_mixed_precision_targets.py b/python/aisimulate/tests/unit/sdk/models/test_modelopt_mixed_precision_targets.py index d508e3d4f4..31be1001ff 100644 --- a/python/aisimulate/tests/unit/sdk/models/test_modelopt_mixed_precision_targets.py +++ b/python/aisimulate/tests/unit/sdk/models/test_modelopt_mixed_precision_targets.py @@ -261,12 +261,11 @@ def test_conflicting_algorithms_for_one_module_raise(self): } -@pytest.mark.parametrize("precise", [False, True]) @pytest.mark.parametrize("case", list(_ATTENTION_CASES)) -def test_attention_groups_absent_from_quantized_layers_are_excluded(case, precise): +def test_attention_groups_absent_from_quantized_layers_are_excluded(case): extra, exclude, expected = _ATTENTION_CASES[case] raw = _raw(_glm53_layers(extra=extra), exclude=("model.layers.78*", *exclude)) - assert H.attention_projection_exclusions(raw, precise_module_paths=precise) == expected + assert H.attention_projection_exclusions(raw) == expected def test_keys_beyond_the_decoder_layers_list_nothing(): diff --git a/python/aisimulate/tools/prediction_regression_gate/collect_static.py b/python/aisimulate/tools/prediction_regression_gate/collect_static.py index 299a30e12a..424425f5bc 100644 --- a/python/aisimulate/tools/prediction_regression_gate/collect_static.py +++ b/python/aisimulate/tools/prediction_regression_gate/collect_static.py @@ -114,7 +114,7 @@ def collect_combo(combo: grid.Combo) -> tuple[grid.Combo, list[dict], float]: parallel = grid.MOE_PARALLEL if is_moe else grid.DENSE_PARALLEL for tp, pp, adp, moe_tp, moe_ep in parallel: - for quant_label, gemm_quant, moe_quant in grid.quant_variants_for(combo.system): + for quant_label, gemm_quant, moe_quant in grid.QUANT_VARIANTS: base = { "model": model_name, "tp": tp, diff --git a/python/aisimulate/tools/prediction_regression_gate/grid.py b/python/aisimulate/tools/prediction_regression_gate/grid.py index 914916f44c..c3836822c1 100644 --- a/python/aisimulate/tools/prediction_regression_gate/grid.py +++ b/python/aisimulate/tools/prediction_regression_gate/grid.py @@ -10,8 +10,7 @@ Axes: - shape: fixed prefill (static_ctx) and decode (static_gen) points - parallelism: per model family (dense vs MoE) - - quant: model default plus explicit overrides (nvfp4 only on systems whose - data can contain it) + - quant: the checkpoint's own quantization only (see QUANT_VARIANTS) """ from __future__ import annotations @@ -105,15 +104,16 @@ def _dir_is_incomplete(path: str) -> bool: # -------------------------------------------------------------------------- # Quant axis: label -> (gemm_quant_mode, moe_quant_mode) overrides. -# "default" keeps the model config's own quantization. +# "default" keeps the model config's own quantization, and it is the only +# variant. Forcing one GEMM/MoE mode onto every bundled model pairs already +# quantized checkpoints with formats they do not ship (per-tensor FP8 or NVFP4 +# projections on a block-FP8 GLM checkpoint), and each such row then demands +# perf data no deployment uses. Bundled FP8 and NVFP4 checkpoints cover those +# precisions as they are actually served. # -------------------------------------------------------------------------- QUANT_VARIANTS: list[tuple[str, str | None, str | None]] = [ ("default", None, None), - ("fp8", "fp8", "fp8"), ] -# Only meaningful where NVFP4 silicon can exist (sm100+ family systems). -NVFP4_SYSTEMS = {"b200_sxm", "b300_sxm", "gb200", "gb300"} -NVFP4_VARIANT: tuple[str, str | None, str | None] = ("nvfp4", "nvfp4", "nvfp4") @dataclass(frozen=True) @@ -127,13 +127,6 @@ def relpath(self) -> str: return os.path.join(self.system, self.backend, f"{self.version}.csv") -def quant_variants_for(system: str) -> list[tuple[str, str | None, str | None]]: - variants = list(QUANT_VARIANTS) - if system in NVFP4_SYSTEMS: - variants.append(NVFP4_VARIANT) - return variants - - def bundled_models() -> list[str]: """Offline-buildable bundled models (HF-download quant variants excluded).""" names = []