From 4a0ab1c22c727d4471ba27a2229624c8d6de4c23 Mon Sep 17 00:00:00 2001 From: booth-algo Date: Thu, 4 Jun 2026 01:22:14 +0100 Subject: [PATCH] feat(frontend): emit per-region hbm_sizes alongside hbm_addrs Expose each HBM region's authoritative byte size (inp.hbm_size) as a parallel name->bytes map in the decoder and vision compile results. Downstream tooling (the simulator's weight-manifest generator) can then size weight / KV / activation regions - including KV stores, which tensor_layouts omits - without recomputing the MXFP8 padding. Additive result key; back-compatible with existing consumers. --- aten/plena_frontend.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/aten/plena_frontend.py b/aten/plena_frontend.py index f091176..85f533e 100644 --- a/aten/plena_frontend.py +++ b/aten/plena_frontend.py @@ -2826,9 +2826,11 @@ def _verbose(message: str = ""): print(f"\nVision compilation complete: {info['isa_lines']} ISA lines, " f"{n_layers} layers, output at VRAM row {o_vram_addr // mlen}") hbm_addrs = {} + hbm_sizes = {} for name, inp in prog._inputs.items(): if hasattr(inp, "hbm_addr"): hbm_addrs[name] = inp.hbm_addr + hbm_sizes[name] = getattr(inp, "hbm_size", None) return { "isa": isa_code, "golden_output": golden_out, @@ -2841,6 +2843,7 @@ def _verbose(message: str = ""): "comparison_params": comparison_params, "info": info, "hbm_addrs": hbm_addrs, + "hbm_sizes": hbm_sizes, "sim_golden_result": { "original_output": padded_golden_output, "tensor_layouts": tensor_layouts, @@ -3571,9 +3574,11 @@ def _verbose(message: str = ""): print(f"\nCompilation complete: {info['isa_lines']} ISA lines, " f"{n_layers} layers, output at VRAM row {o_vram_addr // mlen}") hbm_addrs = {} + hbm_sizes = {} for name, inp in prog._inputs.items(): if hasattr(inp, "hbm_addr"): hbm_addrs[name] = inp.hbm_addr + hbm_sizes[name] = getattr(inp, "hbm_size", None) return { "isa": isa_code, @@ -3588,6 +3593,7 @@ def _verbose(message: str = ""): "info": info, "stage_checkpoints": stage_checkpoint_metadata, "hbm_addrs": hbm_addrs, + "hbm_sizes": hbm_sizes, "sim_golden_result": { "original_output": padded_golden_output, "tensor_layouts": tensor_layouts,