{
  "domain": "com.microsoft",
  "name": "GatherBlockQuantized",
  "sinceVersion": 1,
  "description": "Gathers rows from a block-wise quantized weight matrix and dequantizes them. This inference implementation supports the standard `gather_axis = 0`, `quantize_axis = 1` matrix subset with uint8 `data`, 4-bit packed or 8-bit values, rank-1 non-negative in-bounds int64 `indices` projected to uint32 WebGPU storage, and float32 scales/output. Higher-rank gathers, negative indices, int32 indices, int4/uint4 data, 2-bit data, float16/bfloat16 output, and non-default axes are not implemented.",
  "inputs": [
    {
      "role": "data",
      "dtype": "T1",
      "rank": 2,
      "description": "Constant uint8 weight matrix. With `bits = 4`, each byte stores two values low-nibble first; with `bits = 8`, each byte stores one value."
    },
    {
      "role": "indices",
      "dtype": "Tind",
      "rank": 1,
      "description": "Non-negative logical int64 indices selecting rows from axis 0 of `data`. Every index must be less than the row count; values use checked uint32 WebGPU storage."
    },
    {
      "role": "scales",
      "dtype": "T2",
      "rank": 2,
      "description": "Per-block dequantization scale factors of shape `(rows, ceil(output_columns / block_size))`."
    },
    {
      "role": "zero_points",
      "dtype": "T1",
      "rank": 2,
      "optional": true,
      "description": "Optional uint8 zero points. At 4 bits two zero points are packed per byte along the quantized axis, low-nibble first; at 8 bits the shape matches `scales`. If absent, uint8 data uses 2^(bits-1)."
    }
  ],
  "outputs": [
    {
      "role": "output",
      "dtype": "T2",
      "rank": 2,
      "shape": "[dim(shapes.indices, 0), dim(shapes.data, 1) * (8 / attrs.bits)]",
      "description": "Dequantized floating-point output rows corresponding to the gathered indices."
    }
  ],
  "attributes": { "bits": 4, "block_size": 128, "gather_axis": 0, "quantize_axis": 1 },
  "attributeConstraints": {
    "bits": { "values": [4, 8] },
    "gather_axis": { "values": [0] },
    "quantize_axis": { "values": [1] }
  },
  "attributeDescriptions": {
    "bits": "Bits per quantized value. The schema default is 4; this implementation supports 4 or 8.",
    "block_size": "Number of values sharing a scale. Defaults to 128 and must be a power of two at least 16.",
    "gather_axis": "Axis from which values are gathered. This matrix implementation supports the standard default, axis 0.",
    "quantize_axis": "Axis split into quantization blocks. This matrix implementation supports the standard default, axis 1."
  },
  "typeConstraints": { "T1": ["uint8"], "T2": ["float32"], "Tind": ["int64"] },
  "args": {
    "dataT": { "kind": "tensor", "semantic": "data", "role": "input" },
    "indicesT": { "kind": "tensor", "semantic": "indices", "role": "input", "dtype": "uint32", "narrowing": "checked" },
    "scalesT": { "kind": "tensor", "semantic": "scales", "role": "input" },
    "zeroPointsT": { "kind": "tensor", "semantic": "zero_points", "role": "input", "required": false },
    "outputT": { "kind": "tensor", "semantic": "output", "role": "output" }
  },
  "derive": {
    "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
    "bits": "attrs.bits",
    "blockSize": "attrs.block_size",
    "blockSizeOk": "blockSize >= 16 and pow2ceil(blockSize) == blockSize",
    "outBlocks": "ceilDiv(dim(shapes.output, 1), blockSize)",
    "zeroPointCols": "ceilDiv(outBlocks, 2) if bits == 4 else outBlocks",
    "workgroupSize": "min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap)",
    "commonShapeValid": "blockSizeOk and ranks.data == 2 and ranks.indices == 1 and ranks.scales == 2 and ranks.output == 2 and tensorDtypes.indices == \"uint32\" and tensorDtypes.scales == \"float32\" and tensorDtypes.output == \"float32\" and dim(shapes.output, 0) == dim(shapes.indices, 0) and dim(shapes.scales, 0) == dim(shapes.data, 0) and dim(shapes.scales, 1) == outBlocks",
    "q4ShapeValid": "commonShapeValid and tensorDtypes.data == \"uint8\" and dim(shapes.output, 1) == dim(shapes.data, 1) * 2",
    "q8ShapeValid": "commonShapeValid and tensorDtypes.data == \"uint8\" and dim(shapes.output, 1) == dim(shapes.data, 1)",
    "zeroPointsValid": "present.zeroPointsT and ranks.zero_points == 2 and tensorDtypes.zero_points == \"uint8\" and dim(shapes.zero_points, 0) == dim(shapes.data, 0) and dim(shapes.zero_points, 1) == zeroPointCols",
    "noZeroMode": "not present.zeroPointsT",
    "zeroMode": "zeroPointsValid",
    "workgroupFits": "workgroupSize > 0",
    "foldedDispatchFits": "ceil(ceil(numel(shapes.output) / device.limits.maxComputeWorkgroupsPerDimension) / workgroupSize) <= device.limits.maxComputeWorkgroupsPerDimension"
  },
  "tunables": { "WORKGROUP_SIZE": 64 },
  "bindingSets": {
    "noZero": [
      {
        "name": "data",
        "arg": "dataT",
        "semantic": "data",
        "buffer": { "type": "read-only-storage" },
        "elementType": "$dataElement"
      },
      {
        "name": "indices",
        "arg": "indicesT",
        "semantic": "indices",
        "buffer": { "type": "read-only-storage" },
        "elementType": "$indexScalar"
      },
      {
        "name": "scales",
        "arg": "scalesT",
        "semantic": "scales",
        "buffer": { "type": "read-only-storage" },
        "elementType": "$scaleScalar"
      },
      {
        "name": "output",
        "arg": "outputT",
        "semantic": "output",
        "buffer": { "type": "storage" },
        "elementType": "$outputElement"
      },
      {
        "name": "params",
        "semantic": "kernel.params",
        "buffer": { "type": "uniform" },
        "struct": {
          "name": "Params",
          "fields": [
            { "name": "indexCount", "type": "u32", "value": "dim(shapes.indices, 0)" },
            { "name": "cols", "type": "u32", "value": "dim(shapes.output, 1)" },
            { "name": "packedCols", "type": "u32", "value": "dim(shapes.data, 1)" },
            { "name": "blocks", "type": "u32", "value": "dim(shapes.scales, 1)" },
            { "name": "blockSize", "type": "u32", "value": "blockSize" },
            { "name": "rows", "type": "u32", "value": "dim(shapes.data, 0)" }
          ]
        }
      }
    ],
    "zero": [
      {
        "name": "data",
        "arg": "dataT",
        "semantic": "data",
        "buffer": { "type": "read-only-storage" },
        "elementType": "$dataElement"
      },
      {
        "name": "indices",
        "arg": "indicesT",
        "semantic": "indices",
        "buffer": { "type": "read-only-storage" },
        "elementType": "$indexScalar"
      },
      {
        "name": "scales",
        "arg": "scalesT",
        "semantic": "scales",
        "buffer": { "type": "read-only-storage" },
        "elementType": "$scaleScalar"
      },
      {
        "name": "zero_points",
        "arg": "zeroPointsT",
        "semantic": "zero_points",
        "buffer": { "type": "read-only-storage" },
        "elementType": "$zeroPointElement"
      },
      {
        "name": "output",
        "arg": "outputT",
        "semantic": "output",
        "buffer": { "type": "storage" },
        "elementType": "$outputElement"
      },
      {
        "name": "params",
        "semantic": "kernel.params",
        "buffer": { "type": "uniform" },
        "struct": {
          "name": "Params",
          "fields": [
            { "name": "indexCount", "type": "u32", "value": "dim(shapes.indices, 0)" },
            { "name": "cols", "type": "u32", "value": "dim(shapes.output, 1)" },
            { "name": "packedCols", "type": "u32", "value": "dim(shapes.data, 1)" },
            { "name": "blocks", "type": "u32", "value": "dim(shapes.scales, 1)" },
            { "name": "blockSize", "type": "u32", "value": "blockSize" },
            { "name": "rows", "type": "u32", "value": "dim(shapes.data, 0)" }
          ]
        }
      }
    ],
    "noZeroIo": [
      {
        "name": "data",
        "arg": "dataT",
        "semantic": "data",
        "buffer": { "type": "read-only-storage" },
        "elementType": "$dataElement"
      },
      {
        "name": "indices",
        "arg": "indicesT",
        "semantic": "indices",
        "buffer": { "type": "read-only-storage" },
        "elementType": "$indexScalar"
      },
      {
        "name": "scales",
        "arg": "scalesT",
        "semantic": "scales",
        "buffer": { "type": "read-only-storage" },
        "elementType": "$scaleScalar"
      },
      {
        "name": "output",
        "arg": "outputT",
        "semantic": "output",
        "buffer": { "type": "storage" },
        "elementType": "$outputElement"
      }
    ],
    "zeroIo": [
      {
        "name": "data",
        "arg": "dataT",
        "semantic": "data",
        "buffer": { "type": "read-only-storage" },
        "elementType": "$dataElement"
      },
      {
        "name": "indices",
        "arg": "indicesT",
        "semantic": "indices",
        "buffer": { "type": "read-only-storage" },
        "elementType": "$indexScalar"
      },
      {
        "name": "scales",
        "arg": "scalesT",
        "semantic": "scales",
        "buffer": { "type": "read-only-storage" },
        "elementType": "$scaleScalar"
      },
      {
        "name": "zero_points",
        "arg": "zeroPointsT",
        "semantic": "zero_points",
        "buffer": { "type": "read-only-storage" },
        "elementType": "$zeroPointElement"
      },
      {
        "name": "output",
        "arg": "outputT",
        "semantic": "output",
        "buffer": { "type": "storage" },
        "elementType": "$outputElement"
      }
    ],
    "q4NoZero": [
      {
        "name": "data",
        "arg": "dataT",
        "semantic": "data",
        "buffer": { "type": "read-only-storage" },
        "elementType": "$dataElement"
      },
      {
        "name": "indices",
        "arg": "indicesT",
        "semantic": "indices",
        "buffer": { "type": "read-only-storage" },
        "elementType": "$indexScalar"
      },
      {
        "name": "scales",
        "arg": "scalesT",
        "semantic": "scales",
        "buffer": { "type": "read-only-storage" },
        "elementType": "$scaleScalar"
      },
      {
        "name": "output",
        "arg": "outputT",
        "semantic": "output",
        "buffer": { "type": "storage" },
        "elementType": "$outputElement"
      },
      {
        "name": "params",
        "semantic": "kernel.params",
        "buffer": { "type": "uniform" },
        "struct": {
          "name": "Params",
          "fields": [
            { "name": "indexCount", "type": "u32", "value": "dim(shapes.indices, 0)" },
            { "name": "packedCols", "type": "u32", "value": "dim(shapes.data, 1)" },
            { "name": "blocks", "type": "u32", "value": "dim(shapes.scales, 1)" },
            { "name": "blockSize", "type": "u32", "value": "blockSize" },
            { "name": "rows", "type": "u32", "value": "dim(shapes.data, 0)" }
          ]
        }
      }
    ],
    "q4Zero": [
      {
        "name": "data",
        "arg": "dataT",
        "semantic": "data",
        "buffer": { "type": "read-only-storage" },
        "elementType": "$dataElement"
      },
      {
        "name": "indices",
        "arg": "indicesT",
        "semantic": "indices",
        "buffer": { "type": "read-only-storage" },
        "elementType": "$indexScalar"
      },
      {
        "name": "scales",
        "arg": "scalesT",
        "semantic": "scales",
        "buffer": { "type": "read-only-storage" },
        "elementType": "$scaleScalar"
      },
      {
        "name": "zero_points",
        "arg": "zeroPointsT",
        "semantic": "zero_points",
        "buffer": { "type": "read-only-storage" },
        "elementType": "$zeroPointElement"
      },
      {
        "name": "output",
        "arg": "outputT",
        "semantic": "output",
        "buffer": { "type": "storage" },
        "elementType": "$outputElement"
      },
      {
        "name": "params",
        "semantic": "kernel.params",
        "buffer": { "type": "uniform" },
        "struct": {
          "name": "Params",
          "fields": [
            { "name": "indexCount", "type": "u32", "value": "dim(shapes.indices, 0)" },
            { "name": "packedCols", "type": "u32", "value": "dim(shapes.data, 1)" },
            { "name": "blocks", "type": "u32", "value": "dim(shapes.scales, 1)" },
            { "name": "zeroPointCols", "type": "u32", "value": "zeroPointCols" },
            { "name": "blockSize", "type": "u32", "value": "blockSize" },
            { "name": "rows", "type": "u32", "value": "dim(shapes.data, 0)" }
          ]
        }
      }
    ],
    "q8Vec4NoZero": [
      {
        "name": "data",
        "arg": "dataT",
        "semantic": "data",
        "buffer": { "type": "read-only-storage" },
        "elementType": "$dataElement"
      },
      {
        "name": "indices",
        "arg": "indicesT",
        "semantic": "indices",
        "buffer": { "type": "read-only-storage" },
        "elementType": "$indexScalar"
      },
      {
        "name": "scales",
        "arg": "scalesT",
        "semantic": "scales",
        "buffer": { "type": "read-only-storage" },
        "elementType": "$scaleScalar"
      },
      {
        "name": "output",
        "arg": "outputT",
        "semantic": "output",
        "buffer": { "type": "storage" },
        "elementType": "$outputElement"
      },
      {
        "name": "params",
        "semantic": "kernel.params",
        "buffer": { "type": "uniform" },
        "struct": {
          "name": "Params",
          "fields": [
            { "name": "indexCount", "type": "u32", "value": "dim(shapes.indices, 0)" },
            { "name": "cols", "type": "u32", "value": "dim(shapes.output, 1)" },
            { "name": "blocks", "type": "u32", "value": "dim(shapes.scales, 1)" },
            { "name": "blockSize", "type": "u32", "value": "blockSize" },
            { "name": "rows", "type": "u32", "value": "dim(shapes.data, 0)" }
          ]
        }
      }
    ],
    "q8Vec4Zero": [
      {
        "name": "data",
        "arg": "dataT",
        "semantic": "data",
        "buffer": { "type": "read-only-storage" },
        "elementType": "$dataElement"
      },
      {
        "name": "indices",
        "arg": "indicesT",
        "semantic": "indices",
        "buffer": { "type": "read-only-storage" },
        "elementType": "$indexScalar"
      },
      {
        "name": "scales",
        "arg": "scalesT",
        "semantic": "scales",
        "buffer": { "type": "read-only-storage" },
        "elementType": "$scaleScalar"
      },
      {
        "name": "zero_points",
        "arg": "zeroPointsT",
        "semantic": "zero_points",
        "buffer": { "type": "read-only-storage" },
        "elementType": "$zeroPointElement"
      },
      {
        "name": "output",
        "arg": "outputT",
        "semantic": "output",
        "buffer": { "type": "storage" },
        "elementType": "$outputElement"
      },
      {
        "name": "params",
        "semantic": "kernel.params",
        "buffer": { "type": "uniform" },
        "struct": {
          "name": "Params",
          "fields": [
            { "name": "indexCount", "type": "u32", "value": "dim(shapes.indices, 0)" },
            { "name": "cols", "type": "u32", "value": "dim(shapes.output, 1)" },
            { "name": "blocks", "type": "u32", "value": "dim(shapes.scales, 1)" },
            { "name": "blockSize", "type": "u32", "value": "blockSize" },
            { "name": "rows", "type": "u32", "value": "dim(shapes.data, 0)" }
          ]
        }
      }
    ]
  },
  "variants": [
    {
      "id": "q8_no_zero_vec4",
      "priority": 10,
      "when": ["q8ShapeValid", "noZeroMode", "bits == 8", "blockSize % 4 == 0", "dim(shapes.output, 1) % 4 == 0", "workgroupFits", "foldedDispatchFits"],
      "constants": {
        "hasZero": false,
        "scalarTail": false,
        "workgroupSize": "workgroupSize",
        "dataElement": "\"vec4<u32>\"",
        "indexScalar": "\"u32\"",
        "scaleScalar": "\"f32\"",
        "outputElement": "\"vec4<f32>\""
      },
      "passes": [
        {
          "id": "main",
          "shader": "gather-block-quantized-q8-vec4.wgsl.jinja",
          "bindings": "q8Vec4NoZero",
          "dispatch": {
            "threads": "dim(shapes.indices, 0) * (dim(shapes.output, 1) / 4)",
            "workgroupSize": "constants.workgroupSize"
          }
        }
      ]
    },
    {
      "id": "q4_no_zero_pair",
      "priority": 10,
      "when": ["q4ShapeValid", "noZeroMode", "bits == 4", "workgroupFits", "foldedDispatchFits"],
      "constants": {
        "hasZero": false,
        "workgroupSize": "workgroupSize",
        "dataElement": "\"u32\"",
        "indexScalar": "\"u32\"",
        "scaleScalar": "\"f32\"",
        "outputElement": "\"vec2<f32>\""
      },
      "passes": [
        {
          "id": "main",
          "shader": "gather-block-quantized-q4-pair.wgsl.jinja",
          "bindings": "q4NoZero",
          "dispatch": {
            "threads": "dim(shapes.indices, 0) * dim(shapes.data, 1)",
            "workgroupSize": "constants.workgroupSize"
          }
        }
      ]
    },
    {
      "id": "q4_zero_pair",
      "priority": 10,
      "when": ["q4ShapeValid", "zeroMode", "bits == 4", "workgroupFits", "foldedDispatchFits"],
      "constants": {
        "hasZero": true,
        "workgroupSize": "workgroupSize",
        "dataElement": "\"u32\"",
        "indexScalar": "\"u32\"",
        "scaleScalar": "\"f32\"",
        "outputElement": "\"vec2<f32>\"",
        "zeroPointElement": "\"u32\""
      },
      "passes": [
        {
          "id": "main",
          "shader": "gather-block-quantized-q4-pair.wgsl.jinja",
          "bindings": "q4Zero",
          "dispatch": {
            "threads": "dim(shapes.indices, 0) * dim(shapes.data, 1)",
            "workgroupSize": "constants.workgroupSize"
          }
        }
      ]
    },
    {
      "id": "q8_zero_vec4",
      "priority": 10,
      "when": ["q8ShapeValid", "zeroMode", "bits == 8", "blockSize % 4 == 0", "dim(shapes.output, 1) % 4 == 0", "workgroupFits", "foldedDispatchFits"],
      "constants": {
        "hasZero": true,
        "scalarTail": false,
        "workgroupSize": "workgroupSize",
        "dataElement": "\"vec4<u32>\"",
        "zeroPointElement": "\"u32\"",
        "indexScalar": "\"u32\"",
        "scaleScalar": "\"f32\"",
        "outputElement": "\"vec4<f32>\""
      },
      "passes": [
        {
          "id": "main",
          "shader": "gather-block-quantized-q8-vec4.wgsl.jinja",
          "bindings": "q8Vec4Zero",
          "dispatch": {
            "threads": "dim(shapes.indices, 0) * (dim(shapes.output, 1) / 4)",
            "workgroupSize": "constants.workgroupSize"
          }
        }
      ]
    },
    {
      "id": "q8_no_zero_tail4",
      "priority": 5,
      "when": ["q8ShapeValid", "noZeroMode", "bits == 8", "workgroupFits", "foldedDispatchFits"],
      "constants": {
        "hasZero": false,
        "scalarTail": true,
        "workgroupSize": "workgroupSize",
        "dataElement": "\"u32\"",
        "indexScalar": "\"u32\"",
        "scaleScalar": "\"f32\"",
        "outputElement": "\"f32\""
      },
      "passes": [
        {
          "id": "main",
          "shader": "gather-block-quantized-q8-vec4.wgsl.jinja",
          "bindings": "noZero",
          "dispatch": {
            "threads": "dim(shapes.indices, 0) * ceilDiv(dim(shapes.output, 1), 4)",
            "workgroupSize": "constants.workgroupSize"
          }
        }
      ]
    },
    {
      "id": "q8_zero_tail4",
      "priority": 5,
      "when": ["q8ShapeValid", "zeroMode", "bits == 8", "workgroupFits", "foldedDispatchFits"],
      "constants": {
        "hasZero": true,
        "scalarTail": true,
        "workgroupSize": "workgroupSize",
        "dataElement": "\"u32\"",
        "indexScalar": "\"u32\"",
        "scaleScalar": "\"f32\"",
        "outputElement": "\"f32\"",
        "zeroPointElement": "\"u32\""
      },
      "passes": [
        {
          "id": "main",
          "shader": "gather-block-quantized-q8-vec4.wgsl.jinja",
          "bindings": "zero",
          "dispatch": {
            "threads": "dim(shapes.indices, 0) * ceilDiv(dim(shapes.output, 1), 4)",
            "workgroupSize": "constants.workgroupSize"
          }
        }
      ]
    }
  ]
}