Support sub-64 MLEN in the emulator (relax 64-byte HBM burst model)

transactional_emulator/src/main.rs

@@ -249,7 +249,11 @@ impl Accelerator {
             assert!(element_bits.is_power_of_two());
 
             let len_in_bits_per_load = element_bits as u32 * load_dim;
-            assert!(len_in_bits_per_load.is_multiple_of(8 * 64));
+            // A load must be a whole number of bytes. (Was a multiple of 8*64
+            // bits, i.e. a full 64-byte HBM burst — relaxed so sub-64 MLEN, which
+            // packs <64 bytes per row, is supported; the element-read loop below
+            // reads aligned 64-byte words and extracts the real bytes.)
+            assert!(len_in_bits_per_load.is_multiple_of(8));
             let len_in_bytes_per_load = len_in_bits_per_load / 8;
 
             // Calculate scale bytes per load iteration (for Mx types)
@@ -297,16 +301,34 @@ impl Accelerator {
                         as usize
                         + block_idx as usize * scale_len_in_bytes_per_load as usize;
 
-                    // Element chunks:
-                    for i in 0..(len_in_bytes_per_load as usize + 63) / 64 {
-                        let chunk_offset = byte_offset + i * 64;
-                        let chunk_size = std::cmp::min(64, total_bytes - chunk_offset);
-                        let addr = element_addr + (i * 64) as u64;
-                        assert!(addr.is_multiple_of(64));
-                        futures.push(Box::pin(async move {
-                            let data = hbm_clone.read(addr).await;
-                            ChunkType::Element(chunk_offset, data, chunk_size)
-                        }));
+                    // Element chunks: read the 64-byte-aligned HBM word(s) that
+                    // cover this load's byte range [element_addr, element_addr+len)
+                    // and copy only the real bytes into the packed `bytes` buffer.
+                    // Handles element_addr that is NOT 64-aligned (sub-64 MLEN, where
+                    // HBM rows are packed tighter than one 64-byte burst) and loads
+                    // that straddle a 64-byte boundary. At MLEN>=64 element_addr is
+                    // 64-aligned and len is a multiple of 64, so this reduces to the
+                    // previous one-aligned-read-per-64-bytes behaviour.
+                    let load_len = len_in_bytes_per_load as u64;
+                    let load_end = element_addr + load_len; // exclusive HBM byte
+                    let mut blk = (element_addr / 64) * 64;
+                    while blk < load_end {
+                        let copy_start = std::cmp::max(blk, element_addr);
+                        let copy_end = std::cmp::min(blk + 64, load_end);
+                        let within = (copy_start - blk) as usize;
+                        let dst = byte_offset + (copy_start - element_addr) as usize;
+                        let mut n = (copy_end - copy_start) as usize;
+                        n = std::cmp::min(n, total_bytes.saturating_sub(dst));
+                        let addr = blk;
+                        if n > 0 {
+                            futures.push(Box::pin(async move {
+                                let data = hbm_clone.read(addr).await;
+                                let mut sel = [0u8; 64];
+                                sel[..n].copy_from_slice(&data[within..within + n]);
+                                ChunkType::Element(dst, sel, n)
+                            }));
+                        }
+                        blk += 64;
                     }
 
                     // Scale chunks (if Mx type)
@@ -456,7 +478,9 @@ impl Accelerator {
         assert!(element_bits.is_power_of_two());
 
         let len_in_bits_per_store = element_bits as u32 * store_dim;
-        assert!(len_in_bits_per_store.is_multiple_of(8 * 64));
+        // Whole-bytes only (was a full 64-byte HBM burst); sub-64 MLEN stores
+        // are read-modify-written into aligned 64-byte words below.
+        assert!(len_in_bits_per_store.is_multiple_of(8));
         let len_in_bytes_per_store = len_in_bits_per_store / 8;
 
         // Calculate scale bytes per store iteration (for Mx types)
@@ -530,20 +554,28 @@ impl Accelerator {
             let element_addr = index + (store_iter * stride) as u64;
             let scale_addr = scale_index + (store_iter as f32 * stride_scale) as u64;
 
-            // Write element bytes to HBM (64-byte aligned chunks)
-            for i in 0..(len_in_bytes_per_store as usize + 63) / 64 {
-                let chunk_offset = i * 64;
-                let chunk_size = std::cmp::min(64, len_in_bytes_per_store as usize - chunk_offset);
-                let addr = element_addr + (i * 64) as u64;
-                assert!(addr.is_multiple_of(64));
-
-                let mut chunk = [0u8; 64];
-                if chunk_offset < element_bytes.len() {
-                    let copy_len = std::cmp::min(chunk_size, element_bytes.len() - chunk_offset);
-                    chunk[..copy_len]
-                        .copy_from_slice(&element_bytes[chunk_offset..chunk_offset + copy_len]);
+            // Write element bytes to HBM. Read-modify-write the 64-byte-aligned
+            // word(s) covering [element_addr, element_addr+len) so sub-64 MLEN
+            // stores (which pack <64 bytes per row and may be unaligned or share a
+            // 64-byte word with the next row) do not clobber neighbouring bytes.
+            // At MLEN>=64 each store fills whole 64-byte words, so the RMW
+            // overwrites all 64 bytes — identical to the previous plain write.
+            let store_len = len_in_bytes_per_store as u64;
+            let store_end = element_addr + store_len;
+            let mut blk = (element_addr / 64) * 64;
+            while blk < store_end {
+                let copy_start = std::cmp::max(blk, element_addr);
+                let copy_end = std::cmp::min(blk + 64, store_end);
+                let within = (copy_start - blk) as usize;
+                let src = (copy_start - element_addr) as usize;
+                let n = (copy_end - copy_start) as usize;
+                let mut chunk = hbm_clone.read(blk).await;
+                if src < element_bytes.len() {
+                    let cn = std::cmp::min(n, element_bytes.len() - src);
+                    chunk[within..within + cn].copy_from_slice(&element_bytes[src..src + cn]);
                 }
-                hbm_clone.write(addr, chunk).await;
+                hbm_clone.write(blk, chunk).await;
+                blk += 64;
             }
 
             // Write scale bytes to HBM (if Mx type)

transactional_emulator/testbench/aten/configurable.py

@@ -98,13 +98,10 @@ def from_args(
         mlen = int(args.mlen)
         vlen = int(args.vlen if args.vlen is not None else mlen)
         blen = int(args.blen)
-        # HBM is modelled as 64-byte bursts; an 8-bit-element MLEN row must fill a
-        # whole number of bursts -> MLEN must be a multiple of 64 (see setup_hw).
-        if mlen % 64 != 0:
-            raise ValueError(
-                f"MLEN ({mlen}) must be a multiple of 64 (emulator models HBM as "
-                f"64-byte bursts). Use MLEN in {{64, 128, 192, 256, ...}}."
-            )
+        # Sub-64 MLEN is supported: the emulator packs HBM rows tightly and reads
+        # aligned 64-byte words with byte-granular extraction (transfer_mx_from_hbm).
+        # The only HBM requirement is a byte-aligned element row, which 8-bit MXFP
+        # always satisfies. (mlen % blen alignment is checked in setup_hw.)
         hlen = int(args.hlen if args.hlen is not None else default_hlen or base["HLEN"])
         broadcast_amount = int(
             args.broadcast_amount
@@ -215,17 +212,10 @@ def setup_hw(args: argparse.Namespace, build_dir: Path) -> HardwareConfig:
         raise ValueError(f"MLEN ({mlen}) must be divisible by BLEN ({blen})")
     if vlen != mlen:
         raise ValueError(f"VLEN ({vlen}) must equal MLEN ({mlen}) for ATen tests")
-    # HBM is modelled as fixed 64-byte (512-bit) bursts. With 8-bit MXFP
-    # elements, one MLEN-wide row must fill a whole number of bursts, i.e.
-    # 8*MLEN must be a multiple of 512 -> MLEN must be a multiple of 64.
-    # Sub-64 MLEN would read <64B per row and trips a Rust panic deep in the
-    # emulator (transfer_mx_from_hbm); fail early here with a clear message.
-    if mlen % 64 != 0:
-        raise ValueError(
-            f"MLEN ({mlen}) must be a multiple of 64: the emulator models HBM as "
-            f"64-byte bursts and an 8-bit-element row narrower than 64 bytes cannot "
-            f"fill one burst. Use MLEN in {{64, 128, 192, 256, ...}}."
-        )
+    # Sub-64 MLEN is supported: the Python HBM writer packs rows tightly and the
+    # emulator reads aligned 64-byte words with byte-granular extraction
+    # (transfer_mx_from_hbm / transfer_mx_to_hbm). 8-bit MXFP rows are always
+    # byte-aligned, so no MLEN%64 constraint is needed.
 
     base = read_behavior_config()
     hlen = args.hlen if args.hlen is not None else base["HLEN"]

transactional_emulator/testbench/sim_env_utils.py

@@ -88,7 +88,6 @@ def _map_mx_byte_aligned(
     if blocks_array.ndim == 1:
         blocks_array = blocks_array.reshape(-1, 1)
     bias_array = np.asarray(bias).reshape(-1)
-    total_bytes_written = 0
 
     with open(output_file, mode) as f:
         for row_idx in range(logical_rows):
@@ -102,7 +101,6 @@ def _map_mx_byte_aligned(
                 raise ValueError(f"Packed element row ({len(row_bytes)} B) > HBM row ({hbm_row_bytes} B)")
             row_bytes += b"\x00" * (hbm_row_bytes - len(row_bytes))
             f.write(row_bytes)
-            total_bytes_written += len(row_bytes)
 
         for row_idx in range(logical_rows):
             row_start = row_idx * blocks_per_source_row
@@ -115,11 +113,18 @@ def _map_mx_byte_aligned(
                 raise ValueError(f"Packed scale row ({len(row_bytes)} B) > scale row ({scale_row_bytes} B)")
             row_bytes += b"\x00" * (scale_row_bytes - len(row_bytes))
             f.write(row_bytes)
-            total_bytes_written += len(row_bytes)
 
-        remainder = total_bytes_written % 64
-        if remainder != 0:
-            f.write(b"\x00" * (64 - remainder))
+        # NOTE: no per-tensor pad to a 64-byte boundary. Each tensor is written at
+        # its compiler-assigned hbm_addr (forward-padded above) and packed at its
+        # exact element+scale byte length, matching the compiler's tight HBM base
+        # allocation (advances by hbm_size, not rounded to 64). A per-tensor tail
+        # pad to 64 would shift the NEXT tensor past its compiler base whenever a
+        # tensor's size is not a 64-multiple (only happens at sub-64 MLEN, e.g.
+        # X(16,16) = 288 B -> padded 320 -> next tensor read at the wrong addr ->
+        # zero weights). The emulator's MemoryBacked capacity (hbm-size, a 64-byte
+        # multiple sized to ~2x the preload) zero-covers the final aligned-block
+        # read past the file end. At MLEN>=64 tensor sizes are already 64-multiples,
+        # so dropping this pad is a no-op.
 
 
 def map_mx_data_to_hbm_for_behave_sim(
@@ -157,9 +162,16 @@ def map_mx_data_to_hbm_for_behave_sim(
     if source_row_elements > logical_row_elements:
         raise ValueError(f"source_row_elements ({source_row_elements}) > logical_row_elements ({logical_row_elements})")
 
-    physical_hbm_row_bytes = (hbm_row_width + 7) // 8
+    # Pack each element row tightly at its real byte width. The compiler
+    # addresses HBM rows at the logical element stride (e.g. MLEN bytes for 8-bit
+    # elements) and places scales at the tight element-region offset, so the
+    # writer must match that stride rather than padding rows up to a 64-byte HBM
+    # burst. The emulator reads aligned 64-byte words and extracts the real bytes
+    # (transfer_mx_from_hbm), so a tight layout is read correctly. At MLEN>=64 the
+    # element row is already >=64 bytes, so this equals the previous value
+    # (max with the 64-byte burst floor was a no-op there).
     element_row_bits = logical_row_elements * element_width
-    hbm_row_bytes = max(physical_hbm_row_bytes, (element_row_bits + 7) // 8)
+    hbm_row_bytes = (element_row_bits + 7) // 8
     blocks_per_source_row = (source_row_elements + block_width - 1) // block_width
     blocks_per_logical_row = (logical_row_elements + block_width - 1) // block_width
     inferred_source_rows = (len(blocks) + blocks_per_source_row - 1) // blocks_per_source_row