flashinfer.comm._ulysses_lowp.verify_duplicate_scale_slots

flashinfer.comm._ulysses_lowp.verify_duplicate_scale_slots(recv_u8: Tensor, *, batch_size: int, local_sequence: int, local_heads: int, head_dim: int, world_size: int, q_group: int = 32, k_group: int = 64, spec: Dict[str, int | float] | None = None) → bool

Debug/test-only check: every cross-boundary scale slot is bit-identical on all sources that carry it. Runs as a read-only pass after the A2A, so there is no concurrent-writer hazard. Never call on the hot path.

q_group / k_group select the grid the payload was packed with, like the group= argument on boundary_descriptors() and merge_boundary_amax(). They default to the SM89/SM120 constants; pass Q_GROUP_SM90 / K_GROUP_SM90 for an SM90 payload. The grid fixes both the scale byte offsets and which global groups straddle a rank edge, so checking an SM90 payload on the SM120 grid reads the wrong bytes and finds no shared groups at all – the check then passes whatever it is given. spec must come from the matching layout when the grid is not the default one; UlyssesLowpSageLayoutSM90.verify_duplicate_scale_slots() wires all three together.

Parameters:

  • recv_u8 – Contiguous CUDA uint8 source-major buffer [world_size, chunk_bytes] received by all-to-all.

  • batch_size – Positive batch size B.

  • local_sequence – Number of tokens per source shard, including zero padding.

  • local_heads – Heads owned by this destination after the all-to-all, H // world_size.

  • head_dim – Head dimension, 64 or 128; must agree on every rank.

  • world_size – Ulysses group size, one of 2, 4, 6 or 8; 6 is supported structurally but has not been validated.

  • q_group – Q tokens per quantization group: 32 for SM89/SM120, 16 for SM90.

  • k_group – K tokens per quantization group: 64 for SM89/SM120, 128 for SM90.

  • spec – Matching payload geometry. None derives the SM89/SM120 geometry; supply the SM90 layout specification for the SM90 grid.