Add initial results

Author: MohsenDehghankar

README.md

@@ -14,4 +14,32 @@ RSR-core/
 │   └── cuda/       # CUDA GPU kernels
 ├── tests/          # Unit and integration tests
 └── benchmarks/     # Performance benchmarks
-```
+```
+
+## Benchmark Results
+
+### Matrix-Vector Multiplication
+
+#### CPU:
+
+#### CUDA:
+
+### Ternary (1.58bit) LLMs
+
+Speedup is computed from `Avg Time` against the `HF bfloat16` baseline for the same model.
+
+#### CPU 🖥️
+| Model | HF Time | RSR (ours) Time | HF Tok/s | RSR (ours) Tok/s | Speedup vs HF |
+| --- | ---: | ---: | ---: | ---: | ---: |
+| Falcon3-10B-Instruct-1.58bit | 351.215s | **5.663s** | 0.2 | **11.3** | **62.0x** |
+| Llama3-8B-1.58-100B-tokens | 261.557s | **4.862s** | 0.2 | **13.4** | **53.8x** |
+| bitnet-b1.58-2B-4T-bf16 | 31.446s | **2.258s** | 2.1 | **28.8** | **13.9x** |
+| bitnet-b1.58-2B-4T | 4.582s | **2.221s** | 14.2 | **29.3** | **2.1x** |
+
+#### CUDA ⚡
+| Model | HF Time | RSR (ours) Time | HF Tok/s | RSR (ours) Tok/s | Speedup vs HF |
+| --- | ---: | ---: | ---: | ---: | ---: |
+| Falcon3-10B-Instruct-1.58bit | 2.536s | **1.351s** | 25.2 | **47.4** | **1.9x** |
+| Llama3-8B-1.58-100B-tokens | 2.035s | **1.097s** | 31.9 | **59.3** | **1.9x** |
+| bitnet-b1.58-2B-4T-bf16 | 1.966s | **1.133s** | 33.1 | **57.4** | **1.7x** |
+| bitnet-b1.58-2B-4T | 1.563s | **1.139s** | 41.6 | **57.1** | **1.4x** |

integrations/bitnet/COMING_SOON.md

@@ -0,0 +1 @@
+Coming soon!

integrations/hf/model_infer.py

@@ -49,6 +49,11 @@
 )
 
 
+def _bitnet_act_quant(activation: torch.Tensor) -> torch.Tensor:
+    """Compatibility wrapper around the shared BitNet activation quantizer."""
+    return bitnet_act_quant(activation)
+
+
 def _read_json(path: Path) -> dict[str, Any]:
     return json.loads(path.read_text())
 
@@ -175,10 +180,13 @@ def forward(self, inputs: torch.Tensor) -> torch.Tensor:
         if self.rms_norm is not None:
             inputs = self.rms_norm(inputs)
 
+        input_device = inputs.device
+        input_dtype = inputs.dtype
+
         if self._cuda_backend:
-            # CUDA path: unchanged — the kernel does its own quantization.
-            inputs = inputs.float()
-            flat_inputs = inputs.reshape(-1, self.in_features)
+            # CUDA path: run the kernel in float32, then restore the caller's
+            # activation dtype so surrounding bf16 modules keep matching dtypes.
+            flat_inputs = inputs.float().reshape(-1, self.in_features)
             out_device = self.multiplier.device
             flat_outputs = torch.empty(
                 (flat_inputs.shape[0], self.out_features),
@@ -231,8 +239,8 @@ def forward(self, inputs: torch.Tensor) -> torch.Tensor:
             else:
                 output = output * ws
 
-        if inputs.device.type != "cpu" or inputs.dtype != output.dtype:
-            output = output.to(device=inputs.device, dtype=inputs.dtype)
+        if output.device != input_device or output.dtype != input_dtype:
+            output = output.to(device=input_device, dtype=input_dtype)
         return output
 
     def extra_repr(self) -> str:

tests/test_model_infer.py

@@ -189,6 +189,20 @@ def test_forward_batch(self):
         out = layer(x)
         assert out.shape == (4, 32)
 
+    def test_preserves_bfloat16_for_downstream_cuda_linear(self):
+        layer, _ = _make_cuda_rsr_linear(32, 64)
+        x = torch.randn(4, 64, device="cuda", dtype=torch.bfloat16)
+
+        out = layer(x)
+
+        assert out.device.type == "cuda"
+        assert out.dtype == torch.bfloat16
+
+        downstream = nn.Linear(32, 16, bias=False, device="cuda", dtype=torch.bfloat16)
+        y = downstream(out)
+        assert y.dtype == torch.bfloat16
+        assert y.shape == (4, 16)
+
     def test_weight_scale_multiply(self):
         layer, _ = _make_cuda_rsr_linear(32, 64, weight_scale_mode="multiply",
                                           weight_scale_val=2.0)