refactor: drop num_moe_layers/top_k from r3/v1 deserializer

Mirrors the gateway-side r3_serializer change: the per-token matrix
shape (num_moe_layers, top_k) is no longer required and is no longer
written into the r3/v1 binary header.  Per-token matrix byte size is
recovered as matrix_byte_length / replayed_token_count.

- HEADER_FORMAT: "<4sBBBBIIHHIIQ" (36 bytes) -> "<4sBBBBIIIIQ" (32 bytes).
- Drop num_moe_layers/top_k from _parse_header() and the metadata dict
  returned by decompress_and_parse_r3().
- Compute matrix_elem_size from matrix_byte_length / replayed_token_count
  with a divisibility check that surfaces malformed payloads early.
- Update unit tests to use matrix_elem_size as the parameter and drop
  assertions on the removed header fields; round-trip test no longer
  passes num_moe_layers/top_k to RouterReplayData.

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: SunnySoldier357

eval_protocol/adapters/r3_deserializer.py

@@ -15,7 +15,6 @@
 
 import base64
 import logging
-import math
 import struct
 from enum import IntEnum
 from typing import Any, Dict, List, Optional, Tuple
@@ -25,8 +24,8 @@
 logger = logging.getLogger(__name__)
 
 MAGIC = b"R3V1"
-HEADER_FORMAT = "<4sBBBBIIHHIIQ"
-HEADER_SIZE = struct.calcsize(HEADER_FORMAT)  # 36 bytes
+HEADER_FORMAT = "<4sBBBBIIIIQ"
+HEADER_SIZE = struct.calcsize(HEADER_FORMAT)  # 32 bytes
 
 
 class _SelectorMode(IntEnum):
@@ -62,8 +61,6 @@ def _parse_header(raw: bytes) -> Dict[str, Any]:
         flags,
         total_token_count,
         replayed_token_count,
-        num_moe_layers,
-        top_k,
         replay_start_token,
         selector_byte_length,
         matrix_byte_length,
@@ -80,8 +77,6 @@ def _parse_header(raw: bytes) -> Dict[str, Any]:
         "flags": flags,
         "total_token_count": total_token_count,
         "replayed_token_count": replayed_token_count,
-        "num_moe_layers": num_moe_layers,
-        "top_k": top_k,
         "replay_start_token": replay_start_token,
         "selector_byte_length": selector_byte_length,
         "matrix_byte_length": matrix_byte_length,
@@ -117,9 +112,9 @@ def decompress_and_parse_r3(
           ``total_token_count``.  Each present position contains a
           base64-encoded routing matrix (matching the format returned by
           the direct inference path); absent positions are ``None``.
-        - ``metadata`` is a dict with keys ``num_moe_layers``, ``top_k``,
-          ``routing_dtype``, ``selector_mode``, ``total_token_count``,
-          ``replayed_token_count``, ``replay_start_token``.
+        - ``metadata`` is a dict with keys ``routing_dtype``,
+          ``selector_mode``, ``total_token_count``, ``replayed_token_count``,
+          ``replay_start_token``.
     """
     compressed = base64.b64decode(data_b64)
 
@@ -132,14 +127,23 @@ def decompress_and_parse_r3(
     routing_dtype = header["routing_dtype"]
     total_token_count = header["total_token_count"]
     replayed_token_count = header["replayed_token_count"]
-    num_moe_layers = header["num_moe_layers"]
-    top_k = header["top_k"]
     replay_start_token = header["replay_start_token"]
     selector_byte_length = header["selector_byte_length"]
     matrix_byte_length = header["matrix_byte_length"]
 
-    dtype_byte_width = _RoutingDtype(routing_dtype).byte_width
-    matrix_elem_size = num_moe_layers * top_k * dtype_byte_width
+    # Per-token matrix byte size is implicit in the payload: all replayed
+    # tokens share the same matrix length, so we can recover it from the
+    # matrix section total length divided by the replayed-token count.
+    if replayed_token_count > 0:
+        if matrix_byte_length % replayed_token_count != 0:
+            raise ValueError(
+                f"matrix_byte_length ({matrix_byte_length}) is not a multiple of "
+                f"replayed_token_count ({replayed_token_count}); cannot split "
+                "into per-token matrices"
+            )
+        matrix_elem_size = matrix_byte_length // replayed_token_count
+    else:
+        matrix_elem_size = 0
 
     body = raw[HEADER_SIZE:]
     selector_bytes = body[:selector_byte_length]
@@ -173,8 +177,6 @@ def decompress_and_parse_r3(
         matrices[pos] = base64.b64encode(matrix_bytes[start:end]).decode("ascii")
 
     metadata: Dict[str, Any] = {
-        "num_moe_layers": num_moe_layers,
-        "top_k": top_k,
         "routing_dtype": _ROUTING_DTYPE_NAMES.get(
             _RoutingDtype(routing_dtype), str(routing_dtype)
         ),

tests/adapters/test_r3_deserializer.py

@@ -28,17 +28,20 @@ def _make_raw_r3(
     routing_dtype: int = _RoutingDtype.UINT8,
     total_token_count: int = 4,
     replayed_token_count: int = 4,
-    num_moe_layers: int = 2,
-    top_k: int = 2,
+    matrix_elem_size: Optional[int] = None,
     replay_start_token: int = 0,
     selector_bytes: bytes = b"",
     matrix_data: Optional[bytes] = None,
 ) -> bytes:
-    """Build a raw (uncompressed) R3/v1 payload for testing."""
-    dtype_byte_width = _RoutingDtype(routing_dtype).byte_width
-    matrix_elem_size = num_moe_layers * top_k * dtype_byte_width
+    """Build a raw (uncompressed) R3/v1 payload for testing.
 
+    ``matrix_elem_size`` is the per-token matrix byte length; when not given
+    and no explicit ``matrix_data`` is supplied, defaults to 4 bytes/token
+    (a minimal placeholder for tests that don't care about shape).
+    """
     if matrix_data is None:
+        if matrix_elem_size is None:
+            matrix_elem_size = 4
         matrix_data = bytes(range(matrix_elem_size)) * replayed_token_count
 
     header = struct.pack(
@@ -50,8 +53,6 @@ def _make_raw_r3(
         0x01,  # flags: little-endian
         total_token_count,
         replayed_token_count,
-        num_moe_layers,
-        top_k,
         replay_start_token,
         len(selector_bytes),
         len(matrix_data),
@@ -86,7 +87,7 @@ def test_too_short(self):
     def test_unsupported_version(self):
         raw = struct.pack(
             HEADER_FORMAT,
-            MAGIC, 99, 0, 1, 0, 4, 4, 2, 2, 0, 0, 16,
+            MAGIC, 99, 0, 1, 0, 4, 4, 0, 0, 16,
         )
         with pytest.raises(ValueError, match="Unsupported R3 header version"):
             _parse_header(raw)
@@ -118,10 +119,8 @@ def test_multi_byte(self):
 
 class TestDecompressAndParseR3:
     def test_all_mode_uint8(self):
-        num_moe_layers = 2
-        top_k = 2
+        matrix_elem_size = 4  # e.g. 2 MoE layers * 2 top-k * 1 byte (uint8)
         total_tokens = 4
-        matrix_elem_size = num_moe_layers * top_k  # 4 bytes per token
 
         matrices_raw = []
         for i in range(total_tokens):
@@ -131,17 +130,13 @@ def test_all_mode_uint8(self):
         raw = _make_raw_r3(
             total_token_count=total_tokens,
             replayed_token_count=total_tokens,
-            num_moe_layers=num_moe_layers,
-            top_k=top_k,
             matrix_data=matrix_data,
         )
         blob = _compress_and_b64(raw)
 
         matrices, metadata = decompress_and_parse_r3(blob)
 
         assert len(matrices) == total_tokens
-        assert metadata["num_moe_layers"] == num_moe_layers
-        assert metadata["top_k"] == top_k
         assert metadata["routing_dtype"] == "uint8"
         assert metadata["selector_mode"] == "all"
         assert metadata["total_token_count"] == total_tokens
@@ -153,12 +148,10 @@ def test_all_mode_uint8(self):
             assert decoded == matrices_raw[i]
 
     def test_suffix_mode(self):
-        num_moe_layers = 2
-        top_k = 2
+        matrix_elem_size = 4
         total_tokens = 8
         replayed = 3
         start_token = 5
-        matrix_elem_size = num_moe_layers * top_k
 
         matrices_raw = []
         for i in range(replayed):
@@ -169,8 +162,6 @@ def test_suffix_mode(self):
             selector_mode=_SelectorMode.SUFFIX,
             total_token_count=total_tokens,
             replayed_token_count=replayed,
-            num_moe_layers=num_moe_layers,
-            top_k=top_k,
             replay_start_token=start_token,
             matrix_data=matrix_data,
         )
@@ -194,10 +185,8 @@ def test_suffix_mode(self):
             assert decoded == matrices_raw[i]
 
     def test_bitmap_mode(self):
-        num_moe_layers = 2
-        top_k = 2
+        matrix_elem_size = 4
         total_tokens = 8
-        matrix_elem_size = num_moe_layers * top_k
 
         # Replay tokens at positions 1, 3, 6
         replayed_positions = [1, 3, 6]
@@ -218,8 +207,6 @@ def test_bitmap_mode(self):
             selector_mode=_SelectorMode.BITMAP,
             total_token_count=total_tokens,
             replayed_token_count=replayed,
-            num_moe_layers=num_moe_layers,
-            top_k=top_k,
             selector_bytes=selector_bytes,
             matrix_data=matrix_data,
         )
@@ -241,10 +228,8 @@ def test_bitmap_mode(self):
                 assert matrices[i] is None
 
     def test_uint16_dtype(self):
-        num_moe_layers = 2
-        top_k = 2
+        matrix_elem_size = 8  # e.g. 2 MoE layers * 2 top-k * 2 bytes (uint16)
         total_tokens = 2
-        matrix_elem_size = num_moe_layers * top_k * 2  # 2 bytes per element for uint16
 
         matrices_raw = []
         for i in range(total_tokens):
@@ -255,8 +240,6 @@ def test_uint16_dtype(self):
             routing_dtype=_RoutingDtype.UINT16,
             total_token_count=total_tokens,
             replayed_token_count=total_tokens,
-            num_moe_layers=num_moe_layers,
-            top_k=top_k,
             matrix_data=matrix_data,
         )
         blob = _compress_and_b64(raw)
@@ -336,8 +319,6 @@ def test_round_trip_with_serializer(self):
         data = RouterReplayData(
             routing_matrices=original_matrices,
             total_token_count=total_tokens,
-            num_moe_layers=num_moe_layers,
-            top_k=top_k,
             routing_dtype="uint8",
         )
 
@@ -350,8 +331,7 @@ def test_round_trip_with_serializer(self):
         matrices, metadata = decompress_and_parse_r3(blob_b64)
 
         assert len(matrices) == total_tokens
-        assert metadata["num_moe_layers"] == num_moe_layers
-        assert metadata["top_k"] == top_k
+        assert metadata["total_token_count"] == total_tokens
 
         for i in range(total_tokens):
             if original_b64[i] is None:
@@ -367,10 +347,8 @@ class TestConvertTraceDictWithPayloads:
     def test_trace_with_router_replay_payload(self):
         from eval_protocol.adapters.fireworks_tracing import convert_trace_dict_to_evaluation_row
 
-        num_moe_layers = 2
-        top_k = 2
+        matrix_elem_size = 4
         total_tokens = 4
-        matrix_elem_size = num_moe_layers * top_k
 
         matrices_raw = []
         for i in range(total_tokens):
@@ -380,8 +358,6 @@ def test_trace_with_router_replay_payload(self):
         raw = _make_raw_r3(
             total_token_count=total_tokens,
             replayed_token_count=total_tokens,
-            num_moe_layers=num_moe_layers,
-            top_k=top_k,
             matrix_data=matrix_data,
         )
         blob = _compress_and_b64(raw)
@@ -424,8 +400,8 @@ def test_trace_with_router_replay_payload(self):
             assert decoded == matrices_raw[i]
 
         meta = row.execution_metadata.extra["routing_metadata"]
-        assert meta["num_moe_layers"] == num_moe_layers
-        assert meta["top_k"] == top_k
+        assert meta["routing_dtype"] == "uint8"
+        assert meta["total_token_count"] == total_tokens
 
     def test_trace_without_payloads(self):
         from eval_protocol.adapters.fireworks_tracing import convert_trace_dict_to_evaluation_row