Skip to content
Merged
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion benchmarks/kernels/benchmark_moe.py
Original file line number Diff line number Diff line change
Expand Up @@ -560,7 +560,7 @@ def save_configs(
filename = os.path.join(save_dir, filename)
print(f"Writing best config to {filename}...")
with open(filename, "w") as f:
json.dump(configs, f, indent=4)
json.dump({"triton_version": triton.__version__, **configs}, f, indent=4)
f.write("\n")


Expand Down
5 changes: 4 additions & 1 deletion vllm/model_executor/layers/fused_moe/fused_moe.py
Original file line number Diff line number Diff line change
Expand Up @@ -720,7 +720,10 @@ def get_moe_configs(
logger.info("Using configuration from %s for MoE layer.",
config_file_path)
# If a configuration has been found, return it
return {int(key): val for key, val in json.load(f).items()}
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

The proposed change to explicitly pop triton_version is brittle. If other metadata is added to the configuration file in the future, this code will fail with a ValueError when trying to convert a non-integer string key to an integer.

A more robust approach is to filter for keys that are valid integers. By checking key.isdigit(), you can automatically ignore any non-numeric keys, including triton_version and any future metadata, making the code more resilient to future changes.

Suggested change
return {int(key): val for key, val in json.load(f).items()}
return {int(key): val for key, val in json.load(f).items() if key.isdigit()}

tuned_config = json.load(f)
# Delete triton_version from tuned_config
tuned_config.pop("triton_version", None)
return {int(key): val for key, val in tuned_config.items()}

# If no optimized configuration is available, we will use the default
# configuration
Expand Down
2 changes: 1 addition & 1 deletion vllm/triton_utils/importing.py
Original file line number Diff line number Diff line change
Expand Up @@ -68,7 +68,7 @@ class TritonPlaceholder(types.ModuleType):

def __init__(self):
super().__init__("triton")
self.__version__ = "3.3.0"
self.__version__ = "3.4.0"
self.jit = self._dummy_decorator("jit")
self.autotune = self._dummy_decorator("autotune")
self.heuristics = self._dummy_decorator("heuristics")
Expand Down