orbitquant 0.2.1__tar.gz → 0.2.2__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {orbitquant-0.2.1 → orbitquant-0.2.2}/.github/workflows/publish-pypi.yml +1 -1
- {orbitquant-0.2.1 → orbitquant-0.2.2}/PKG-INFO +1 -1
- {orbitquant-0.2.1 → orbitquant-0.2.2}/docs/kernel-audit.md +21 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cuda/packed_matmul.cu +21 -22
- {orbitquant-0.2.1 → orbitquant-0.2.2}/pyproject.toml +1 -1
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/__init__.py +1 -1
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/artifacts/model_card.py +18 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/kernels/triton_cuda.py +114 -8
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/layers.py +34 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/modeling.py +3 -1
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/pipeline.py +2 -2
- {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_diffusers_modelmixin_integration.py +18 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_kernels.py +36 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_model_card.py +25 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_orbit_linear.py +41 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/uv.lock +1 -1
- {orbitquant-0.2.1 → orbitquant-0.2.2}/.github/workflows/ci.yml +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/.gitignore +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/LICENSE +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/README.md +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/docs/paper-methodology-audit.md +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/native-kernels/orbitquant-packed-matmul/CARD.md +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/native-kernels/orbitquant-packed-matmul/benchmarks/benchmark.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/native-kernels/orbitquant-packed-matmul/build.toml +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/native-kernels/orbitquant-packed-matmul/example.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/native-kernels/orbitquant-packed-matmul/flake.lock +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/native-kernels/orbitquant-packed-matmul/flake.nix +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_metal/packed_matmul.metal +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_metal/packed_matmul.mm +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/native-kernels/orbitquant-packed-matmul/tests/__init__.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/native-kernels/orbitquant-packed-matmul/tests/test_packed_matmul.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/native-kernels/orbitquant-packed-matmul/torch-ext/orbitquant_packed_matmul/__init__.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/native-kernels/orbitquant-packed-matmul/torch-ext/torch_binding.cpp +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/native-kernels/orbitquant-packed-matmul/torch-ext/torch_binding.h +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/scripts/run_cuda_kernel_checks.sh +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/scripts/run_hf_compat_checks.sh +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/scripts/run_mps_kernel_checks.sh +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/scripts/run_paper_methodology_checks.sh +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/scripts/runpod_ssh_health.sh +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/scripts/verify_hf_kernel_model_artifact.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/adaln.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/artifacts/__init__.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/artifacts/assets.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/artifacts/benchmark.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/artifacts/checksums.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/artifacts/comparisons.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/artifacts/loader.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/artifacts/manifest.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/artifacts/refresh.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/artifacts/repair.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/artifacts/validator.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/artifacts/writer.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/benchmarks.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/cli/__init__.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/cli/main.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/codebooks/__init__.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/codebooks/lloyd_max.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/config.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/errors.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/eval/__init__.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/eval/assets.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/eval/external_export.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/eval/external_metrics.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/eval/external_plan.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/eval/metrics.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/eval/native_plan.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/eval/native_runner.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/eval/native_settings.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/eval/prompts.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/eval/report.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/functional.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/hub.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/kernels/__init__.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/kernels/dispatch.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/kernels/mps.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/kernels/native_packed_matmul.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/linear_adapters.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/packing/__init__.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/packing/bitpack.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/policies/__init__.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/policies/flux.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/policies/flux2.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/policies/generic_dit.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/policies/wan.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/policies/z_image.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/quantizer.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/recipes.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/rotations/__init__.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/rotations/fwht.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/rotations/rpbh.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/transformers_ops.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_adaln_rtn.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_artifact_writer.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_benchmarks.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_bitpack.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_ci.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_cli.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_cli_report.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_codebooks.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_config.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_distribution.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_documentation.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_eval_assets.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_eval_prompts.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_eval_report.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_external_eval_plan.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_external_export.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_hub.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_kernel_model_verifier.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_linear_adapters.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_manifest.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_model_quantization.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_native_kernel_package.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_native_packed_matmul.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_native_runner.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_native_settings.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_optional_dependencies.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_paper_alignment.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_paper_methodology.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_pipeline_helpers.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_policies.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_prompt_jobs.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_prompt_selection.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_quantizer_adapter.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_rpbh.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_target_policies.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_transformers_pretrained_integration.py +0 -0
- {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_universal_transformers.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: orbitquant
|
|
3
|
-
Version: 0.2.
|
|
3
|
+
Version: 0.2.2
|
|
4
4
|
Summary: Calibration-free OrbitQuant for transformer linear projections
|
|
5
5
|
Project-URL: Homepage, https://github.com/iamwavecut/OrbitQuant
|
|
6
6
|
Project-URL: Repository, https://github.com/iamwavecut/OrbitQuant
|
|
@@ -107,6 +107,27 @@ The native CUDA package was also built and tested on an NVIDIA RTX 4090
|
|
|
107
107
|
All 34 applicable package tests passed for W2/W3/W4/W6, FP16/BF16, bias and
|
|
108
108
|
no-bias paths, partial output tiles, short rows, and tensor-core rows.
|
|
109
109
|
|
|
110
|
+
The same local package build passed its 34 CUDA tests on an NVIDIA A40
|
|
111
|
+
(`sm_86`) with Torch 2.9.1+cu128 and CUDA 12.8. A full FLUX.2 Klein 9B W4A4
|
|
112
|
+
pipeline exercised 396 packed linears across the transformer and Qwen3 text
|
|
113
|
+
encoder. Native 1024x1024 generation at four steps used
|
|
114
|
+
`native_packed_matmul` for every packed linear and the Triton activation path;
|
|
115
|
+
no full weight matrix was materialized.
|
|
116
|
+
|
|
117
|
+
Representative A40 W4A4 layer timings from that pipeline are:
|
|
118
|
+
|
|
119
|
+
| Projection | Rows | Shape | RPBH + activation quantization | Native packed matmul | Full layer |
|
|
120
|
+
| --- | ---: | ---: | ---: | ---: | ---: |
|
|
121
|
+
| Double-stream Q | 4096 | 4096 -> 4096 | 0.698 ms | 2.607 ms | 3.279 ms |
|
|
122
|
+
| Single-stream fused input | 4608 | 4096 -> 36864 | 0.770 ms | 25.962 ms | 26.677 ms |
|
|
123
|
+
| Single-stream output | 4608 | 16384 -> 4096 | 4.257 ms | 11.244 ms | 15.514 ms |
|
|
124
|
+
|
|
125
|
+
The corresponding full-pipeline hot-generation mean was 7.418 seconds with a
|
|
126
|
+
17.66 GB NVML peak. On the same A40 and prompt pack, BF16 measured 3.911
|
|
127
|
+
seconds and 40.83 GB, while SDNQ UINT4 measured 3.546 seconds and 17.55 GB.
|
|
128
|
+
These results establish the memory reduction and the remaining throughput gap;
|
|
129
|
+
they do not support a model-level speedup claim for OrbitQuant on A40.
|
|
130
|
+
|
|
110
131
|
Measured W4 BF16 operator latency for `in_features=768` and
|
|
111
132
|
`out_features=2304`:
|
|
112
133
|
|
|
@@ -58,7 +58,7 @@ __device__ __forceinline__ __nv_bfloat16 orbitquant_mma_from_float<__nv_bfloat16
|
|
|
58
58
|
return __float2bfloat16(value);
|
|
59
59
|
}
|
|
60
60
|
|
|
61
|
-
template <int Bits, typename mma_t>
|
|
61
|
+
template <int Bits, int Values, typename mma_t>
|
|
62
62
|
__device__ __forceinline__ void decode_mma64_weight_segment(
|
|
63
63
|
mma_t *__restrict__ destination,
|
|
64
64
|
uint8_t const *__restrict__ packed_weight_indices,
|
|
@@ -68,9 +68,9 @@ __device__ __forceinline__ void decode_mma64_weight_segment(
|
|
|
68
68
|
int64_t global_k,
|
|
69
69
|
int64_t in_features,
|
|
70
70
|
bool valid_segment) {
|
|
71
|
-
constexpr int values = 32;
|
|
72
71
|
constexpr uint32_t mask = (1u << Bits) - 1u;
|
|
73
|
-
|
|
72
|
+
constexpr int word_count = (Values * Bits + 31) / 32;
|
|
73
|
+
uint32_t packed_words[word_count] = {};
|
|
74
74
|
float norm = 0.0f;
|
|
75
75
|
if (valid_segment) {
|
|
76
76
|
const int64_t value_offset = global_col * in_features + global_k;
|
|
@@ -78,14 +78,14 @@ __device__ __forceinline__ void decode_mma64_weight_segment(
|
|
|
78
78
|
auto const *words = reinterpret_cast<uint32_t const *>(
|
|
79
79
|
packed_weight_indices + byte_index);
|
|
80
80
|
#pragma unroll
|
|
81
|
-
for (int word = 0; word <
|
|
81
|
+
for (int word = 0; word < word_count; ++word) {
|
|
82
82
|
packed_words[word] = words[word];
|
|
83
83
|
}
|
|
84
84
|
norm = static_cast<float>(row_norms[global_col]);
|
|
85
85
|
}
|
|
86
86
|
|
|
87
87
|
#pragma unroll
|
|
88
|
-
for (int index_offset = 0; index_offset <
|
|
88
|
+
for (int index_offset = 0; index_offset < Values; ++index_offset) {
|
|
89
89
|
const int bit_start = index_offset * Bits;
|
|
90
90
|
const int word_index = bit_start >> 5;
|
|
91
91
|
const int shift = bit_start & 31;
|
|
@@ -111,23 +111,22 @@ __global__ void orbitquant_packed_matmul_mma64_kernel(
|
|
|
111
111
|
int64_t rows,
|
|
112
112
|
int64_t out_features,
|
|
113
113
|
int64_t in_features) {
|
|
114
|
-
constexpr int tile_m =
|
|
115
|
-
constexpr int tile_n =
|
|
114
|
+
constexpr int tile_m = 128;
|
|
115
|
+
constexpr int tile_n = 128;
|
|
116
116
|
constexpr int tile_k = 64;
|
|
117
117
|
constexpr int padded_k = 72;
|
|
118
|
-
constexpr int warps_per_block =
|
|
118
|
+
constexpr int warps_per_block = 8;
|
|
119
119
|
constexpr int warp_tile = 16;
|
|
120
120
|
constexpr int col_tiles = tile_n / warp_tile;
|
|
121
121
|
constexpr int x_vector_values = 8;
|
|
122
122
|
constexpr int x_vectors_per_row = tile_k / x_vector_values;
|
|
123
|
-
constexpr int weight_segment_values = 32;
|
|
123
|
+
constexpr int weight_segment_values = Bits == 3 ? 32 : 16;
|
|
124
124
|
constexpr int weight_segments_per_row = tile_k / weight_segment_values;
|
|
125
125
|
static_assert(sizeof(storage_t) == sizeof(mma_t));
|
|
126
126
|
|
|
127
127
|
__shared__ mma_t x_tile[tile_m * padded_k];
|
|
128
128
|
__shared__ mma_t weight_tile[tile_n * padded_k];
|
|
129
|
-
__shared__ float accumulator_tile[
|
|
130
|
-
warps_per_block * col_tiles * warp_tile * warp_tile];
|
|
129
|
+
__shared__ float accumulator_tile[warps_per_block * warp_tile * warp_tile];
|
|
131
130
|
|
|
132
131
|
const int warp_id = threadIdx.x / warpSize;
|
|
133
132
|
const int lane = threadIdx.x & (warpSize - 1);
|
|
@@ -160,14 +159,14 @@ __global__ void orbitquant_packed_matmul_mma64_kernel(
|
|
|
160
159
|
}
|
|
161
160
|
|
|
162
161
|
constexpr int weight_tasks = tile_n * weight_segments_per_row;
|
|
163
|
-
|
|
164
|
-
|
|
162
|
+
for (int weight_task = threadIdx.x; weight_task < weight_tasks;
|
|
163
|
+
weight_task += blockDim.x) {
|
|
165
164
|
const int local_col = weight_task / weight_segments_per_row;
|
|
166
165
|
const int local_segment =
|
|
167
166
|
weight_task - local_col * weight_segments_per_row;
|
|
168
167
|
const int local_k = local_segment * weight_segment_values;
|
|
169
168
|
const int64_t global_col = block_col + local_col;
|
|
170
|
-
decode_mma64_weight_segment<Bits>(
|
|
169
|
+
decode_mma64_weight_segment<Bits, weight_segment_values>(
|
|
171
170
|
weight_tile + local_col * padded_k + local_k,
|
|
172
171
|
packed_weight_indices,
|
|
173
172
|
row_norms,
|
|
@@ -204,10 +203,9 @@ __global__ void orbitquant_packed_matmul_mma64_kernel(
|
|
|
204
203
|
__syncthreads();
|
|
205
204
|
}
|
|
206
205
|
|
|
206
|
+
float *warp_accumulator = accumulator_tile + warp_id * warp_tile * warp_tile;
|
|
207
207
|
#pragma unroll
|
|
208
208
|
for (int col_tile = 0; col_tile < col_tiles; ++col_tile) {
|
|
209
|
-
float *warp_accumulator = accumulator_tile +
|
|
210
|
-
(warp_id * col_tiles + col_tile) * warp_tile * warp_tile;
|
|
211
209
|
wmma::store_matrix_sync(
|
|
212
210
|
warp_accumulator,
|
|
213
211
|
accumulators[col_tile],
|
|
@@ -228,6 +226,7 @@ __global__ void orbitquant_packed_matmul_mma64_kernel(
|
|
|
228
226
|
static_cast<storage_t>(value);
|
|
229
227
|
}
|
|
230
228
|
}
|
|
229
|
+
__syncwarp();
|
|
231
230
|
}
|
|
232
231
|
}
|
|
233
232
|
|
|
@@ -648,9 +647,9 @@ void matmul_packed_weight(
|
|
|
648
647
|
if (x.scalar_type() == at::kBFloat16 && x.size(0) >= 9) {
|
|
649
648
|
if (in_features % 64 == 0 &&
|
|
650
649
|
(bits == 2 || bits == 3 || bits == 4 || bits == 6)) {
|
|
651
|
-
constexpr int mma_tile_m =
|
|
652
|
-
constexpr int mma_tile_n =
|
|
653
|
-
const dim3 mma_block(
|
|
650
|
+
constexpr int mma_tile_m = 128;
|
|
651
|
+
constexpr int mma_tile_n = 128;
|
|
652
|
+
const dim3 mma_block(256);
|
|
654
653
|
const dim3 mma_grid(
|
|
655
654
|
(out_features + mma_tile_n - 1) / mma_tile_n,
|
|
656
655
|
(x.size(0) + mma_tile_m - 1) / mma_tile_m);
|
|
@@ -740,9 +739,9 @@ void matmul_packed_weight(
|
|
|
740
739
|
if (x.scalar_type() == at::kHalf && x.size(0) >= 9) {
|
|
741
740
|
if (in_features % 64 == 0 &&
|
|
742
741
|
(bits == 2 || bits == 3 || bits == 4 || bits == 6)) {
|
|
743
|
-
constexpr int mma_tile_m =
|
|
744
|
-
constexpr int mma_tile_n =
|
|
745
|
-
const dim3 mma_block(
|
|
742
|
+
constexpr int mma_tile_m = 128;
|
|
743
|
+
constexpr int mma_tile_n = 128;
|
|
744
|
+
const dim3 mma_block(256);
|
|
746
745
|
const dim3 mma_grid(
|
|
747
746
|
(out_features + mma_tile_n - 1) / mma_tile_n,
|
|
748
747
|
(x.size(0) + mma_tile_m - 1) / mma_tile_m);
|
|
@@ -408,6 +408,22 @@ def _comparison_prompt_section(
|
|
|
408
408
|
return lines
|
|
409
409
|
|
|
410
410
|
|
|
411
|
+
def _observed_quality_section(
|
|
412
|
+
benchmark_summary: dict[str, Any] | None,
|
|
413
|
+
) -> list[str]:
|
|
414
|
+
if not isinstance(benchmark_summary, dict):
|
|
415
|
+
return []
|
|
416
|
+
observed_quality = benchmark_summary.get("observed_quality")
|
|
417
|
+
if not isinstance(observed_quality, str) or not observed_quality.strip():
|
|
418
|
+
return []
|
|
419
|
+
return [
|
|
420
|
+
"## Observed Quality",
|
|
421
|
+
"",
|
|
422
|
+
f"**Warning:** {observed_quality.strip()}",
|
|
423
|
+
"",
|
|
424
|
+
]
|
|
425
|
+
|
|
426
|
+
|
|
411
427
|
def render_model_card(
|
|
412
428
|
manifest: OrbitQuantManifest,
|
|
413
429
|
*,
|
|
@@ -424,6 +440,7 @@ def render_model_card(
|
|
|
424
440
|
comparison_prompt_lines = _comparison_prompt_section(
|
|
425
441
|
data["source_model_id"], benchmark_summary
|
|
426
442
|
)
|
|
443
|
+
observed_quality_lines = _observed_quality_section(benchmark_summary)
|
|
427
444
|
adaln_group_size = int(data.get("adaln_group_size", 64))
|
|
428
445
|
adaln_default_note = (
|
|
429
446
|
"- AdaLN group-size note: paper default."
|
|
@@ -531,6 +548,7 @@ def render_model_card(
|
|
|
531
548
|
"- Calibration data: none",
|
|
532
549
|
"- Text encoders and VAE: left in source precision by default",
|
|
533
550
|
"",
|
|
551
|
+
*observed_quality_lines,
|
|
534
552
|
*comparison_lines,
|
|
535
553
|
"## Source",
|
|
536
554
|
"",
|
|
@@ -236,6 +236,50 @@ def _fused_rpbh_quantize_activation_kernel(
|
|
|
236
236
|
tl.store(output_ptr + row * dim + cols, quantized * norm, mask=mask)
|
|
237
237
|
|
|
238
238
|
|
|
239
|
+
@triton.jit
|
|
240
|
+
def _fused_rpbh_fwht_1024_chunk_kernel(
|
|
241
|
+
input_ptr,
|
|
242
|
+
norms_ptr,
|
|
243
|
+
permutation_ptr,
|
|
244
|
+
signs_ptr,
|
|
245
|
+
work_ptr,
|
|
246
|
+
rows: tl.constexpr,
|
|
247
|
+
dim: tl.constexpr,
|
|
248
|
+
num_blocks: tl.constexpr,
|
|
249
|
+
orbit_block_size: tl.constexpr,
|
|
250
|
+
chunks_per_orbit_block: tl.constexpr,
|
|
251
|
+
eps: tl.constexpr,
|
|
252
|
+
):
|
|
253
|
+
program = tl.program_id(0)
|
|
254
|
+
chunk = program % chunks_per_orbit_block
|
|
255
|
+
row_block = program // chunks_per_orbit_block
|
|
256
|
+
row = row_block // num_blocks
|
|
257
|
+
orbit_block = row_block % num_blocks
|
|
258
|
+
local_cols = tl.arange(0, 1024)
|
|
259
|
+
cols = orbit_block * orbit_block_size + chunk * 1024 + local_cols
|
|
260
|
+
mask = (row < rows) & (cols < dim)
|
|
261
|
+
source_cols = tl.load(permutation_ptr + cols, mask=mask, other=0).to(tl.int64)
|
|
262
|
+
signs = tl.load(signs_ptr + cols, mask=mask, other=1).to(tl.float32)
|
|
263
|
+
norm = tl.load(norms_ptr + row, mask=row < rows, other=0.0).to(tl.float32)
|
|
264
|
+
values = tl.load(
|
|
265
|
+
input_ptr + row * dim + source_cols,
|
|
266
|
+
mask=mask,
|
|
267
|
+
other=0.0,
|
|
268
|
+
).to(tl.float32)
|
|
269
|
+
values = values * signs / (norm + eps)
|
|
270
|
+
values = _fwht_local_stage(values, 1024, 1)
|
|
271
|
+
values = _fwht_local_stage(values, 1024, 2)
|
|
272
|
+
values = _fwht_local_stage(values, 1024, 4)
|
|
273
|
+
values = _fwht_local_stage(values, 1024, 8)
|
|
274
|
+
values = _fwht_local_stage(values, 1024, 16)
|
|
275
|
+
values = _fwht_local_stage(values, 1024, 32)
|
|
276
|
+
values = _fwht_local_stage(values, 1024, 64)
|
|
277
|
+
values = _fwht_local_stage(values, 1024, 128)
|
|
278
|
+
values = _fwht_local_stage(values, 1024, 256)
|
|
279
|
+
values = _fwht_local_stage(values, 1024, 512)
|
|
280
|
+
tl.store(work_ptr + row * dim + cols, values, mask=mask)
|
|
281
|
+
|
|
282
|
+
|
|
239
283
|
@triton.jit
|
|
240
284
|
def _codebook_rescale_kernel(
|
|
241
285
|
rotated_ptr,
|
|
@@ -868,24 +912,86 @@ def quantize_activations_with_triton(
|
|
|
868
912
|
|
|
869
913
|
orbit_block_size = int(rotation.block_size)
|
|
870
914
|
num_blocks = int(rotation.num_blocks)
|
|
871
|
-
|
|
872
|
-
|
|
915
|
+
if orbit_block_size <= 1024:
|
|
916
|
+
fwht_stages = orbit_block_size.bit_length() - 1
|
|
917
|
+
_fused_rpbh_quantize_activation_kernel[(rows * num_blocks,)](
|
|
918
|
+
input_contiguous,
|
|
919
|
+
norms,
|
|
920
|
+
permutation,
|
|
921
|
+
signs,
|
|
922
|
+
centroids,
|
|
923
|
+
boundaries,
|
|
924
|
+
output,
|
|
925
|
+
rows=rows,
|
|
926
|
+
dim=dim,
|
|
927
|
+
num_blocks=num_blocks,
|
|
928
|
+
orbit_block_size=orbit_block_size,
|
|
929
|
+
fwht_stages=fwht_stages,
|
|
930
|
+
levels=centroids.numel(),
|
|
931
|
+
eps=float(eps),
|
|
932
|
+
inv_sqrt_block=float(rotation.normalization),
|
|
933
|
+
num_warps=8 if orbit_block_size >= 512 else 4,
|
|
934
|
+
)
|
|
935
|
+
return output.reshape(original_shape)
|
|
936
|
+
|
|
937
|
+
element_block_size = 256
|
|
938
|
+
work = torch.empty(total, device=x.device, dtype=torch.float32)
|
|
939
|
+
chunks_per_orbit_block = orbit_block_size // 1024
|
|
940
|
+
_fused_rpbh_fwht_1024_chunk_kernel[
|
|
941
|
+
(rows * num_blocks * chunks_per_orbit_block,)
|
|
942
|
+
](
|
|
873
943
|
input_contiguous,
|
|
874
944
|
norms,
|
|
875
945
|
permutation,
|
|
876
946
|
signs,
|
|
877
|
-
|
|
878
|
-
boundaries,
|
|
879
|
-
output,
|
|
947
|
+
work,
|
|
880
948
|
rows=rows,
|
|
881
949
|
dim=dim,
|
|
882
950
|
num_blocks=num_blocks,
|
|
883
951
|
orbit_block_size=orbit_block_size,
|
|
884
|
-
|
|
885
|
-
levels=centroids.numel(),
|
|
952
|
+
chunks_per_orbit_block=chunks_per_orbit_block,
|
|
886
953
|
eps=float(eps),
|
|
954
|
+
num_warps=8,
|
|
955
|
+
)
|
|
956
|
+
|
|
957
|
+
stage_width = 1024
|
|
958
|
+
while stage_width * 2 < orbit_block_size:
|
|
959
|
+
total_quads = rows * num_blocks * (orbit_block_size // 4)
|
|
960
|
+
_fwht_two_stage_kernel[(triton.cdiv(total_quads, element_block_size),)](
|
|
961
|
+
work,
|
|
962
|
+
total_quads=total_quads,
|
|
963
|
+
in_features=dim,
|
|
964
|
+
num_blocks=num_blocks,
|
|
965
|
+
orbit_block_size=orbit_block_size,
|
|
966
|
+
stage_width=stage_width,
|
|
967
|
+
block_size=element_block_size,
|
|
968
|
+
)
|
|
969
|
+
stage_width *= 4
|
|
970
|
+
if stage_width < orbit_block_size:
|
|
971
|
+
total_pairs = rows * num_blocks * (orbit_block_size // 2)
|
|
972
|
+
_fwht_stage_kernel[(triton.cdiv(total_pairs, element_block_size),)](
|
|
973
|
+
work,
|
|
974
|
+
total_pairs=total_pairs,
|
|
975
|
+
in_features=dim,
|
|
976
|
+
num_blocks=num_blocks,
|
|
977
|
+
orbit_block_size=orbit_block_size,
|
|
978
|
+
stage_width=stage_width,
|
|
979
|
+
block_size=element_block_size,
|
|
980
|
+
)
|
|
981
|
+
|
|
982
|
+
_quantize_activation_work_rescale_kernel[
|
|
983
|
+
(triton.cdiv(total, element_block_size),)
|
|
984
|
+
](
|
|
985
|
+
work,
|
|
986
|
+
norms,
|
|
987
|
+
centroids,
|
|
988
|
+
boundaries,
|
|
989
|
+
output,
|
|
990
|
+
total=total,
|
|
991
|
+
dim=dim,
|
|
992
|
+
levels=centroids.numel(),
|
|
887
993
|
inv_sqrt_block=float(rotation.normalization),
|
|
888
|
-
|
|
994
|
+
block_size=element_block_size,
|
|
889
995
|
)
|
|
890
996
|
return output.reshape(original_shape)
|
|
891
997
|
|
|
@@ -223,6 +223,7 @@ class OrbitQuantLinear(nn.Module):
|
|
|
223
223
|
self.last_effective_runtime_mode: str | None = None
|
|
224
224
|
self.last_activation_kernel_backend: str | None = None
|
|
225
225
|
self.last_forward_device_type: str | None = None
|
|
226
|
+
self._derived_constants_valid = True
|
|
226
227
|
|
|
227
228
|
@classmethod
|
|
228
229
|
def from_linear(
|
|
@@ -341,6 +342,38 @@ class OrbitQuantLinear(nn.Module):
|
|
|
341
342
|
self._dequantized_weight_cache = None
|
|
342
343
|
self._dequantized_weight_cache_key = None
|
|
343
344
|
|
|
345
|
+
def _apply(self, fn, recurse: bool = True):
|
|
346
|
+
result = super()._apply(fn, recurse=recurse)
|
|
347
|
+
self._derived_constants_valid = False
|
|
348
|
+
self.clear_dequantized_cache()
|
|
349
|
+
return result
|
|
350
|
+
|
|
351
|
+
def _ensure_derived_constants(self, device: torch.device) -> None:
|
|
352
|
+
if self._derived_constants_valid and self._rotation_permutation.device == device:
|
|
353
|
+
return
|
|
354
|
+
constants = (
|
|
355
|
+
("_rotation_permutation", self.rotation.permutation, torch.int64),
|
|
356
|
+
("_rotation_signs", self.rotation.signs, torch.int8),
|
|
357
|
+
(
|
|
358
|
+
"_activation_codebook_centroids",
|
|
359
|
+
self.activation_codebook.centroids,
|
|
360
|
+
torch.float32,
|
|
361
|
+
),
|
|
362
|
+
(
|
|
363
|
+
"_activation_codebook_boundaries",
|
|
364
|
+
self.activation_codebook.boundaries,
|
|
365
|
+
torch.float32,
|
|
366
|
+
),
|
|
367
|
+
(
|
|
368
|
+
"_weight_codebook_centroids",
|
|
369
|
+
self.weight_codebook.centroids,
|
|
370
|
+
torch.float32,
|
|
371
|
+
),
|
|
372
|
+
)
|
|
373
|
+
for name, source, dtype in constants:
|
|
374
|
+
setattr(self, name, source.to(device=device, dtype=dtype).clone())
|
|
375
|
+
self._derived_constants_valid = True
|
|
376
|
+
|
|
344
377
|
def _constant_buffer(
|
|
345
378
|
self,
|
|
346
379
|
name: str,
|
|
@@ -518,6 +551,7 @@ class OrbitQuantLinear(nn.Module):
|
|
|
518
551
|
return dequantized
|
|
519
552
|
|
|
520
553
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
554
|
+
self._ensure_derived_constants(x.device)
|
|
521
555
|
runtime_mode = (
|
|
522
556
|
self._resolve_auto_fused_runtime(x)
|
|
523
557
|
if self.runtime_mode == "auto_fused"
|
|
@@ -353,7 +353,9 @@ def quantize_model(
|
|
|
353
353
|
model.quantization_config = resolved_config
|
|
354
354
|
model.orbitquant_summary = summary
|
|
355
355
|
model_config = getattr(model, "config", None)
|
|
356
|
-
if
|
|
356
|
+
if hasattr(model, "register_to_config"):
|
|
357
|
+
model.register_to_config(quantization_config=resolved_config.to_dict())
|
|
358
|
+
elif model_config is not None:
|
|
357
359
|
model_config.quantization_config = resolved_config.to_dict()
|
|
358
360
|
return summary
|
|
359
361
|
|
|
@@ -12,7 +12,7 @@ from orbitquant.artifacts import (
|
|
|
12
12
|
save_orbitquant_artifact,
|
|
13
13
|
)
|
|
14
14
|
from orbitquant.config import OrbitQuantConfig
|
|
15
|
-
from orbitquant.modeling import QuantizationSummary,
|
|
15
|
+
from orbitquant.modeling import QuantizationSummary, quantize_model
|
|
16
16
|
from orbitquant.quantizer import register_hf_quantizers
|
|
17
17
|
|
|
18
18
|
|
|
@@ -124,7 +124,7 @@ def quantize_pipeline(
|
|
|
124
124
|
synchronize_per_module: bool = False,
|
|
125
125
|
) -> QuantizationSummary:
|
|
126
126
|
target = _get_pipeline_component(pipeline, component)
|
|
127
|
-
return
|
|
127
|
+
return quantize_model(
|
|
128
128
|
target,
|
|
129
129
|
config,
|
|
130
130
|
quantization_device=quantization_device,
|
|
@@ -5,6 +5,7 @@ import torch
|
|
|
5
5
|
|
|
6
6
|
from orbitquant.config import OrbitQuantConfig
|
|
7
7
|
from orbitquant.layers import OrbitQuantLinear
|
|
8
|
+
from orbitquant.modeling import quantize_model
|
|
8
9
|
from orbitquant.quantizer import register_hf_quantizers
|
|
9
10
|
|
|
10
11
|
diffusers = pytest.importorskip("diffusers")
|
|
@@ -101,3 +102,20 @@ def test_diffusers_modelmixin_save_pretrained_round_trips_pre_quantized_model(
|
|
|
101
102
|
assert isinstance(restored.proj_out, torch.nn.Linear)
|
|
102
103
|
x = torch.randn(2, 3, 16)
|
|
103
104
|
assert torch.isfinite(restored(x)).all()
|
|
105
|
+
|
|
106
|
+
|
|
107
|
+
def test_diffusers_manual_quantize_model_persists_quantization_config(tmp_path):
|
|
108
|
+
model = TinyDiffusersTransformer()
|
|
109
|
+
config = OrbitQuantConfig(
|
|
110
|
+
block_size=8,
|
|
111
|
+
runtime_mode="dequant_bf16",
|
|
112
|
+
activation_kernel_backend="cpu",
|
|
113
|
+
)
|
|
114
|
+
|
|
115
|
+
quantize_model(model, config, quantization_device="cpu")
|
|
116
|
+
model.save_pretrained(tmp_path)
|
|
117
|
+
saved_config = json.loads((tmp_path / "config.json").read_text())
|
|
118
|
+
restored = TinyDiffusersTransformer.from_pretrained(tmp_path)
|
|
119
|
+
|
|
120
|
+
assert saved_config["quantization_config"] == config.to_dict()
|
|
121
|
+
assert isinstance(restored.transformer_blocks[0]["attn"]["to_q"], OrbitQuantLinear)
|
|
@@ -813,6 +813,42 @@ def test_triton_cuda_activation_kernel_matches_paper_sized_rpbh_block():
|
|
|
813
813
|
assert torch.allclose(actual, expected, atol=1e-5, rtol=1e-5)
|
|
814
814
|
|
|
815
815
|
|
|
816
|
+
def test_triton_cuda_activation_kernel_matches_large_rpbh_block():
|
|
817
|
+
if not torch.cuda.is_available() or not available_backends()["triton_cuda"]:
|
|
818
|
+
pytest.skip("CUDA/Triton backend is not available")
|
|
819
|
+
|
|
820
|
+
torch.manual_seed(3)
|
|
821
|
+
x = torch.randn(3, 4096, device="cuda", dtype=torch.bfloat16)
|
|
822
|
+
rotation = RPBHRotation(dim=4096, seed=0, block_size="paper")
|
|
823
|
+
codebook = get_codebook(dim=4096, bits=4)
|
|
824
|
+
expected = quantize_activations(x, rotation=rotation, codebook=codebook, eps=1e-12)
|
|
825
|
+
actual = quantize_activations_kernel(
|
|
826
|
+
x, rotation=rotation, codebook=codebook, eps=1e-12, backend="triton_cuda"
|
|
827
|
+
)
|
|
828
|
+
|
|
829
|
+
assert rotation.block_size == 4096
|
|
830
|
+
assert actual.dtype == torch.bfloat16
|
|
831
|
+
assert torch.allclose(actual.float(), expected.float(), atol=1e-2, rtol=1e-2)
|
|
832
|
+
|
|
833
|
+
|
|
834
|
+
def test_triton_cuda_activation_kernel_matches_very_large_rpbh_block():
|
|
835
|
+
if not torch.cuda.is_available() or not available_backends()["triton_cuda"]:
|
|
836
|
+
pytest.skip("CUDA/Triton backend is not available")
|
|
837
|
+
|
|
838
|
+
torch.manual_seed(4)
|
|
839
|
+
x = torch.randn(2, 16384, device="cuda", dtype=torch.bfloat16)
|
|
840
|
+
rotation = RPBHRotation(dim=16384, seed=0, block_size="paper")
|
|
841
|
+
codebook = get_codebook(dim=16384, bits=4)
|
|
842
|
+
expected = quantize_activations(x, rotation=rotation, codebook=codebook, eps=1e-12)
|
|
843
|
+
actual = quantize_activations_kernel(
|
|
844
|
+
x, rotation=rotation, codebook=codebook, eps=1e-12, backend="triton_cuda"
|
|
845
|
+
)
|
|
846
|
+
|
|
847
|
+
assert rotation.block_size == 16384
|
|
848
|
+
assert actual.dtype == torch.bfloat16
|
|
849
|
+
assert torch.allclose(actual.float(), expected.float(), atol=1e-2, rtol=1e-2)
|
|
850
|
+
|
|
851
|
+
|
|
816
852
|
def test_cuda_quantize_linear_modules_keeps_packed_buffers_on_gpu_until_serialization():
|
|
817
853
|
if not torch.cuda.is_available() or not available_backends()["triton_cuda"]:
|
|
818
854
|
pytest.skip("CUDA/Triton backend is not available")
|
|
@@ -120,6 +120,31 @@ def test_model_card_renders_native_validation_evidence_without_raw_records():
|
|
|
120
120
|
assert "original.metrics.jsonl" not in card
|
|
121
121
|
|
|
122
122
|
|
|
123
|
+
def test_model_card_renders_observed_quality_warning_from_benchmark_summary():
|
|
124
|
+
manifest = _manifest_for_model("Tongyi-MAI/Z-Image-Turbo", bits=(2, 3))
|
|
125
|
+
observed_quality = (
|
|
126
|
+
"The OrbitQuant W2A3 outputs show severe visual collapse across all ten "
|
|
127
|
+
"native comparison prompts. Use W3A3 or W4A4 for this model."
|
|
128
|
+
)
|
|
129
|
+
|
|
130
|
+
card = render_model_card(
|
|
131
|
+
manifest,
|
|
132
|
+
benchmark_summary={"observed_quality": observed_quality},
|
|
133
|
+
)
|
|
134
|
+
|
|
135
|
+
assert "## Observed Quality" in card
|
|
136
|
+
assert f"**Warning:** {observed_quality}" in card
|
|
137
|
+
|
|
138
|
+
|
|
139
|
+
def test_model_card_omits_observed_quality_without_a_nonempty_string():
|
|
140
|
+
manifest = _manifest_for_model("Tongyi-MAI/Z-Image-Turbo", bits=(2, 3))
|
|
141
|
+
|
|
142
|
+
for benchmark_summary in (None, {}, {"observed_quality": " "}):
|
|
143
|
+
card = render_model_card(manifest, benchmark_summary=benchmark_summary)
|
|
144
|
+
|
|
145
|
+
assert "## Observed Quality" not in card
|
|
146
|
+
|
|
147
|
+
|
|
123
148
|
def test_model_card_renders_imported_geneval_release_metrics():
|
|
124
149
|
manifest = _manifest_for_model("black-forest-labs/FLUX.1-schnell")
|
|
125
150
|
benchmark_summary = {
|
|
@@ -827,6 +827,47 @@ def test_orbit_linear_explicit_dequant_bf16_still_uses_reference_path(monkeypatc
|
|
|
827
827
|
assert calls == {"dequant": 1, "linear": 1}
|
|
828
828
|
|
|
829
829
|
|
|
830
|
+
def test_orbit_linear_rebuilds_derived_constants_after_to_empty():
|
|
831
|
+
torch.manual_seed(19)
|
|
832
|
+
source = torch.nn.Linear(16, 7, bias=True, dtype=torch.float32)
|
|
833
|
+
x = torch.randn(2, 5, 16)
|
|
834
|
+
config = OrbitQuantConfig(
|
|
835
|
+
weight_bits=4,
|
|
836
|
+
activation_bits=4,
|
|
837
|
+
rotation_seed=11,
|
|
838
|
+
block_size=8,
|
|
839
|
+
runtime_mode="dequant_bf16",
|
|
840
|
+
activation_kernel_backend="cpu",
|
|
841
|
+
)
|
|
842
|
+
reference = OrbitQuantLinear.from_linear(
|
|
843
|
+
source, config=config, module_name="block.ff.linear"
|
|
844
|
+
)
|
|
845
|
+
state = {name: tensor.clone() for name, tensor in reference.state_dict().items()}
|
|
846
|
+
restored = copy.deepcopy(reference)
|
|
847
|
+
restored.to_empty(device="cpu")
|
|
848
|
+
restored.load_state_dict(state)
|
|
849
|
+
|
|
850
|
+
assert not restored._derived_constants_valid
|
|
851
|
+
actual = restored(x)
|
|
852
|
+
|
|
853
|
+
assert restored._derived_constants_valid
|
|
854
|
+
assert torch.equal(restored._rotation_permutation, restored.rotation.permutation)
|
|
855
|
+
assert torch.equal(restored._rotation_signs, restored.rotation.signs)
|
|
856
|
+
assert torch.equal(
|
|
857
|
+
restored._activation_codebook_centroids,
|
|
858
|
+
restored.activation_codebook.centroids,
|
|
859
|
+
)
|
|
860
|
+
assert torch.equal(
|
|
861
|
+
restored._activation_codebook_boundaries,
|
|
862
|
+
restored.activation_codebook.boundaries,
|
|
863
|
+
)
|
|
864
|
+
assert torch.equal(
|
|
865
|
+
restored._weight_codebook_centroids,
|
|
866
|
+
restored.weight_codebook.centroids,
|
|
867
|
+
)
|
|
868
|
+
assert torch.allclose(actual, reference(x))
|
|
869
|
+
|
|
870
|
+
|
|
830
871
|
def test_orbit_linear_native_packed_matmul_mps_matches_dequant_bf16(monkeypatch):
|
|
831
872
|
if not torch.backends.mps.is_available():
|
|
832
873
|
pytest.skip("MPS backend is not available")
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{orbitquant-0.2.1 → orbitquant-0.2.2}/native-kernels/orbitquant-packed-matmul/tests/__init__.py
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|