Add support for bitnets to ORT WebGPU EP

[sushraja-msft]

Description

This change introduces support for BitNet models—specifically microsoft/bitnet-b1.58-2B-4T-bf16—by adding a new operator, BitLinear, which handles the unique weight format used by BitNets for matrix multiplication.

Converted onnx model for testing purposes - https://huggingface.co/sushraja/bitnet-b1.58-2B-4T-fp16-onnx

Motivation and Context

BitNets significantly reduce memory usage due to their compact parameter representation, making them well-suited for client-side inference scenarios.

BitLinear Operator

BitNets encode matrix weights as ternary values (+1, 0, -1) along with a scale factor. These ternary values are represented in base-3 and packed such that 5 weights fit into a single uint8.

Inference Workflow

The inference process involves two main steps:

Step 1: Quantization of Input A

• Input tensor A (in fp16) is quantized to int8 with a single scale per token.
• Four int8 values are packed into a u32.
• Every 5th value of A is extracted and stored in a separate tensor (A5) to align with the BitNet weight packing.
• Result: For every 20 values of A, you get:
  • A vec4 (4 × 4 packed values)
  • One u32 for the 5th values

Step 2: Multiplication with Weights B

• Weights B are stored transposed as a stream of uint8s, each encoding 5 ternary weights.
• Each uint8 is unpacked using a lookup table into:
  • A u32 containing 4 packed weights
  • 1 extra weight (the 5th), which is collected across 4 uint8s into a separateu32
• This results in:
  • A vec4 from the packed weights
  • One u32 from the extra weights
    These are then multiplied using the DP4A instruction, leveraging shared memory for efficient cooperative matmul.

Key Notes

The BitLinear operator does not enforce a specific layout for B.
Weights are stored using ternary packing (5 weights per uint8), and decompression is handled dynamically at runtime.

[github-actions]

You can commit the suggested changes from lintrunner.

[github-actions]

Suggested change

	const uint32_t quantize_5th_output_size = M * K_PADDED / kQuantizationBlockSize; // every 5th element packed int u32
	const uint32_t quantize_5th_output_size = M * K_PADDED / kQuantizationBlockSize; // every 5th element packed int u32

[github-actions]

Suggested change

	Status GenerateShaderCode(ShaderHelper& sh) const override;
	Status GenerateShaderCode(ShaderHelper& sh) const override;

[sushraja-msft]

Going to work on tests in the next couple of days. Sharing this PR for early review on operator shape and any feedback.