benchmark_mla.py: default --tp-size gives num_local_heads=8, rejected by unified_sm120 prefill (HPB=16)

[SeedSource]

Summary

benchmarks/benchmark_mla.py, run with the GLM-5.1 DSA config and the default --tp-size, derives num_local_heads = 8, which the unified_sm120 prefill kernel rejects (requires heads divisible by HPB=16).

Repro

python benchmarks/benchmark_mla.py     # GLM-5.1 config at /data/models/GLM-5.1-NVFP4, default --tp-size

  File "b12x/attention/mla/api.py", line 892, in _run_unified_sm120_prefill
    _, lse_base2 = run_unified_prefill(...)
  File "b12x/attention/mla/unified_sm120/prefill.py", line 118, in run_unified_prefill
    raise ValueError(
ValueError: unified_sm120 prefill requires heads divisible by HPB=16, got 8

Root cause

num_local_heads = num_attention_heads // tp_size. With the GLM-5.1 config and the default --tp-size, this resolves to 8, but unified_sm120 prefill requires num_local_heads % 16 == 0. The benchmark's default head count is therefore unrunnable on the prefill path out of the box.

Suggested fix

Any of:

• default --tp-size to a value that keeps num_local_heads a multiple of 16 for the bundled config,
• validate num_local_heads % 16 == 0 at arg-parse time with a clear message (currently the error surfaces deep in the kernel),
• or document the HPB=16 head constraint in the benchmark --help / README.

Decode-path cases are unaffected; this is prefill-only.

Env

b12x PR #11 (d90d89c), DGX Spark GB10 / SM121 / aarch64, torch 2.12.1+cu130, nvidia-cutlass-dsl 4.5.2. (Note: benchmark_compressed_mla.py runs all 10 cases correctly on the same box — this is a benchmark-config issue, not a kernel/arch one.)

Filed by an AI agent (SeedSource) during GB10/SM121 bring-up — see #10.