Model Shape Gen

source /root/wch/Hopper_benchmark/.venv/bin/activate
cd /root/wch/Hopper_benchmark

# Single operator, single model
python scripts/generate_model_shapes.py \
    --op rms_norm --models llama3-8b \
  --batch-sizes 32 64 128 --seq-lens 1024 2048 4096

# Multiple operators, multiple models
python scripts/generate_model_shapes.py \
    --op softmax add rms_norm \
    --models llama3-8b deepseek-v3 qwen3-8b \
    --batch-sizes 32 64 128 --seq-lens 1024 2048 4096

# All operators for a model set
python scripts/generate_model_shapes.py \
    --op all --models llama3-8b qwen2-72b \
  --batch-sizes 32 64 128 --seq-lens 1024 2048 4096

# Dry-run (print shapes without writing)
python scripts/generate_model_shapes.py \
    --op embedding --models deepseek-v3 \
    --batch-sizes 32 64 128 --seq-lens 1024 2048 4096 \
    --dry-run

# Append to existing shapes (don't overwrite)
python scripts/generate_model_shapes.py \
    --op rms_norm --models qwen3-8b \
    --batch-sizes 32 64 --seq-lens 1024 4096 \
    --append

# List available models / operators
python scripts/generate_model_shapes.py --list-models
python scripts/generate_model_shapes.py --list-ops

The script handles all steps automatically:

1. Loads model architecture parameters from transformers config classes
2. Maps each operator to its shape formula (flattened T=B*S, batched (B,NH,S,S), special encoding, etc.)
3. Generates all batch × seq_len combinations
4. Deduplicates within and across models
5. No 40 GB filter — all shapes are written regardless of size. Shapes that exceed GPU memory are caught by OOM handling at benchmark runtime and skipped gracefully.
6. Writes model_shape/<op>_shape.py with docstring and per-model grouping
7. Verifies the written file loads correctly via load_model_shapes()

Adding a New Model to the Script

To add a new model, add a loader function and register it in _MODEL_LOADERS in scripts/generate_model_shapes.py:

def _load_new_model() -> ModelConfig:
    from transformers import SomeConfig
    c = SomeConfig(hidden_size=..., ...)
    return ModelConfig(
        name="NewModel-XB", hf_id="org/NewModel-XB",
        hidden_size=c.hidden_size, num_attention_heads=c.num_attention_heads,
        num_kv_heads=c.num_key_value_heads,
        head_dim=c.hidden_size // c.num_attention_heads,
        intermediate_size=c.intermediate_size, vocab_size=c.vocab_size,
    )

_MODEL_LOADERS["newmodel-xb"] = _load_new_model

Manual Reference (Background for the automation script)

When to Use

• Populating model_shape/<op>_shape.py with real model-derived shapes (currently all empty)
• Benchmarking operators at shapes that match actual LLM inference workloads
• Adding Model Shape coverage for a new or existing operator

Inputs Required

Procedure

1. Look Up Model Architecture Parameters

Retrieve the model's architecture config. Primary source: vLLM model configs (installed at /usr/local/lib/python3.10/dist-packages/vllm) or HuggingFace transformers.

Method A — vLLM Config (Preferred)

vLLM 0.17.1 is installed in the project venv (.venv). Use it to load model configs directly — no internet or weights download required. Always activate the venv first:

source /root/wch/Hopper_benchmark/.venv/bin/activate
# or use the venv python directly:
/root/wch/Hopper_benchmark/.venv/bin/python

For non-standard models (DeepSeek-V3, ChatGLM, Mistral-vLLM, Falcon, etc.) that have custom config classes in vLLM:

# List available custom vLLM config files
ls /root/wch/Hopper_benchmark/.venv/lib/python3.10/site-packages/vllm/transformers_utils/configs/
# Includes: deepseek_v3.py, chatglm.py, mistral.py, falcon.py, nemotron.py, ...

# DeepSeek-V3 example (vLLM custom config)
from vllm.transformers_utils.configs.deepseek_v3 import DeepseekV3Config
c = DeepseekV3Config()
print(c.hidden_size, c.num_attention_heads, c.n_routed_experts,
      c.num_experts_per_tok, c.n_group, c.topk_group)
# → 7168, 128, 256, 8, 8, 4

For standard models (LLaMA, Qwen2, GPT-2, Mistral-HF) that use HuggingFace transformers config classes:

# LLaMA example (HuggingFace config — no weights needed)
from transformers import LlamaConfig
c = LlamaConfig()  # defaults are LLaMA-2-7B
# Override with specific model parameters:
c = LlamaConfig(hidden_size=4096, num_attention_heads=32,
                num_key_value_heads=8, intermediate_size=14336, vocab_size=128256)

# Qwen2 example
from transformers import Qwen2Config
c = Qwen2Config(hidden_size=3584, num_attention_heads=28,
                num_key_value_heads=4, intermediate_size=18944, vocab_size=152064)

# Key field names (standard HuggingFace config):
# c.hidden_size, c.num_attention_heads, c.num_key_value_heads,
# c.intermediate_size, c.vocab_size
# c.head_dim  (= hidden_size // num_attention_heads if not explicit)

Auto-detection script — if you have the model ID but are unsure which class to use:

from transformers import AutoConfig
# Only works if model is already cached locally or HF Hub is accessible:
config = AutoConfig.from_pretrained("<hf_model_id>", trust_remote_code=True)
print(type(config).__name__, vars(config))

Method B — Known Model Registry

Use the table below for well-known models. If the model is not listed, look it up from HuggingFace/vLLM first.

Record all extracted parameters before proceeding to Step 2.

2. Map Operator to Model Component

Each operator appears at specific locations in the transformer forward pass. Use this mapping to determine the shape formula for the given operator.

Operator → Shape Formula Reference

Convention: B = batch size, S = sequence length, H = hidden_size, NH = num_heads, NKV = num_kv_heads, HD = head_dim, V = vocab_size, I = intermediate_size, T = num_tokens (used only for ops with inherently flattened token semantics: embedding, argmax, sort, topk, grouped_topk).

3. Check Operator-Specific Shape Encoding

Critical: Not all operators use simple tensor shapes. Read benchmarks/<op>/config.py to confirm the shape tuple convention:

# Check how the operator's config.py interprets the shape tuple
cat benchmarks/<op>/config.py

Special-encoding operators:

• embedding: (num_tokens, vocab_size, embedding_dim) — 3-tuple, vocab_size is metadata; T = B*S (flattened, no batch dim)
• topk: (num_tokens, V, K) — 3-tuple, K is the k parameter; num_tokens = B*S
• grouped_topk: (num_tokens, num_experts, k, num_expert_group, topk_group) — 5-tuple with MoE routing params
• concat / stack: (num_tensors, B, S, H) — 4-tuple, first dim is tensor count, remaining are per-tensor 3D shape

4. Generate Shape Combinations

For each (batch, seq_len) pair, compute the operator's input tensor shape using the formula from Step 2.

Guidelines:

• Generate shapes for all combinations of batch_sizes × seq_lens
• Use explicit batch dims (B, S, dim) for all standard operators — preserve batch and sequence separately
• Use flattened token count T = B * S only for operators with inherently token-indexed semantics (embedding, argmax, sort, topk, grouped_topk)
• Deduplicate: If two (B, S) pairs produce identical shapes, keep both only if operator semantics differ with batch structure
• Check memory: Ensure each shape's input tensor fits in GPU memory. Skip shapes where product(shape) * element_size(fp32) > 40 GB. For attention-pattern operators like softmax with (B, NH, S, S), the S*S dimension grows quadratically — watch for large S values.
• For MoE operators (grouped_topk), the (batch, seq_len) maps to num_tokens = B * S, and the expert/group params come from the model config.

Example — rms_norm for LLaMA-3-8B:

Model: hidden_size=4096
batch_sizes=[32, 64, 128], seq_lens=[1024, 2048, 4096]

Shapes (B, S, H) — batch and sequence kept separate:
  (32, 1024, 4096)     # B=32, S=1024
  (32, 2048, 4096)     # B=32, S=2048
  (32, 4096, 4096)     # B=32, S=4096
  (64, 1024, 4096)     # B=64, S=1024
  ...deduplicate if shapes repeat across (B,S) combinations

Example — softmax for LLaMA-3-8B (attention scores):

Model: num_heads=32
batch_sizes=[32, 64, 128], seq_lens=[1024, 2048, 4096]

Shapes (B, NH, S, S):
  (32, 32, 1024, 1024)    # ~4 GB FP32 — OK
  (32, 32, 2048, 2048)    # ~16 GB FP32 — OK
  (64, 32, 1024, 1024)    # ~8 GB FP32 — OK
  (64, 32, 2048, 2048)    # ~32 GB FP32 — OK
  (128, 32, 1024, 1024)   # ~16 GB FP32 — OK
  # (32, 32, 4096, 4096)  # ~64 GB FP32 — skip (>40 GB)

Example — grouped_topk for DeepSeek-V3 (MoE routing):

Model: num_experts=256, k=8, num_expert_group=8, topk_group=4
batch_sizes=[32, 64, 128], seq_lens=[1024, 2048, 4096]

Shapes (T, num_experts, k, num_expert_group, topk_group):
  (32768, 256, 8, 8, 4)     # B=32, S=1024
  (65536, 256, 8, 8, 4)     # B=32, S=2048 (or B=64, S=1024)
  (131072, 256, 8, 8, 4)    # B=32, S=4096 (or B=64, S=2048, or B=128, S=1024)
  (262144, 256, 8, 8, 4)    # B=64, S=4096 (or B=128, S=2048)
  (524288, 256, 8, 8, 4)    # B=128, S=4096
  ...

5. Build and Validate the SHAPES List

Assemble the deduplicated shapes into a Python list. Add a comment header indicating the model and parameters used.

"""Model shapes for <op> operator.

Shapes derived from <model_name> (<hf_model_id>):
  hidden_size={H}, num_heads={NH}, num_kv_heads={NKV},
  head_dim={HD}, intermediate_size={I}, vocab_size={V}

Batch sizes: {batch_sizes}
Sequence lengths: {seq_lens}
"""

SHAPES: list[tuple[int, ...]] = [
    # B=1, S=128
    (128, 4096),
    # B=1, S=1024
    (1024, 4096),
    # ... etc
]

Validation checklist before writing:

• Every shape matches the operator's config.py tuple convention (check benchmarks/<op>/config.py)
• No duplicates in the SHAPES list
• No shape exceeds 40 GB when computed as product(shape) * 4 (FP32 element size)
• For operators with special encoding (embedding, topk, grouped_topk), the extra parameters match the model config
• The SHAPES list is not excessively large (aim for 10–30 shapes per model; if combining multiple models, keep total ≤ 50)

6. Write to model_shape/<op>_shape.py

Replace the contents of the existing shape file:

# The file already exists (created when the operator was added)
cat model_shape/<op>_shape.py

Write the generated SHAPES list. Preserve the module-level docstring format.

7. Verify Integration

Run a quick sanity check:

python -c "
from model_shape import load_model_shapes
import torch
cases = load_model_shapes('<op>', torch.float32)
print(f'Loaded {len(cases)} model shapes for <op>:')
for c in cases:
    print(f'  {c.label}: shape={c.shape}, dtype={c.dtype}, category={c.workload_category}')
"

Confirm:

• All shapes load correctly
• workload_category is "Model Shape" for all cases
• Labels are formatted as model_<dim1>x<dim2>x...

8. (Optional) Run Benchmark with Model Shapes

nvidia-smi  # find a free GPU
CUDA_VISIBLE_DEVICES=<free_gpu> python scripts/run_bench.py --op <op> --dtype fp16 fp32

The model shapes will appear in the Model Shape section of the report.

Multiple Models

When generating shapes for multiple models, combine shapes from all models into a single SHAPES list:

"""Model shapes for <op> operator.

Combined shapes from multiple models:
  - LLaMA-3-8B: hidden_size=4096, num_heads=32, ...
  - LLaMA-3-70B: hidden_size=8192, num_heads=64, ...
  - Qwen2-72B: hidden_size=8192, num_heads=64, ...

Batch sizes: [32, 64, 128]
Sequence lengths: [1024, 2048, 4096]
"""

SHAPES: list[tuple[int, ...]] = [
    # --- LLaMA-3-8B ---
    (128, 4096),
    (1024, 4096),
    ...
    # --- LLaMA-3-70B ---
    (128, 8192),
    (1024, 8192),
    ...
]

Deduplicate across models — if two models produce the same shape for an operator, keep only one entry. Add a comment noting which models share that shape.

Operator Applicability

Not all operators appear in all models. Use this quick reference:

If the operator is not applicable to the requested model, inform the user and suggest relevant models instead.

Handling vLLM-Specific Patterns

vLLM uses flattened token representations during inference for most tensor ops. The benchmark convention is:

• Standard ops (rms_norm, add, mul, neg, sigmoid, swi_glu, etc.): use (B, S, H) — preserve batch and sequence separately, even though vLLM flattens at runtime.
• Inherently token-indexed ops (embedding, argmax, sort, topk, grouped_topk): use flattened T = B * S as first dim, matching vLLM semantics. These ops index into vocab or expert tables and don't have a meaningful (B, S) decomposition in context.

PagedAttention in vLLM changes the attention shape:

• Standard: Q shape (T, NH, HD), K/V paged in blocks
• For softmax benchmarking, use the standard (B, NH, S, S) attention score shape — this represents the per-head attention matrix regardless of paging.

Reference Files

Automation script

• scripts/generate_model_shapes.py — Primary tool: generates shapes, deduplicates, validates, writes files

Project files

• model_shape/init.py — load_model_shapes() loader
• model_shape/softmax_shape.py — Example shape file (currently empty)
• benchmarks/softmax/config.py — Shape convention reference
• benchmarks/embedding/config.py — Special encoding example
• benchmarks/grouped_topk/config.py — MoE shape encoding
• common/shape_gen.py — BenchmarkCase and classify_workload

vLLM (v0.17.1, installed in project venv at /root/wch/Hopper_benchmark/.venv/)

• /root/wch/Hopper_benchmark/.venv/lib/python3.10/site-packages/vllm/transformers_utils/configs/ — Custom model config classes for non-standard models
  • deepseek_v3.py → DeepseekV3Config (MoE fields: n_routed_experts, num_experts_per_tok, n_group, topk_group)
  • chatglm.py → ChatGLMConfig
  • mistral.py → vLLM Mistral variant
  • falcon.py → RWConfig (original Falcon-7B/40B)
  • nemotron.py → NemotronConfig
• Standard models (LLaMA, Qwen2, GPT-2) use HuggingFace transformers config classes directly — no vLLM-specific file needed
• venv python: /root/wch/Hopper_benchmark/.venv/bin/python