Skip to content

Commit 7fc7837

Browse files
committed
Address review feedback: always emit qk_norm_epsilon, cache mxfp4 weight map, drop stray whitespace
- base.py: always add qk_norm_epsilon GQA attribute (matches convention of unconditionally-added attributes; value only used when q/k norm weights are present) - base.py: remove whitespace-only line after CudaQuantizer import - gptoss.py: add explicit _mxfp4_weight_map_loaded flag so the no-index case caches instead of re-globbing the snapshot dir on every tensor load
1 parent 7edb10a commit 7fc7837

2 files changed

Lines changed: 5 additions & 4 deletions

File tree

src/python/py/models/builders/base.py

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@
3333
)
3434

3535
from .cuda_quantizer import CudaQuantizer
36-
36+
3737

3838
class Model:
3939
def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
@@ -2751,9 +2751,8 @@ def make_group_query_attention(self, name, **kwargs):
27512751
"softcap": self.attention_attrs["softcap"],
27522752
"do_rotary": self.attention_attrs["use_rope_in_attn"],
27532753
"rotary_interleaved": self.rope_attrs["interleaved"],
2754+
"qk_norm_epsilon": kwargs.get("qk_norm_epsilon", self.attention_attrs["qk_norm_epsilon"]),
27542755
}
2755-
if q_norm_weight:
2756-
attributes["qk_norm_epsilon"] = kwargs.get("qk_norm_epsilon", self.attention_attrs["qk_norm_epsilon"])
27572756
self.make_node(
27582757
"GroupQueryAttention",
27592758
inputs=inputs,

src/python/py/models/builders/gptoss.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
2222
self.moe_attrs["normalize_routing_weights"] = True
2323
self.moe_attrs["swiglu_fusion"] = 1
2424
self._mxfp4_weight_map = None
25+
self._mxfp4_weight_map_loaded = False
2526
self._mxfp4_snapshot_dir = None
2627
self._mxfp4_tensor_cache = {}
2728

@@ -50,7 +51,7 @@ def _get_mxfp4_snapshot_dir(self):
5051
return self._mxfp4_snapshot_dir
5152

5253
def _get_mxfp4_weight_map(self):
53-
if self._mxfp4_weight_map is not None:
54+
if self._mxfp4_weight_map_loaded:
5455
return self._mxfp4_weight_map
5556

5657
snapshot_dir = self._get_mxfp4_snapshot_dir()
@@ -67,6 +68,7 @@ def _get_mxfp4_weight_map(self):
6768
f"Could not locate original GPT-OSS MXFP4 safetensors index in {snapshot_dir}."
6869
)
6970
self._mxfp4_weight_map = None
71+
self._mxfp4_weight_map_loaded = True
7072
return self._mxfp4_weight_map
7173

7274
def _load_original_mxfp4_tensor(self, tensor_name):

0 commit comments

Comments
 (0)