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.
Files changed (128) hide show
  1. {orbitquant-0.2.1 → orbitquant-0.2.2}/.github/workflows/publish-pypi.yml +1 -1
  2. {orbitquant-0.2.1 → orbitquant-0.2.2}/PKG-INFO +1 -1
  3. {orbitquant-0.2.1 → orbitquant-0.2.2}/docs/kernel-audit.md +21 -0
  4. {orbitquant-0.2.1 → orbitquant-0.2.2}/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cuda/packed_matmul.cu +21 -22
  5. {orbitquant-0.2.1 → orbitquant-0.2.2}/pyproject.toml +1 -1
  6. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/__init__.py +1 -1
  7. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/artifacts/model_card.py +18 -0
  8. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/kernels/triton_cuda.py +114 -8
  9. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/layers.py +34 -0
  10. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/modeling.py +3 -1
  11. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/pipeline.py +2 -2
  12. {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_diffusers_modelmixin_integration.py +18 -0
  13. {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_kernels.py +36 -0
  14. {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_model_card.py +25 -0
  15. {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_orbit_linear.py +41 -0
  16. {orbitquant-0.2.1 → orbitquant-0.2.2}/uv.lock +1 -1
  17. {orbitquant-0.2.1 → orbitquant-0.2.2}/.github/workflows/ci.yml +0 -0
  18. {orbitquant-0.2.1 → orbitquant-0.2.2}/.gitignore +0 -0
  19. {orbitquant-0.2.1 → orbitquant-0.2.2}/LICENSE +0 -0
  20. {orbitquant-0.2.1 → orbitquant-0.2.2}/README.md +0 -0
  21. {orbitquant-0.2.1 → orbitquant-0.2.2}/docs/paper-methodology-audit.md +0 -0
  22. {orbitquant-0.2.1 → orbitquant-0.2.2}/native-kernels/orbitquant-packed-matmul/CARD.md +0 -0
  23. {orbitquant-0.2.1 → orbitquant-0.2.2}/native-kernels/orbitquant-packed-matmul/benchmarks/benchmark.py +0 -0
  24. {orbitquant-0.2.1 → orbitquant-0.2.2}/native-kernels/orbitquant-packed-matmul/build.toml +0 -0
  25. {orbitquant-0.2.1 → orbitquant-0.2.2}/native-kernels/orbitquant-packed-matmul/example.py +0 -0
  26. {orbitquant-0.2.1 → orbitquant-0.2.2}/native-kernels/orbitquant-packed-matmul/flake.lock +0 -0
  27. {orbitquant-0.2.1 → orbitquant-0.2.2}/native-kernels/orbitquant-packed-matmul/flake.nix +0 -0
  28. {orbitquant-0.2.1 → orbitquant-0.2.2}/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_metal/packed_matmul.metal +0 -0
  29. {orbitquant-0.2.1 → orbitquant-0.2.2}/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_metal/packed_matmul.mm +0 -0
  30. {orbitquant-0.2.1 → orbitquant-0.2.2}/native-kernels/orbitquant-packed-matmul/tests/__init__.py +0 -0
  31. {orbitquant-0.2.1 → orbitquant-0.2.2}/native-kernels/orbitquant-packed-matmul/tests/test_packed_matmul.py +0 -0
  32. {orbitquant-0.2.1 → orbitquant-0.2.2}/native-kernels/orbitquant-packed-matmul/torch-ext/orbitquant_packed_matmul/__init__.py +0 -0
  33. {orbitquant-0.2.1 → orbitquant-0.2.2}/native-kernels/orbitquant-packed-matmul/torch-ext/torch_binding.cpp +0 -0
  34. {orbitquant-0.2.1 → orbitquant-0.2.2}/native-kernels/orbitquant-packed-matmul/torch-ext/torch_binding.h +0 -0
  35. {orbitquant-0.2.1 → orbitquant-0.2.2}/scripts/run_cuda_kernel_checks.sh +0 -0
  36. {orbitquant-0.2.1 → orbitquant-0.2.2}/scripts/run_hf_compat_checks.sh +0 -0
  37. {orbitquant-0.2.1 → orbitquant-0.2.2}/scripts/run_mps_kernel_checks.sh +0 -0
  38. {orbitquant-0.2.1 → orbitquant-0.2.2}/scripts/run_paper_methodology_checks.sh +0 -0
  39. {orbitquant-0.2.1 → orbitquant-0.2.2}/scripts/runpod_ssh_health.sh +0 -0
  40. {orbitquant-0.2.1 → orbitquant-0.2.2}/scripts/verify_hf_kernel_model_artifact.py +0 -0
  41. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/adaln.py +0 -0
  42. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/artifacts/__init__.py +0 -0
  43. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/artifacts/assets.py +0 -0
  44. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/artifacts/benchmark.py +0 -0
  45. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/artifacts/checksums.py +0 -0
  46. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/artifacts/comparisons.py +0 -0
  47. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/artifacts/loader.py +0 -0
  48. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/artifacts/manifest.py +0 -0
  49. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/artifacts/refresh.py +0 -0
  50. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/artifacts/repair.py +0 -0
  51. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/artifacts/validator.py +0 -0
  52. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/artifacts/writer.py +0 -0
  53. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/benchmarks.py +0 -0
  54. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/cli/__init__.py +0 -0
  55. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/cli/main.py +0 -0
  56. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/codebooks/__init__.py +0 -0
  57. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/codebooks/lloyd_max.py +0 -0
  58. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/config.py +0 -0
  59. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/errors.py +0 -0
  60. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/eval/__init__.py +0 -0
  61. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/eval/assets.py +0 -0
  62. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/eval/external_export.py +0 -0
  63. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/eval/external_metrics.py +0 -0
  64. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/eval/external_plan.py +0 -0
  65. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/eval/metrics.py +0 -0
  66. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/eval/native_plan.py +0 -0
  67. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/eval/native_runner.py +0 -0
  68. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/eval/native_settings.py +0 -0
  69. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/eval/prompts.py +0 -0
  70. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/eval/report.py +0 -0
  71. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/functional.py +0 -0
  72. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/hub.py +0 -0
  73. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/kernels/__init__.py +0 -0
  74. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/kernels/dispatch.py +0 -0
  75. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/kernels/mps.py +0 -0
  76. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/kernels/native_packed_matmul.py +0 -0
  77. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/linear_adapters.py +0 -0
  78. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/packing/__init__.py +0 -0
  79. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/packing/bitpack.py +0 -0
  80. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/policies/__init__.py +0 -0
  81. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/policies/flux.py +0 -0
  82. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/policies/flux2.py +0 -0
  83. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/policies/generic_dit.py +0 -0
  84. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/policies/wan.py +0 -0
  85. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/policies/z_image.py +0 -0
  86. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/quantizer.py +0 -0
  87. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/recipes.py +0 -0
  88. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/rotations/__init__.py +0 -0
  89. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/rotations/fwht.py +0 -0
  90. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/rotations/rpbh.py +0 -0
  91. {orbitquant-0.2.1 → orbitquant-0.2.2}/src/orbitquant/transformers_ops.py +0 -0
  92. {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_adaln_rtn.py +0 -0
  93. {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_artifact_writer.py +0 -0
  94. {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_benchmarks.py +0 -0
  95. {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_bitpack.py +0 -0
  96. {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_ci.py +0 -0
  97. {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_cli.py +0 -0
  98. {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_cli_report.py +0 -0
  99. {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_codebooks.py +0 -0
  100. {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_config.py +0 -0
  101. {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_distribution.py +0 -0
  102. {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_documentation.py +0 -0
  103. {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_eval_assets.py +0 -0
  104. {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_eval_prompts.py +0 -0
  105. {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_eval_report.py +0 -0
  106. {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_external_eval_plan.py +0 -0
  107. {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_external_export.py +0 -0
  108. {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_hub.py +0 -0
  109. {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_kernel_model_verifier.py +0 -0
  110. {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_linear_adapters.py +0 -0
  111. {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_manifest.py +0 -0
  112. {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_model_quantization.py +0 -0
  113. {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_native_kernel_package.py +0 -0
  114. {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_native_packed_matmul.py +0 -0
  115. {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_native_runner.py +0 -0
  116. {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_native_settings.py +0 -0
  117. {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_optional_dependencies.py +0 -0
  118. {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_paper_alignment.py +0 -0
  119. {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_paper_methodology.py +0 -0
  120. {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_pipeline_helpers.py +0 -0
  121. {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_policies.py +0 -0
  122. {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_prompt_jobs.py +0 -0
  123. {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_prompt_selection.py +0 -0
  124. {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_quantizer_adapter.py +0 -0
  125. {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_rpbh.py +0 -0
  126. {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_target_policies.py +0 -0
  127. {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_transformers_pretrained_integration.py +0 -0
  128. {orbitquant-0.2.1 → orbitquant-0.2.2}/tests/test_universal_transformers.py +0 -0
@@ -6,7 +6,7 @@ on:
6
6
  version:
7
7
  description: Version to publish
8
8
  required: true
9
- default: "0.2.1"
9
+ default: "0.2.2"
10
10
 
11
11
  jobs:
12
12
  build:
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: orbitquant
3
- Version: 0.2.1
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
- uint32_t packed_words[Bits] = {};
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 < Bits; ++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 < values; ++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 = 64;
115
- constexpr int tile_n = 64;
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 = 4;
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
- const int weight_task = threadIdx.x;
164
- if (weight_task < weight_tasks) {
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 = 64;
652
- constexpr int mma_tile_n = 64;
653
- const dim3 mma_block(128);
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 = 64;
744
- constexpr int mma_tile_n = 64;
745
- const dim3 mma_block(128);
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);
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
4
4
 
5
5
  [project]
6
6
  name = "orbitquant"
7
- version = "0.2.1"
7
+ version = "0.2.2"
8
8
  description = "Calibration-free OrbitQuant for transformer linear projections"
9
9
  readme = "README.md"
10
10
  requires-python = ">=3.11"
@@ -18,7 +18,7 @@ from orbitquant.pipeline import (
18
18
  from orbitquant.quantizer import register_hf_quantizers
19
19
  from orbitquant.recipes import recipe
20
20
 
21
- __version__ = "0.2.1"
21
+ __version__ = "0.2.2"
22
22
 
23
23
  register_hf_quantizers()
24
24
 
@@ -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
- fwht_stages = orbit_block_size.bit_length() - 1
872
- _fused_rpbh_quantize_activation_kernel[(rows * num_blocks,)](
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
- centroids,
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
- fwht_stages=fwht_stages,
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
- num_warps=8 if orbit_block_size >= 512 else 4,
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 model_config is not None:
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, quantize_linear_modules
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 quantize_linear_modules(
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")
@@ -1095,7 +1095,7 @@ wheels = [
1095
1095
 
1096
1096
  [[package]]
1097
1097
  name = "orbitquant"
1098
- version = "0.2.1"
1098
+ version = "0.2.2"
1099
1099
  source = { editable = "." }
1100
1100
  dependencies = [
1101
1101
  { name = "huggingface-hub" },
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes