Skip to content
Phobos

@padstage

Shared memory is split into banks, and two threads that hit the same bank in the same access wait for each other. The bank of an address repeats every 128 bytes, so in a tile whose rows are a multiple of 128 bytes long, the same column of every row lies in the same bank. Reading a column, as a dot or a transpose does, then serializes.

@padstage adds a few unused elements to the end of each row of such a tile. The rows then start in different banks, and a column spreads across them.

@padstage

It takes no arguments. See the Attributes table in the spec.

Which tiles it reaches

Section titled “Which tiles it reaches”

Only staging tiles, and only those whose rows would collide:

  • The tile is a copy of a tensor slice: var t = <tensor slice>. A let view of a tensor that a dot reads more than once is staged by the compiler anyway, and that staging tile counts too.
  • Its row pitch, the last dimension times the element size, is a static multiple of 128 bytes.

Every other tile is unchanged: a tile filled from an expression, an accumulator, or a staging tile whose rows are, say, 64 bytes.

The padding is 8 elements per row. It is invisible to the source: the tile keeps its shape, and only its allocation and its row stride grow.

A minimal kernel

Section titled “A minimal kernel”
@padstage
kernel stage(K: tensor<f16>[64, 128], O: tensor<f16>[8, 128]) {
  var k = K[0 :+ 8, 0 :+ 128]
  O[0 :+ 8, 0 :+ 128] = k
}

k is [8, 128] f16, 256 bytes per row, a multiple of 128. Under @padstage its allocation is [8, 136]. With rows of 32 f16 (64 bytes) instead, the tile would stay as declared.

The pair below is q8_matvec without and with @padstage. Its dot reads a, a let view of the activations, so the compiler stages a into a shared tile of 1 x 32 f32, one row of 128 bytes. That pitch is a multiple of 128 bytes, so under @padstage the stage buffer grows to 160 bytes, and every other buffer keeps its size. Step to the dot on either side to see the stage filled and read.

without @padstage

// Q8_0 matvec: a decode-step projection that reads the weights quantized.
//
// Q8_0 stores signed bytes with one scale per block of 32 values along the
// contraction axis. Reading the bytes directly moves a quarter of the data
// an f32 copy would.
//
// The scale factors out of a block's dot product, so it is applied once per
// block and column, not once per weight.
@autotune(TILE_N in [64], BLK in [32])
@launch(256)
kernel q8_matvec(A: tensor<f32>[1, K],
                 W: tensor<i8>[K, N],
                 S: tensor<f32>[KB, N],
                 C: tensor<f32>[1, N]) {
  let pn = ctx.block.x
  var acc: tile<f32>[1, TILE_N] = 0.0


  for k in range(0, K, BLK) {
    let b = k / BLK
    let a = A[0 :+ 1, k :+ BLK]
    let w = W[k :+ BLK, pn * TILE_N :+ TILE_N]
    let s = S[b :+ 1, pn * TILE_N :+ TILE_N]


    var wf: tile<f32>[BLK, TILE_N] = f32(w)
    var part: tile<f32>[1, TILE_N] = dot(a, wf)
    acc += part * s
  }


  C[0 :+ 1, pn * TILE_N :+ TILE_N] = acc
}

phobos-lang/examples/q8_matvec.phOpen full view

The q8_matvec kernel, stepped through as the compiler built it for sm_75A 2 grid of CTAs with 256 threads each. The drawing shows the tensors A, W, S, C in global memory with the picked CTA's current boxes, shared memory as a byte ruler of 10880 bytes with each buffer at its planned offset, the barriers kept and elided, and how the CTA's threads and warp 0's lanes take part. The readout below describes each step.

CTA 0
  • global memory
  • shared memory
  • threads and lanes
  • extra 128-byte lines, masked elements

with @padstage

// Q8_0 matvec: a decode-step projection that reads the weights quantized.
//
// Q8_0 stores signed bytes with one scale per block of 32 values along the
// contraction axis. Reading the bytes directly moves a quarter of the data
// an f32 copy would.
//
// The scale factors out of a block's dot product, so it is applied once per
// block and column, not once per weight.
@autotune(TILE_N in [64], BLK in [32])
@launch(256)
@padstage
kernel q8_matvec(A: tensor<f32>[1, K],
                 W: tensor<i8>[K, N],
                 S: tensor<f32>[KB, N],
                 C: tensor<f32>[1, N]) {
  let pn = ctx.block.x
  var acc: tile<f32>[1, TILE_N] = 0.0


  for k in range(0, K, BLK) {
    let b = k / BLK
    let a = A[0 :+ 1, k :+ BLK]
    let w = W[k :+ BLK, pn * TILE_N :+ TILE_N]
    let s = S[b :+ 1, pn * TILE_N :+ TILE_N]


    var wf: tile<f32>[BLK, TILE_N] = f32(w)
    var part: tile<f32>[1, TILE_N] = dot(a, wf)
    acc += part * s
  }


  C[0 :+ 1, pn * TILE_N :+ TILE_N] = acc
}

phobos-lang/examples/q8_matvec.ph@padstage addedOpen full view

The q8_matvec kernel with @padstage added, stepped through as the compiler built it for sm_75A 2 grid of CTAs with 256 threads each. The drawing shows the tensors A, W, S, C in global memory with the picked CTA's current boxes, shared memory as a byte ruler of 10880 bytes with each buffer at its planned offset, the barriers kept and elided, and how the CTA's threads and warp 0's lanes take part. The readout below describes each step.

CTA 0
  • global memory
  • shared memory
  • threads and lanes
  • extra 128-byte lines, masked elements

Cost and interactions

Section titled “Cost and interactions”

The padding costs shared memory: 8 * rows elements per padded tile. A kernel at the edge of the 48 KB budget can fail to fit, or lose double buffering under @pipeline, once its stages grow.

A kernel benefits when it reads its staged tiles by column. Attention, which contracts a staged key tile with dot_t, is the typical case: the GGUF backend’s attention kernel carries @padstage. A kernel that only streams rows through elementwise sweeps reads every row whole and gains nothing.