orbitquant 0.2.0__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.0 → orbitquant-0.2.2}/.github/workflows/publish-pypi.yml +1 -1
  2. {orbitquant-0.2.0 → orbitquant-0.2.2}/PKG-INFO +4 -3
  3. {orbitquant-0.2.0 → orbitquant-0.2.2}/README.md +3 -2
  4. {orbitquant-0.2.0 → orbitquant-0.2.2}/docs/kernel-audit.md +21 -0
  5. {orbitquant-0.2.0 → orbitquant-0.2.2}/docs/paper-methodology-audit.md +3 -1
  6. {orbitquant-0.2.0 → orbitquant-0.2.2}/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_cuda/packed_matmul.cu +21 -22
  7. {orbitquant-0.2.0 → orbitquant-0.2.2}/pyproject.toml +1 -1
  8. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/__init__.py +1 -1
  9. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/artifacts/model_card.py +96 -5
  10. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/eval/prompts.py +89 -25
  11. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/eval/report.py +8 -5
  12. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/kernels/triton_cuda.py +114 -8
  13. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/layers.py +34 -0
  14. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/modeling.py +3 -1
  15. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/pipeline.py +2 -2
  16. {orbitquant-0.2.0 → orbitquant-0.2.2}/tests/test_artifact_writer.py +3 -1
  17. {orbitquant-0.2.0 → orbitquant-0.2.2}/tests/test_cli.py +8 -8
  18. {orbitquant-0.2.0 → orbitquant-0.2.2}/tests/test_diffusers_modelmixin_integration.py +18 -0
  19. {orbitquant-0.2.0 → orbitquant-0.2.2}/tests/test_eval_prompts.py +7 -1
  20. {orbitquant-0.2.0 → orbitquant-0.2.2}/tests/test_eval_report.py +4 -0
  21. {orbitquant-0.2.0 → orbitquant-0.2.2}/tests/test_kernels.py +36 -0
  22. {orbitquant-0.2.0 → orbitquant-0.2.2}/tests/test_model_card.py +61 -1
  23. {orbitquant-0.2.0 → orbitquant-0.2.2}/tests/test_orbit_linear.py +41 -0
  24. {orbitquant-0.2.0 → orbitquant-0.2.2}/uv.lock +1 -1
  25. {orbitquant-0.2.0 → orbitquant-0.2.2}/.github/workflows/ci.yml +0 -0
  26. {orbitquant-0.2.0 → orbitquant-0.2.2}/.gitignore +0 -0
  27. {orbitquant-0.2.0 → orbitquant-0.2.2}/LICENSE +0 -0
  28. {orbitquant-0.2.0 → orbitquant-0.2.2}/native-kernels/orbitquant-packed-matmul/CARD.md +0 -0
  29. {orbitquant-0.2.0 → orbitquant-0.2.2}/native-kernels/orbitquant-packed-matmul/benchmarks/benchmark.py +0 -0
  30. {orbitquant-0.2.0 → orbitquant-0.2.2}/native-kernels/orbitquant-packed-matmul/build.toml +0 -0
  31. {orbitquant-0.2.0 → orbitquant-0.2.2}/native-kernels/orbitquant-packed-matmul/example.py +0 -0
  32. {orbitquant-0.2.0 → orbitquant-0.2.2}/native-kernels/orbitquant-packed-matmul/flake.lock +0 -0
  33. {orbitquant-0.2.0 → orbitquant-0.2.2}/native-kernels/orbitquant-packed-matmul/flake.nix +0 -0
  34. {orbitquant-0.2.0 → orbitquant-0.2.2}/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_metal/packed_matmul.metal +0 -0
  35. {orbitquant-0.2.0 → orbitquant-0.2.2}/native-kernels/orbitquant-packed-matmul/orbitquant_packed_matmul_metal/packed_matmul.mm +0 -0
  36. {orbitquant-0.2.0 → orbitquant-0.2.2}/native-kernels/orbitquant-packed-matmul/tests/__init__.py +0 -0
  37. {orbitquant-0.2.0 → orbitquant-0.2.2}/native-kernels/orbitquant-packed-matmul/tests/test_packed_matmul.py +0 -0
  38. {orbitquant-0.2.0 → orbitquant-0.2.2}/native-kernels/orbitquant-packed-matmul/torch-ext/orbitquant_packed_matmul/__init__.py +0 -0
  39. {orbitquant-0.2.0 → orbitquant-0.2.2}/native-kernels/orbitquant-packed-matmul/torch-ext/torch_binding.cpp +0 -0
  40. {orbitquant-0.2.0 → orbitquant-0.2.2}/native-kernels/orbitquant-packed-matmul/torch-ext/torch_binding.h +0 -0
  41. {orbitquant-0.2.0 → orbitquant-0.2.2}/scripts/run_cuda_kernel_checks.sh +0 -0
  42. {orbitquant-0.2.0 → orbitquant-0.2.2}/scripts/run_hf_compat_checks.sh +0 -0
  43. {orbitquant-0.2.0 → orbitquant-0.2.2}/scripts/run_mps_kernel_checks.sh +0 -0
  44. {orbitquant-0.2.0 → orbitquant-0.2.2}/scripts/run_paper_methodology_checks.sh +0 -0
  45. {orbitquant-0.2.0 → orbitquant-0.2.2}/scripts/runpod_ssh_health.sh +0 -0
  46. {orbitquant-0.2.0 → orbitquant-0.2.2}/scripts/verify_hf_kernel_model_artifact.py +0 -0
  47. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/adaln.py +0 -0
  48. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/artifacts/__init__.py +0 -0
  49. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/artifacts/assets.py +0 -0
  50. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/artifacts/benchmark.py +0 -0
  51. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/artifacts/checksums.py +0 -0
  52. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/artifacts/comparisons.py +0 -0
  53. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/artifacts/loader.py +0 -0
  54. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/artifacts/manifest.py +0 -0
  55. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/artifacts/refresh.py +0 -0
  56. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/artifacts/repair.py +0 -0
  57. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/artifacts/validator.py +0 -0
  58. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/artifacts/writer.py +0 -0
  59. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/benchmarks.py +0 -0
  60. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/cli/__init__.py +0 -0
  61. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/cli/main.py +0 -0
  62. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/codebooks/__init__.py +0 -0
  63. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/codebooks/lloyd_max.py +0 -0
  64. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/config.py +0 -0
  65. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/errors.py +0 -0
  66. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/eval/__init__.py +0 -0
  67. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/eval/assets.py +0 -0
  68. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/eval/external_export.py +0 -0
  69. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/eval/external_metrics.py +0 -0
  70. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/eval/external_plan.py +0 -0
  71. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/eval/metrics.py +0 -0
  72. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/eval/native_plan.py +0 -0
  73. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/eval/native_runner.py +0 -0
  74. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/eval/native_settings.py +0 -0
  75. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/functional.py +0 -0
  76. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/hub.py +0 -0
  77. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/kernels/__init__.py +0 -0
  78. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/kernels/dispatch.py +0 -0
  79. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/kernels/mps.py +0 -0
  80. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/kernels/native_packed_matmul.py +0 -0
  81. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/linear_adapters.py +0 -0
  82. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/packing/__init__.py +0 -0
  83. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/packing/bitpack.py +0 -0
  84. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/policies/__init__.py +0 -0
  85. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/policies/flux.py +0 -0
  86. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/policies/flux2.py +0 -0
  87. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/policies/generic_dit.py +0 -0
  88. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/policies/wan.py +0 -0
  89. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/policies/z_image.py +0 -0
  90. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/quantizer.py +0 -0
  91. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/recipes.py +0 -0
  92. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/rotations/__init__.py +0 -0
  93. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/rotations/fwht.py +0 -0
  94. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/rotations/rpbh.py +0 -0
  95. {orbitquant-0.2.0 → orbitquant-0.2.2}/src/orbitquant/transformers_ops.py +0 -0
  96. {orbitquant-0.2.0 → orbitquant-0.2.2}/tests/test_adaln_rtn.py +0 -0
  97. {orbitquant-0.2.0 → orbitquant-0.2.2}/tests/test_benchmarks.py +0 -0
  98. {orbitquant-0.2.0 → orbitquant-0.2.2}/tests/test_bitpack.py +0 -0
  99. {orbitquant-0.2.0 → orbitquant-0.2.2}/tests/test_ci.py +0 -0
  100. {orbitquant-0.2.0 → orbitquant-0.2.2}/tests/test_cli_report.py +0 -0
  101. {orbitquant-0.2.0 → orbitquant-0.2.2}/tests/test_codebooks.py +0 -0
  102. {orbitquant-0.2.0 → orbitquant-0.2.2}/tests/test_config.py +0 -0
  103. {orbitquant-0.2.0 → orbitquant-0.2.2}/tests/test_distribution.py +0 -0
  104. {orbitquant-0.2.0 → orbitquant-0.2.2}/tests/test_documentation.py +0 -0
  105. {orbitquant-0.2.0 → orbitquant-0.2.2}/tests/test_eval_assets.py +0 -0
  106. {orbitquant-0.2.0 → orbitquant-0.2.2}/tests/test_external_eval_plan.py +0 -0
  107. {orbitquant-0.2.0 → orbitquant-0.2.2}/tests/test_external_export.py +0 -0
  108. {orbitquant-0.2.0 → orbitquant-0.2.2}/tests/test_hub.py +0 -0
  109. {orbitquant-0.2.0 → orbitquant-0.2.2}/tests/test_kernel_model_verifier.py +0 -0
  110. {orbitquant-0.2.0 → orbitquant-0.2.2}/tests/test_linear_adapters.py +0 -0
  111. {orbitquant-0.2.0 → orbitquant-0.2.2}/tests/test_manifest.py +0 -0
  112. {orbitquant-0.2.0 → orbitquant-0.2.2}/tests/test_model_quantization.py +0 -0
  113. {orbitquant-0.2.0 → orbitquant-0.2.2}/tests/test_native_kernel_package.py +0 -0
  114. {orbitquant-0.2.0 → orbitquant-0.2.2}/tests/test_native_packed_matmul.py +0 -0
  115. {orbitquant-0.2.0 → orbitquant-0.2.2}/tests/test_native_runner.py +0 -0
  116. {orbitquant-0.2.0 → orbitquant-0.2.2}/tests/test_native_settings.py +0 -0
  117. {orbitquant-0.2.0 → orbitquant-0.2.2}/tests/test_optional_dependencies.py +0 -0
  118. {orbitquant-0.2.0 → orbitquant-0.2.2}/tests/test_paper_alignment.py +0 -0
  119. {orbitquant-0.2.0 → orbitquant-0.2.2}/tests/test_paper_methodology.py +0 -0
  120. {orbitquant-0.2.0 → orbitquant-0.2.2}/tests/test_pipeline_helpers.py +0 -0
  121. {orbitquant-0.2.0 → orbitquant-0.2.2}/tests/test_policies.py +0 -0
  122. {orbitquant-0.2.0 → orbitquant-0.2.2}/tests/test_prompt_jobs.py +0 -0
  123. {orbitquant-0.2.0 → orbitquant-0.2.2}/tests/test_prompt_selection.py +0 -0
  124. {orbitquant-0.2.0 → orbitquant-0.2.2}/tests/test_quantizer_adapter.py +0 -0
  125. {orbitquant-0.2.0 → orbitquant-0.2.2}/tests/test_rpbh.py +0 -0
  126. {orbitquant-0.2.0 → orbitquant-0.2.2}/tests/test_target_policies.py +0 -0
  127. {orbitquant-0.2.0 → orbitquant-0.2.2}/tests/test_transformers_pretrained_integration.py +0 -0
  128. {orbitquant-0.2.0 → 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.0"
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.0
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
@@ -232,8 +232,9 @@ save_quantized_pipeline_component(
232
232
  ```
233
233
 
234
234
  Published FLUX, Z-Image, and Wan repositories are compact Diffusers component
235
- artifacts. Their model cards contain the matching pipeline code and native
236
- generation settings.
235
+ artifacts. Image model cards contain the matching pipeline code, native
236
+ generation settings, and a ten-prompt full-resolution BF16-vs-OrbitQuant
237
+ comparison matrix.
237
238
 
238
239
  Load a published component artifact together with its recorded source pipeline:
239
240
 
@@ -188,8 +188,9 @@ save_quantized_pipeline_component(
188
188
  ```
189
189
 
190
190
  Published FLUX, Z-Image, and Wan repositories are compact Diffusers component
191
- artifacts. Their model cards contain the matching pipeline code and native
192
- generation settings.
191
+ artifacts. Image model cards contain the matching pipeline code, native
192
+ generation settings, and a ten-prompt full-resolution BF16-vs-OrbitQuant
193
+ comparison matrix.
193
194
 
194
195
  Load a published component artifact together with its recorded source pipeline:
195
196
 
@@ -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
 
@@ -143,7 +143,9 @@ Native artifact readiness is separate from full GenEval or VBench scoring.
143
143
  Those runs prove metric-table claims. The required artifact-readiness evidence
144
144
  is:
145
145
 
146
- - Native-resolution BF16-vs-OrbitQuant comparison matrix.
146
+ - Ten-prompt native-resolution BF16-vs-OrbitQuant comparison matrix for image
147
+ artifacts, including dense composition, fine detail, artistic style, and
148
+ Latin, Russian, Japanese, and Chinese typography stress cases.
147
149
  - Same prompt and seed for BF16 and OrbitQuant.
148
150
  - Native settings from `src/orbitquant/eval/native_settings.py`.
149
151
  - Finite, nonblank output checks.
@@ -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.0"
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.0"
21
+ __version__ = "0.2.2"
22
22
 
23
23
  register_hf_quantizers()
24
24
 
@@ -4,6 +4,7 @@ from typing import Any
4
4
 
5
5
  from orbitquant.artifacts.manifest import OrbitQuantManifest
6
6
  from orbitquant.eval.native_settings import NativeSuite, list_native_suites
7
+ from orbitquant.eval.prompts import default_prompt_payload
7
8
 
8
9
  _GENEVAL_METRICS = (
9
10
  "geneval_overall",
@@ -65,7 +66,7 @@ def _install_snippet() -> str:
65
66
  return "\n".join(
66
67
  [
67
68
  "```bash",
68
- "pip install \"orbitquant[hf]\"",
69
+ "pip install \"orbitquant[hf,kernels]>=0.2.0\"",
69
70
  "```",
70
71
  ]
71
72
  )
@@ -343,6 +344,86 @@ def _native_validation_proof_section(
343
344
  return lines
344
345
 
345
346
 
347
+ def _comparison_prompt_section(
348
+ source_model_id: str,
349
+ benchmark_summary: dict[str, Any] | None,
350
+ ) -> list[str]:
351
+ suite = _native_suite_for_source_model(source_model_id)
352
+ if suite is None or suite.frames is not None or not isinstance(benchmark_summary, dict):
353
+ return []
354
+ proof = benchmark_summary.get("native_smoke")
355
+ if not isinstance(proof, dict):
356
+ return []
357
+ prompt_payload = default_prompt_payload(
358
+ {
359
+ "flux2-native": "flux2",
360
+ "flux1-schnell-native": "flux",
361
+ "z-image-native": "z_image",
362
+ }[suite.name]
363
+ )
364
+ prompts = prompt_payload["prompts"]
365
+ if proof.get("paired_prompt_seed_count") != len(prompts):
366
+ return []
367
+ paired_keys = proof.get("paired_prompt_seed_keys")
368
+ if not isinstance(paired_keys, list):
369
+ return []
370
+ paired_prompt_ids = {
371
+ str(item[2])
372
+ for item in paired_keys
373
+ if isinstance(item, list) and len(item) == 3
374
+ }
375
+ if paired_prompt_ids != {str(prompt["id"]) for prompt in prompts}:
376
+ return []
377
+
378
+ required_text = {
379
+ "english-text-rendering": "ORBIT QUANT; DATA WITHOUT CALIBRATION",
380
+ "cyrillic-text-rendering": "КВАНТОВАЯ ОРБИТА; МОСКВА 2049; КВАНТОВАНИЕ",
381
+ "style-heavy": "量子の軌道; 東京の未来",
382
+ "occlusion-reflection": "量子轨道; 未来之城",
383
+ }
384
+ lines = [
385
+ "## Comparison Prompt Set",
386
+ "",
387
+ f"The matrix uses all ten prompts from `{prompt_payload['prompt_pack']}` at "
388
+ f"`{suite.width}x{suite.height}`. BF16 and OrbitQuant use the same seed for "
389
+ "each row.",
390
+ "",
391
+ "| # | Stress case | Exact required text |",
392
+ "| ---: | --- | --- |",
393
+ ]
394
+ for index, prompt in enumerate(prompts, start=1):
395
+ text = required_text.get(str(prompt["id"]), "-")
396
+ lines.append(f"| {index} | {prompt['title']} | {text} |")
397
+ lines.extend(
398
+ [
399
+ "",
400
+ "<details>",
401
+ "<summary>Exact comparison prompts</summary>",
402
+ "",
403
+ ]
404
+ )
405
+ for index, prompt in enumerate(prompts, start=1):
406
+ lines.append(f"{index}. **{prompt['title']}**: {prompt['prompt']}")
407
+ lines.extend(["", "</details>", ""])
408
+ return lines
409
+
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
+
346
427
  def render_model_card(
347
428
  manifest: OrbitQuantManifest,
348
429
  *,
@@ -356,6 +437,10 @@ def render_model_card(
356
437
  data["source_model_id"], benchmark_summary
357
438
  )
358
439
  native_validation_proof_lines = _native_validation_proof_section(benchmark_summary)
440
+ comparison_prompt_lines = _comparison_prompt_section(
441
+ data["source_model_id"], benchmark_summary
442
+ )
443
+ observed_quality_lines = _observed_quality_section(benchmark_summary)
359
444
  adaln_group_size = int(data.get("adaln_group_size", 64))
360
445
  adaln_default_note = (
361
446
  "- AdaLN group-size note: paper default."
@@ -425,14 +510,19 @@ def render_model_card(
425
510
  _usage_snippet(data["source_model_id"], bits),
426
511
  "",
427
512
  "`runtime_mode=\"auto_fused\"` is the default optimized runtime. On "
428
- "CUDA it tries the native packed low-bit matmul kernel first, then "
429
- "Triton packed matmul. On MPS it requires the native Metal packed "
430
- "matmul kernel. Use `runtime_mode=\"dequant_bf16\"` only as an "
431
- "explicit compatibility/debug reference path.",
513
+ "CUDA, the `kernels` extra provides the Triton packed fallback; a "
514
+ "locally built native CUDA package is preferred automatically when "
515
+ "installed. On MPS, build and install the native Metal package from "
516
+ "the OrbitQuant source tree. See the [OrbitQuant runtime instructions]"
517
+ "(https://github.com/iamwavecut/OrbitQuant/blob/main/docs/"
518
+ "kernel-audit.md#local-native-package). Use "
519
+ "`runtime_mode=\"dequant_bf16\"` only as an explicit compatibility/"
520
+ "debug reference path.",
432
521
  "",
433
522
  *native_settings_lines,
434
523
  *validation_status_lines,
435
524
  *native_validation_proof_lines,
525
+ *comparison_prompt_lines,
436
526
  "## Quantization",
437
527
  "",
438
528
  f"- Method: `{data['quant_method']}`",
@@ -458,6 +548,7 @@ def render_model_card(
458
548
  "- Calibration data: none",
459
549
  "- Text encoders and VAE: left in source precision by default",
460
550
  "",
551
+ *observed_quality_lines,
461
552
  *comparison_lines,
462
553
  "## Source",
463
554
  "",
@@ -8,62 +8,126 @@ from typing import Any
8
8
  IMAGE_VISUAL_PROMPTS = [
9
9
  {
10
10
  "id": "simple-object",
11
- "category": "simple_object",
12
- "prompt": "A red ceramic mug on a wooden desk, soft daylight, shallow depth of field",
11
+ "title": "01 Fine-detail astrolabe",
12
+ "category": "fine_detail",
13
+ "prompt": (
14
+ "A museum-grade macro photograph of a single ornate brass astronomical clock, "
15
+ "covered with interlocking gears, engraved constellations, enamel moon phases, "
16
+ "hair-thin hands, tiny screws, worn gilt edges, and dust caught in the mechanisms; "
17
+ "dramatic Rembrandt lighting, black velvet background, extreme material detail, "
18
+ "large-format photography, every cog mechanically coherent"
19
+ ),
13
20
  },
14
21
  {
15
22
  "id": "two-object-composition",
16
- "category": "two_object_composition",
17
- "prompt": "Two robots playing chess in a quiet library, one silver and one matte black",
23
+ "title": "02 Layered character composition",
24
+ "category": "layered_character_composition",
25
+ "prompt": (
26
+ "A lacquered white android and an elderly watchmaker facing each other across a "
27
+ "crowded Art Nouveau workshop, jointly repairing a mechanical hummingbird; the "
28
+ "android is on the left, the watchmaker on the right, the bird centered between "
29
+ "them, hundreds of tools and clock parts in the midground, rain and a tram visible "
30
+ "through the window, coherent mirror reflections, cinematic tungsten and cyan light, "
31
+ "intricate faces and hands, editorial realism"
32
+ ),
18
33
  },
19
34
  {
20
35
  "id": "counting",
21
- "category": "counting",
22
- "prompt": "Five glass marbles arranged in a straight line on white paper",
36
+ "title": "03 Exact counting and choreography",
37
+ "category": "counting_and_choreography",
38
+ "prompt": (
39
+ "Exactly seven masked dancers performing on seven separate illuminated platforms "
40
+ "inside a vast flooded opera house, no extra dancers; alternating crimson and ivory "
41
+ "costumes from left to right, each dancer in a distinct pose, balconies reflected in "
42
+ "the water, floating candles, volumetric stage haze, baroque theatrical photography, "
43
+ "sharp foreground and readable background architecture"
44
+ ),
23
45
  },
24
46
  {
25
47
  "id": "color-binding",
26
- "category": "color_binding",
27
- "prompt": "A blue cube beside a yellow sphere and a green pyramid on a gray table",
48
+ "title": "04 Dense color and object binding",
49
+ "category": "dense_color_binding",
50
+ "prompt": (
51
+ "An elaborate surreal fashion tableau with three models: the left model wears a "
52
+ "cobalt-blue coat and holds a yellow glass sphere, the center model wears a saffron "
53
+ "dress and holds a green ceramic pyramid, the right model wears an emerald suit and "
54
+ "holds a red velvet cube; rococo greenhouse, rare orchids, patterned tile floor, "
55
+ "prismatic sunlight, magazine-cover precision, preserve every color-object pairing"
56
+ ),
28
57
  },
29
58
  {
30
59
  "id": "spatial-relationship",
31
- "category": "spatial_relationship",
32
- "prompt": "A blue cube to the left of a yellow sphere, both centered in the frame",
60
+ "title": "05 Nested spatial relationships",
61
+ "category": "nested_spatial_relationships",
62
+ "prompt": (
63
+ "A meticulous cutaway diorama of a vertical city: a glass greenhouse sits directly "
64
+ "above a silver subway car, a red fox stands inside the greenhouse beneath a hanging "
65
+ "moon lamp, a violinist waits below the subway platform, and a yellow airship passes "
66
+ "behind the entire structure; isometric perspective, architectural-section drawing "
67
+ "mixed with photoreal materials, dozens of tiny rooms, stairs and people, clean depth"
68
+ ),
33
69
  },
34
70
  {
35
71
  "id": "long-prompt",
36
- "category": "long_prompt",
72
+ "title": "06 Cinematic night-market panorama",
73
+ "category": "cinematic_panorama",
37
74
  "prompt": (
38
- "A precise editorial photo of a compact robotics workbench with labeled tools, "
39
- "a small oscilloscope, coiled cables, a half-assembled drone, and a warm desk lamp; "
40
- "the scene should remain readable without clutter"
75
+ "A sweeping cinematic panorama of a rain-soaked floating night market at blue hour: "
76
+ "in the foreground a chef plates translucent dumplings under a red silk canopy; in "
77
+ "the midground children chase paper lanterns across narrow bridges while merchants "
78
+ "unload exotic fruit from wooden boats; in the background a terraced megacity rises "
79
+ "through mist beneath a storm, with hundreds of warm windows, wet reflections, steam, "
80
+ "umbrellas, ropes and signage; deep focus, anamorphic highlights, realistic faces, "
81
+ "coherent perspective, painterly color grading with documentary-level detail"
41
82
  ),
42
83
  },
43
84
  {
44
85
  "id": "english-text-rendering",
45
- "category": "text_rendering",
46
- "prompt": 'A clean street sign with the exact text "ORBIT QUANT"',
86
+ "title": "07 Editorial Latin typography",
87
+ "category": "latin_typography",
88
+ "prompt": (
89
+ "A sophisticated Swiss International Style exhibition poster photographed behind "
90
+ "slightly reflective museum glass, with the exact large headline \"ORBIT QUANT\" and "
91
+ "the exact smaller subtitle \"DATA WITHOUT CALIBRATION\"; strict modular grid, red, "
92
+ "black and white screenprint, tiny registration marks, embossed paper fibers, sharp "
93
+ "letterforms, dramatic gallery shadows, no other text"
94
+ ),
47
95
  },
48
96
  {
49
97
  "id": "cyrillic-text-rendering",
50
- "category": "text_rendering",
51
- "prompt": 'A handwritten note in Cyrillic that says "КВАНТОВАНИЕ"',
98
+ "title": "08 Russian Constructivist typography",
99
+ "category": "russian_typography",
100
+ "prompt": (
101
+ "A richly detailed Russian Constructivist science-fiction poster with the exact "
102
+ "Cyrillic headline \"КВАНТОВАЯ ОРБИТА\", the exact subtitle \"МОСКВА 2049\", and a "
103
+ "small exact stamp \"КВАНТОВАНИЕ\"; diagonal red and black geometry, cream paper, "
104
+ "cosmonaut portrait, orbital diagrams, halftone grain, folded corners, layered ink, "
105
+ "museum archival photograph, all letters crisp and correctly ordered"
106
+ ),
52
107
  },
53
108
  {
54
109
  "id": "style-heavy",
55
- "category": "style",
110
+ "title": "09 Japanese typography and mixed style",
111
+ "category": "japanese_typography",
56
112
  "prompt": (
57
- "A crowded train platform at sunset, documentary photo style, high dynamic range, "
58
- "natural grain, realistic faces, detailed architecture"
113
+ "An elaborate Japanese art magazine cover combining Edo woodblock printing with a "
114
+ "futuristic Tokyo skyline, with the exact vertical title \"量子の軌道\" and the exact "
115
+ "subtitle \"東京の未来\"; giant indigo waves curl around glass towers, red-crowned "
116
+ "cranes cross a gold moon, tiny pedestrians and trains fill the lower streets, visible "
117
+ "washi fibers, layered spot colors, precise Japanese glyphs, balanced editorial layout"
59
118
  ),
60
119
  },
61
120
  {
62
121
  "id": "occlusion-reflection",
63
- "category": "occlusion_reflection",
122
+ "title": "10 Chinese typography, reflection, occlusion",
123
+ "category": "chinese_typography_reflection_occlusion",
64
124
  "prompt": (
65
- "A reflective chrome teapot partly hidden behind orange flowers, with the room "
66
- "visible as a coherent reflection on the metal"
125
+ "A luxurious Chinese retro-futurist department-store window at night with the exact "
126
+ "gold title \"量子轨道\" and the exact red subtitle \"未来之城\"; a curved chrome "
127
+ "robot "
128
+ "is partly occluded by peonies and blue-and-white porcelain, the calligraphy and neon "
129
+ "street must appear coherently reflected across its body, multiple glass layers, silk "
130
+ "textures, passing bicycles, cinematic rain, fine product-photography detail"
67
131
  ),
68
132
  },
69
133
  ]
@@ -152,7 +216,7 @@ def default_prompt_payload(target_policy: str) -> dict[str, Any]:
152
216
  "prompts": deepcopy(VIDEO_VISUAL_PROMPTS),
153
217
  }
154
218
  return {
155
- "prompt_pack": "image_visual_v1",
219
+ "prompt_pack": "image_visual_v2",
156
220
  "media_type": "image",
157
221
  "target_policy": target_policy,
158
222
  "prompts": deepcopy(IMAGE_VISUAL_PROMPTS),
@@ -308,6 +308,8 @@ def _write_comparison_matrix(
308
308
  items: list[dict[str, Any]],
309
309
  *,
310
310
  title: str,
311
+ tile_size: tuple[int, int] = (208, 214),
312
+ lossless: bool = False,
311
313
  ) -> Path | None:
312
314
  if not items:
313
315
  return None
@@ -320,12 +322,11 @@ def _write_comparison_matrix(
320
322
  return None
321
323
 
322
324
  by_cell = {(item["row_key"], item["col_key"]): item for item in items}
323
- row_label_width = 210
324
- col_width = 220
325
+ row_label_width = 260 if tile_size[0] >= 512 else 210
326
+ col_width = tile_size[0] + 12
325
327
  title_height = 34
326
328
  header_height = 66
327
- row_height = 230
328
- tile_size = (col_width - 12, row_height - 16)
329
+ row_height = tile_size[1] + 16
329
330
  width = row_label_width + col_width * len(col_keys)
330
331
  height = title_height + header_height + row_height * len(row_keys)
331
332
  sheet = Image.new("RGB", (width, height), "white")
@@ -364,7 +365,7 @@ def _write_comparison_matrix(
364
365
  x = row_label_width + col_index * col_width + 6
365
366
  tile = _thumbnail(item["path"], size=tile_size)
366
367
  sheet.paste(tile, (x, y + 8))
367
- sheet.save(path)
368
+ sheet.save(path, lossless=lossless, method=6)
368
369
  return path
369
370
 
370
371
 
@@ -408,6 +409,8 @@ def _write_report_comparison_matrices(
408
409
  output_path,
409
410
  items,
410
411
  title=f"OrbitQuant {media} generation comparison matrix",
412
+ tile_size=(1024, 1024) if media == "image" else (208, 214),
413
+ lossless=media == "image",
411
414
  )
412
415
  if result is not None:
413
416
  created[media] = result