[BUG] Unsupported: Spyre backend does not support: buf0: per-core tensor span >256M

[thanh-lam]

These failures occurred when running test_benchmark.py with Batch size = 8. Only the spyre tests that failed. The same test that ran with device=cpu didn't fail. Also, the same spyre test that ran with Batch size = 1 didn't fail.

** Detailed error messages: **

test_benchmark.py::test_cmdopt[gemma4_base-08mbs-256len-128tok-float16-spyre-1aiu] Run spyre shape cmd: ['/usr/bin/python3', '/tmp/istbench/torch-spyre-2.0-tests/test_benchmark.py', '--shape-worker', '--worker-model', 'gemma4_base', '--worker-bs', '8', '--worker-seqlen', '256', '--worker-maxtok', '128', '--worker-dtype', 'float16', '--worker-output', 'benchmark_results/20260901/worker_gemma4_base_spyre_1aiu_8_256_128_float16.json', '--device', 'spyre', '--num-aius', '1']
--- worker stdout ---
[ShapeWorker rank=0] Gemma 4 12B Base bs=8 seqlen=256 max_tokens=128 dtype=float16 num_aius=1 validate=True mode=generate
Model on Spyre ready.
  [DEBUG] Starting generate() at 23:50:52
[ShapeWorker rank=0]   Error: Unsupported: Spyre backend does not support: buf0: per-core tensor span 320.000 MB (shape=[8, 320, 262144], dtype=torch.float16, device_size=[320, 4096, 8, 64], splits={d0: 1, d1: 4, d2: 8, d3: 1}) exceeds hardware limit of 256.00 MB
  
Set TORCHDYNAMO_VERBOSE=1 for the internal stack trace (please do this especially if you're reporting a bug to PyTorch). For even more developer context, set TORCH_LOGS="+dynamo"

Outputs from a PASSED test:

test_benchmark.py::test_cmdopt[gemma4_base-08mbs-256len-128tok-float16-cpu-1aiu] [ShapeWorker rank=0] Gemma 4 12B Base bs=8 seqlen=256 max_tokens=128 dtype=float16 num_aius=1 validate=True mode=generate
^MLoading weights:   0%|          | 0/677 [00:00<?, ?it/s]^MLoading weights:  22%|██▏       | 146/677 [00:00<00:00, 1447.44it/s]^MLoading weights:  43%|████▎     | 291/677 [00:00<00:00, 1411.53it/s]^MLoading weights:  64%|██████▍   | 433/677 [00:00<00:00, 1412.22it/s]^MLoading weights:  85%|████████▍ | 575/677 [00:00<00:00, 1363.34it/s]^MLoading weights: 100%|██████████| 677/677 [00:00<00:00, 1413.82it/s]
[ShapeWorker rank=0]   41.005s | 24.97 tok/s
OK Gemma 4 12B Base bs=8 seq=256 tok=128 -> 24.97 tok/s
PASSED

Failed tests included:

• gemma4_base
• gemma4_31b

These tests were running in a 1-AIU pod with latest spyre:dev container image:

• icr.io/ibmaiu_internal/2.0/x86_64/spyre:dev

Latest levels of torch-spyre and hf-adapters on Sept 01.