liu-yiqian · Oct 28, 2026 · Sep 29, 2026 · Sep 29, 2026 · Sep 30, 2026 · Sep 30, 2026
@@ -101,6 +101,25 @@ def isMxf4SubtilePath(kernel) -> bool:
         return False
 
 
+def isSubtileSwizzledTensor(kernel) -> bool:
+    """Scope gate for the subtile global read of a host-swizzled A or B.
+
+    Named off SwizzleTensorA/B, the parameters this reads downstream: the scale
+    blocks have their own swizzle (HasMXScaleSwizzle, graTileAssignmentScaleSwizzled)
+    on a separate layout and a separate read, and nothing here applies to it.
+
+    Distinct from isMxf4SubtilePath even though it currently answers the same
+    question. Nothing about reading a pre-swizzled operand is tied to MXF4 or to
+    a bf16 output -- the layout and the K step are functions of the element size
+    and the matrix instruction, so MXFP8 and BF16 would work the same way. They
+    are simply untuned and untested here, and out of scope for this change, so
+    the narrower gate stands in. Widening is then one edit in one place rather
+    than three, and it will not disturb the store and schedule work, which is
+    genuinely MXF4-and-bf16 specific and stays on isMxf4SubtilePath.
+    """
+    return isMxf4SubtilePath(kernel)
+
+
 def plsinBlockSchedTile(kernel) -> bool:
     """Whether block scheduling may be built for this kernel.
 
@@ -565,6 +584,25 @@ def wmmaV3InputVgprLayout(wmma: Sequence[int], dtypeBitWidth: Optional[int] = No
 # Bytes moved by one buffer_load_dwordx4, the widest global load we issue.
 SWIZZLE_LOAD_BYTES = 16
 
+# MUBUF offset12 is a 12-bit unsigned field.  rocisa does not range-check it, so
+# an oversized step would be emitted verbatim rather than rejected.
+SWIZZLE_DTL_OFFSET12_MAX = 4095
+
+
+def swizzleTensorMaxDepthU(rowsPerBlock: int, bpe: float) -> int:
+    """Largest DepthU the read of a swizzled A or B can encode.
+
+    SwizzleTensor, not the MX scale swizzle: the scale blocks have their own
+    layout and their own read, and this ceiling says nothing about them.
+
+    A K step consumes a K-extent on every row of a swizzle block, so the byte
+    step is rowsPerBlock times the unit-stride count, and that product rides in
+    the DTL immediate.  The ceiling is therefore a property of the element size
+    and the matrix instruction, not of a data type: 512 for MXF4's half-byte
+    element over 16 rows, 256 for FP8, 128 for BF16.
+    """
+    return int((SWIZZLE_DTL_OFFSET12_MAX + 1) // rowsPerBlock / bpe)
+
 def swizzleGeometry(solution, tc: str) -> dict:
     """Layout of the pre-swizzled (pre-tiled) tensor `tc` ("A" or "B").
 
@@ -579,20 +617,32 @@ def swizzleGeometry(solution, tc: str) -> dict:
     `solution` may be a partly derived state; only MIInputPerThread{tc}, MatrixInst{M,N,K},
     WavefrontSize and ProblemType.DataType{tc} are read.
     """
-    bpe       = int(solution["ProblemType"][f"DataType{tc}"].numBytes())
+    # Sub-byte types give a fractional bpe (0.5 for F4), which int() truncates to
+    # 0 and makes a divisor below. Keep it fractional and fold it into the
+    # divisor so the division stays exact.
+    bpe       = solution["ProblemType"][f"DataType{tc}"].numBytes()
     miInput   = solution[f"MIInputPerThread{tc}"]
     miMorN    = solution["MatrixInstM"] if tc == "A" else solution["MatrixInstN"]
     # Pack several MI steps into one load when one operand is narrower than a dwordx4.
-    packK     = max(1, SWIZZLE_LOAD_BYTES // miInput // bpe)
+    packK     = max(1, int(SWIZZLE_LOAD_BYTES // (miInput * bpe)))
     miOperand = miInput * packK
-    laneSize  = min(miOperand, SWIZZLE_LOAD_BYTES // bpe)
+    # Elements the host pre-swizzle vectorises over per lane: half a dwordx4.
+    miKv      = max(1, int(SWIZZLE_LOAD_BYTES // 2 // bpe))
+    laneSize  = min(miOperand, int(SWIZZLE_LOAD_BYTES // bpe))
     # Elements the wave holds vs. distinct elements the instruction consumes.
     dupFactor = max(1, (solution["WavefrontSize"] * miInput) // (miMorN * solution["MatrixInstK"]))
     lanesUsed = solution["WavefrontSize"] // dupFactor
     return {
         "packK":        packK,
         "laneSize":     laneSize,
         "swizzleK":     max(1, lanesUsed // miMorN) * laneSize,
+        # K granularity the host pre-swizzle pads to before reshaping.  It packs
+        # SWIZZLE_LOAD_BYTES per lane out of miKv-element vectors, so the step is
+        # miInput * SWIZZLE_LOAD_BYTES / (miKv * bpe) elements.  This is the
+        # pre-shuffled 16-row block stride and is not swizzleK, which covers a
+        # whole wave and is larger; rounding K to swizzleK instead puts every
+        # block after the first at the wrong offset.
+        "layoutK":      int(miInput * SWIZZLE_LOAD_BYTES // (miKv * bpe)),
         "lanesUsed":    lanesUsed,
         "dupFactor":    dupFactor,
         "loadsPerLane": miOperand // laneSize,

@@ -30,7 +30,7 @@
 from .Subtile.SubtileLREmit import localReadResetOffsetsSubtile
 
 from ..Common import IsaVersion, print2, ceilDivide, log2
-from ..Common.Utilities import isMxf4SubtilePath
+from ..Common.Utilities import isMxf4SubtilePath, isSubtileSwizzledTensor
 from ..Component import Component
 from .TileProcessingStrategy import TileProcessingStrategy, TileWork
 from .WorkAssignment import QueuePartition, StaticPartition
@@ -438,6 +438,28 @@ def _depthUForTc(kernel, tc):
                 return kernel[key]
         return kernel["DepthU"]
 
+    @staticmethod
+    def _preShuffleRows(kernel, tc):
+        """Rows a host pre-swizzled block interleaves into one K step, else 1.
+
+        A pre-shuffled tensor stores a whole 16-row block before moving along K,
+        so stepping one DepthU window consumes that many rows' worth of K.  The
+        unroll loop already advances the SRD by the scaled amount; a workgroup
+        whose StreamK range starts mid-tile has to start from the same scale or
+        it reads a K position 16x too close to the tile.
+
+        Only the subtile global read walks blocks this way.  The classic swizzle
+        path steps one row per K window, so scaling its start would move the
+        workgroup off its own range.
+        """
+        if tc not in ("A", "B"):
+            return 1
+        if not kernel["ProblemType"].get("SwizzleTensor%s" % tc):
+            return 1
+        if not isSubtileSwizzledTensor(kernel):
+            return 1
+        return kernel["MatrixInstM"] if tc == "A" else kernel["MatrixInstN"]
+
     def shiftSrd(self, writer, srdIdx) -> Module:
       module = Module("shiftSrd")
       if writer.states.version[:2] == (12, 5):
@@ -908,7 +930,9 @@ def computeLoadSrdCommon(self, writer, kernel, tP, sTmp):
         tc = tP["tensorChar"]
         depthU = self._depthUForTc(kernel, tc)
         # StreamK partial tile - offset to tile start index
-        module.add(SMulI32(dst=sgpr(sTmp), src0=sgpr("StreamKLocalStart"), src1=depthU, comment="StreamK tile start offset"))
+        module.add(SMulI32(dst=sgpr(sTmp), src0=sgpr("StreamKLocalStart"),
+                           src1=depthU * self._preShuffleRows(kernel, tc),
+                           comment="StreamK tile start offset"))
         strideL = self._summationStride(writer, kernel, tc)
         module.add(writer.s_mul_u64_u32(sgpr(sTmp), sgpr(sTmp+1), sgpr(sTmp), strideL, comment="StreamK tile start offset"))
         # Overflow check removed
@@ -991,7 +1015,9 @@ def graAddressesCommon(self, writer, kernel, tP, vTmp):
         depthU = self._depthUForTc(kernel, tc)
         # StreamK partial tile - offset to tile start index
         tmpOffset = writer.sgprPool.checkOut(2, "skStartOffset")
-        module.add(SMulI32(dst=sgpr(tmpOffset), src0=sgpr("StreamKLocalStart"), src1=int(depthU * tP["bpe"]), comment="StreamK tile start offset"))
+        module.add(SMulI32(dst=sgpr(tmpOffset), src0=sgpr("StreamKLocalStart"),
+                           src1=int(depthU * tP["bpe"]) * self._preShuffleRows(kernel, tc),
+                           comment="StreamK tile start offset"))
         strideL = self._summationStride(writer, kernel, tc)
         module.add(writer.s_mul_u64_u32(sgpr(tmpOffset), sgpr(tmpOffset+1), sgpr(tmpOffset), strideL, comment="StreamK tile start offset"))
         # Overflow check removed

@@ -458,6 +458,11 @@ def __init__(self, geometry: TileGeometry, tc: str, writer, kernel):
     isAB = tc in ['A', 'B']
     isMXSAB = tc in ['MXSA', 'MXSB']
     _tc = 'A' if isA else 'B'
+    # SwizzleTensor{A,B} means the host has already transformed the data into
+    # the MFMA-friendly 16-row interleaved layout.  The GR/LR emitters must not
+    # apply the ordinary LDS swizzle to these tensors.
+    self.isPreShuffled = isAB and kernel["ProblemType"].get(
+        "SwizzleTensor%s" % _tc, False)
 
     # --- Extract kernel config ---
     if isinstance(geometry, (ABTilePair, MXScaleTilePair)):
@@ -1314,8 +1319,10 @@ def _emitMultiDUTailSrdRewind(writer, kernel, numUnroll, tiA, tiB, scaleTiA, sca
   """
   module = Module("MultiDU tail SRD rewind (partial macro tile)")
   scaleInc = lambda ti: int(ti.lrSubtileSize * ti.lrGlobalSubtileGrid[1])
-  incs = [("A", int(numUnroll.get('A', 1)) * int(tiA.depthUBytes)),
-          ("B", int(numUnroll.get('B', 1)) * int(tiB.depthUBytes))]
+  dataInc = lambda ti: int(ti.depthUBytes) * (
+      int(ti.mmaTileShape[0]) if ti.isPreShuffled else 1)
+  incs = [("A", int(numUnroll.get('A', 1)) * dataInc(tiA)),
+          ("B", int(numUnroll.get('B', 1)) * dataInc(tiB))]
   if scaleTiA is not None:
     incs.append(("MXSA", int(numUnroll.get('SA', 1)) * scaleInc(scaleTiA)))
   if scaleTiB is not None:

@@ -18,6 +18,7 @@
 
 import contextlib
 from ...ExecutionPolicy import isPersistentDataParallel
+from ...Common.Utilities import SWIZZLE_DTL_OFFSET12_MAX, swizzleTensorMaxDepthU
 import math
 from functools import singledispatch
 
@@ -313,6 +314,8 @@ def _emitGR_TLU0(tag, tile, ti, writer, kernel):
 
       rl = ti.localSubtilesRegister[min(reg_idx, perpDimSize - 1)]
       offsetK = j * int(ti.mmaTileShape[1] * ti.subtileShape[1] * ti.bpe)
+      if ti.isPreShuffled:
+        offsetK = _preShuffledKStepBytes(ti, offsetK)
 
       module.addComment0(f"GR load {tc} subtile [{i},{j}]")
 
@@ -383,7 +386,10 @@ def _emitGRPtrUpdate_TLU0(tag, tile, ti, writer, kernel, holdOnLastIter=False):
     return module
 
   module = Module(f"GR Ptr Update ({tc})")
-  inc = int(ti.depthUBytes)
+  # Host-pre-shuffled data interleaves one K tile across all rows of an MFMA
+  # tile, so advancing one DepthU consumes 16 contiguous row fragments.
+  rowMultiplier = ti.mmaTileShape[0] if ti.isPreShuffled else 1
+  inc = int(ti.depthUBytes * rowMultiplier)
   emitSrdAdvance(module, tc, inc, writer, kernel, holdOnLastIter)
   return module
 
@@ -564,17 +570,52 @@ def graTileAssignment(writer, kernel, useSwizzling=True):
 
 # --- Legacy interleaved A+B GR offset (temporary, matches reference exactly) ---
 
+def _preShuffledKStepBytes(tileInfo, offsetK):
+  """Scale a K-subtile step for pre-shuffled data, refusing steps that cannot be encoded.
+
+  A K step consumes a K-extent on every row of the block, so the global step is
+  ``mmaTileShape[0]`` times the unit-stride byte count.  That product rides in the
+  DTL immediate, so it has a hard ceiling; past it the step has to move into
+  soffset the way the TLU=1 K-window path already does.
+  """
+  rowsPerBlock = int(tileInfo.mmaTileShape[0])
+  depthU = int(tileInfo.depthUBytes / tileInfo.bpe)
+  maxDepthU = swizzleTensorMaxDepthU(rowsPerBlock, tileInfo.bpe)
+  if depthU > maxDepthU:
+    raise RuntimeError(
+        "SwizzleTensor%s with DepthU %u is not supported: the pre-shuffled global "
+        "read is only validated up to DepthU %u.  Emitting it would read the wrong "
+        "K positions instead of failing."
+        % (tileInfo.tc, depthU, maxDepthU))
+  scaled = offsetK * rowsPerBlock
+  if scaled > SWIZZLE_DTL_OFFSET12_MAX:
+    raise RuntimeError(
+        "SwizzleTensor%s: pre-shuffled K-window step is %u bytes (DepthU %u, %u "
+        "rows per block), which does not fit the %u byte DTL immediate.  Fold the "
+        "step into soffset instead."
+        % (tileInfo.tc, scaled, depthU, rowsPerBlock, SWIZZLE_DTL_OFFSET12_MAX))
+  return scaled
+
+
 def _grComputeOffset_legacy(module, writer, tileInfo, colId, rowId, output):
   tc = tileInfo.tc
   bpeBits = int(8*tileInfo.bpe)
   tmpVgpr = writer.vgprPool.checkOut(2, tag="_grComputeOffset_legacy_tmpVgpr")
   colBytes = tmpVgpr + 1
   loadWidth = tileInfo.loadWidthGR
   module.add(VLShiftLeftB32(dst=vgpr(colBytes), shiftHex=hex(loadWidth.bit_length()-1), src=vgpr(colId), comment="scale col_id by load_width"))
-  strideRef = "StrideA0I" if tc == 'A' else "StrideB1J"
-  module.add(VMulLOU32(dst=vgpr(tmpVgpr), src0=sgpr(strideRef), src1=vgpr(rowId), comment="%s: rowId * stride"%tc))
-  module.add(VLShiftLeftB32(dst=vgpr(tmpVgpr), shiftHex=hex(bpeBits.bit_length()-1), src=vgpr(tmpVgpr), comment="%s: rowId*stride*bpe"%tc))
-  module.add(VLShiftRightB32(dst=vgpr(tmpVgpr), shiftHex=hex(3), src=vgpr(tmpVgpr), comment="to bytes"))
+  if tileInfo.isPreShuffled:
+    # The GR enumeration covers one K-subtile at a time, so a row step is that
+    # subtile's K-extent, not the whole DepthU window.  They are equal only when
+    # localSubtileGrid[1] == 1.
+    subIterKBytes = tileInfo.subIterKBytes
+    module.add(VLShiftLeftB32(dst=vgpr(tmpVgpr), shiftHex=hex(subIterKBytes.bit_length()-1),
+                              src=vgpr(rowId), comment="%s: rowId * subIterKBytes"%tc))
+  else:
+    strideRef = "StrideA0I" if tc == 'A' else "StrideB1J"
+    module.add(VMulLOU32(dst=vgpr(tmpVgpr), src0=sgpr(strideRef), src1=vgpr(rowId), comment="%s: rowId * stride"%tc))
+    module.add(VLShiftLeftB32(dst=vgpr(tmpVgpr), shiftHex=hex(bpeBits.bit_length()-1), src=vgpr(tmpVgpr), comment="%s: rowId*stride*bpe"%tc))
+    module.add(VLShiftRightB32(dst=vgpr(tmpVgpr), shiftHex=hex(3), src=vgpr(tmpVgpr), comment="to bytes"))
   module.add(VAddU32(dst=vgpr(output), src0=vgpr(colBytes), src1=vgpr(tmpVgpr), comment="%s: GR row_offset"%tc))
   writer.vgprPool.checkIn(tmpVgpr)
 
@@ -633,7 +674,7 @@ def _grComputeAllOffsets_legacy(module, writer, tileInfo, colId, rowId, rowOffse
     module.add(VAddU32(dst=vgpr(rowOffset), src0=offset, src1=vgpr(rowOffset), comment="%s: advance row for GR offset %u"%(tileInfo.tc, i)))
     rotatedcolId = writer.vgprPool.checkOut(1, tag="_grComputeAllOffsets_legacy_rotatedcolId")
     loadWidth = tileInfo.loadWidthGR
-    if tileInfo.loadRatioGR == 0.5:
+    if tileInfo.loadRatioGR == 0.5 and not tileInfo.isPreShuffled:
       if tileInfo.bpe == 1:  # FP8: intra-block K_group +2 rotation, preserving block bit
         tmpBlock = writer.vgprPool.checkOut(1, tag="_grComputeAllOffsets_legacy_tmpBlock")
         module.add(VAndB32(dst=vgpr(tmpBlock), src0=vgpr(colId), src1=hex(4), comment="%s: block_bit = colId & 4"%tileInfo.tc))
@@ -782,13 +823,27 @@ def _graTileAssignment_legacy(writer, kernel, useSwizzling=True):
   module.add(VLShiftRightB32(dst=vgpr(waveId), shiftHex=hex(wavesize.bit_length()-1), src=vgpr("Serial"), comment="Wave Id"))
   module.add(VAndB32(dst=vgpr(laneId), src0=vgpr("Serial"), src1=wavesize-1, comment=""))
   module.add(VAndB32(dst=vgpr(colIdA), src0=vgpr("Serial"), src1=(blockSize-1), comment="get col_id in wave for %uB load"%loadWidth))
+  # Without a pre-shuffled tensor the swizzle derives colIdB from colIdA, so only
+  # seed it separately when one of the two skips the swizzle.
+  if tileInfoA.isPreShuffled or tileInfoB.isPreShuffled:
+    module.add(VAndB32(dst=vgpr(colIdB), src0=vgpr("Serial"), src1=(blockSize-1), comment="colIdB base"))
   module.add(VLShiftRightB32(dst=vgpr(rowId), shiftHex=hex(blockSize.bit_length()-1), src=vgpr(laneId), comment="row id within wave"))
-  _grSwizzleColIds(module, writer, [(tileInfoA, colIdA), (tileInfoB, colIdB)],
-                   blockSize, numRowsPerLDSBanks, laneId, waveId)
+  # A pre-shuffled tensor already carries the host's layout, so it must not be
+  # swizzled again; pass only the tensors that still need it.
+  swizzlePairs = [(t, c) for t, c in ((tileInfoA, colIdA), (tileInfoB, colIdB))
+                  if not t.isPreShuffled]
+  if swizzlePairs:
+    _grSwizzleColIds(module, writer, swizzlePairs,
+                     blockSize, numRowsPerLDSBanks, laneId, waveId)
   _grComputeRowPartition_legacy(module, kernel, writer, tileInfoA, waveId, rowOffsetA)
   _grComputeRowPartition_legacy(module, kernel, writer, tileInfoB, waveId, rowOffsetB)
   _grComputeAllOffsets_legacy(module, writer, tileInfoA, colIdA, rowId, rowOffsetA)
   _grComputeAllOffsets_legacy(module, writer, tileInfoB, colIdB, rowId, rowOffsetB)
+
+  for tile, rowOff in ((tileInfoA, rowOffsetA), (tileInfoB, rowOffsetB)):
+    if tile.isPreShuffled:
+      _addPreShuffleCorrectionToGROffsets(module, writer, kernel, tile, rowOff)
+
   writer.vgprPool.checkIn(tmpVgpr)
   _grComputeSubtileOffsets_legacy(writer, module, tileInfoA)
   _grComputeSubtileOffsets_legacy(writer, module, tileInfoB)
@@ -1371,6 +1426,42 @@ def _grComputeSubtileOffsets_tlu(writer, module, tileInfo):
         module.add(VAddU32(dst=vgpr(reg), src0=vgpr(tileInfo.sharedVgprGROffset[i]),
                    src1=hex(off), comment="%s: subtile row %u offset (vgpr)" % (tc, regId)))
   return module
+def _addPreShuffleCorrectionToGROffsets(module, writer, kernel, tile, rowOffset):
+  """Fold aligned 16-row HBM padding into pre-shuffled GR offsets."""
+  tc = tile.tc
+  strideRef = "StrideA0I" if tc == 'A' else "StrideB1J"
+  depthU = int(tile.subIterKBytes / tile.bpe)
+  nTileRows = tile.mmaTileShape[0]
+  corrVgpr = writer.vgprPool.checkOut(
+      1, tag="_addPreShuffleCorrectionToGROffsets_corrVgpr")
+  tmpSgpr = writer.sgprPool.checkOut(
+      1, tag="_addPreShuffleCorrectionToGROffsets_tmpSgpr", preventOverflow=False)
+
+  module.add(VAndB32(dst=vgpr(corrVgpr),
+                     src0=hex(~(nTileRows - 1) & 0xFFFFFFFF),
+                     src1=vgpr(rowOffset), comment="aligned pre-shuffled row"))
+  if tile.bpe == 0.5:
+    # Halve the stride rather than the product: a K shorter than one DepthU
+    # window makes this step legitimately negative, and shifting the product
+    # right is a logical shift, which would turn it into roughly 2**31.
+    module.add(SLShiftRightB32(dst=sgpr(tmpSgpr), shiftHex=hex(1),
+                               src=sgpr(strideRef), comment="FP4 elements to bytes"))
+    module.add(SSubU32(dst=sgpr(tmpSgpr), src0=sgpr(tmpSgpr), src1=hex(depthU // 2),
+                       comment="stride - DepthU"))
+  else:
+    module.add(SSubU32(dst=sgpr(tmpSgpr), src0=sgpr(strideRef), src1=hex(depthU),
+                       comment="stride - DepthU"))
+  module.add(VMulLOU32(dst=vgpr(corrVgpr), src0=sgpr(tmpSgpr),
+                       src1=vgpr(corrVgpr), comment="row padding correction"))
+
+  for i, grOffset in enumerate(tile.sharedVgprGROffset):
+    module.add(VAddU32(dst=vgpr(grOffset), src0=vgpr(grOffset),
+                       src1=vgpr(corrVgpr),
+                       comment="pre-shuffled %s GR[%u] correction" % (tc, i)))
+  writer.vgprPool.checkIn(corrVgpr)
+  writer.sgprPool.checkIn(tmpSgpr)
+
+
 ##################################################
 # Subroutine to generate GR load code
 #
@@ -1427,6 +1518,8 @@ def emitSingleBufferLoad(tileInfo, kernel, sId0, sId1, writer=None):
   useSgpr = regList.is_sgpr
 
   offsetK = sId1 * int(tileInfo.mmaTileShape[1] * tileInfo.subtileShape[1] * tileInfo.bpe)
+  if tileInfo.isPreShuffled:
+    offsetK = _preShuffledKStepBytes(tileInfo, offsetK)
 
   subtileOffset = int(math.ceil(tileInfo.loadRatioGR*tileInfo.subtileSize))
   # loadRatioGR folds in grLoadWaves, so subtileOffset is the bytes one