[Puzzletron] Runtime stats benchmark model computes head size as hidden_size // num_attention_heads instead of reading head_dim

Abierto
  • ¿Algo no está bien?
#3,016 0 comentarios 0 reacciones 1 asignado Ver en GitHub

Los mantenedores suelen responder en 2 días

@grzegorz-k-karch ya está trabajando en esto.

Desde el 26/10/2026.

  • #3112 de @efegokdemir — abierto

Evaluación

Este issue todavía no se ha evaluado.

Descripción

bug

Follow-up to #3009, which fixed the Puzzletron memory objective to use the model's head_dim for the KV cache (#2482).

The runtime objective has the same root cause. calc_runtime_stats.create_benchmark_model builds the synthetic Llama benchmark model with:

model_config = LlamaConfig(
    hidden_size=hidden_size,  # candidate hidden size
    num_attention_heads=num_attention_heads,
    head_dim=None,  # Compute from hidden_size // num_attention_heads instead of using default 128
    ...
)

So when runtime_stats_enabled=True, the benchmarked attention uses hidden_size // num_attention_heads as its head size, which is wrong in two cases:

  • Models that set head_dim independently, e.g. Mistral-Small-24B (160 instead of 128) or gpt-oss-20b (45 instead of 64).
  • Every candidate with a pruned hidden_size. Pruning leaves head_dim unchanged, but here it shrinks with the candidate hidden size.

As a result, the latency numbers fed to the MIP are skewed the same way the memory numbers were before #3009.

Suggested fix: pass the parent's head size, get_attention_head_dim(descriptor.get_language_model_config(model_config)) from modelopt/torch/puzzletron/utils/misc.py, into create_benchmark_model as head_dim. Validate with a vLLM runtime-stats run.

Raised in the review of #3009: https://github.com/NVIDIA/Model-Optimizer/pull/3009#discussion_r4374334841

Lenguaje dominante
Python
Estrellas
5.5k
Forks
785
Merge medio
3 d 22 h
PR fusionados (30 d)
277

Preparar el entorno

Primeros pasos

  1. Lee el issue completo y luego la guía de contribución del proyecto.
  2. Comenta en el issue que vas a ocuparte — evita que dos personas hagan lo mismo.
  3. Haz un fork del repositorio y trabaja en una rama.
  4. Abre un pull request que haga referencia al número del issue.

Más de NVIDIA/Model-Optimizer

Todos los issues de NVIDIA/Model-Optimizer

Issues similares

Más issues de Python

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.