Adds support for int8 w8a8_gemlite quantization

README.md

@@ -138,6 +138,21 @@ img = pipeline.append_frame(uint8_img)  # returns passed image
 
 Note: returned `img` is always on the same device as `engine.device`
 
+## Quantization 
+
+Model can be quantized by passing quant argument to WorldEngine
+```
+engine = WorldEngine("Overworld/Waypoint-1.5-1B", quant="intw8a8", device="cuda")
+```
+Supported inference quantization schemes are:
+
+| Config | Description | Supported GPUs |
+|--------|-------------|----------------|
+| `intw8a8` | INT8 weights + INT8 dynamic per-token activations | NVIDIA (30xx, 40xx, Ampere+) |
+| `fp8w8a8` | FP8 (e4m3) weights + FP8 per-tensor activations via `torch._scaled_mm` | NVIDIA Ada Lovelace / Hopper+ (RTX 40xx, H100) |
+| `nvfp4` | NVFP4 weights + FP4 activations via FlashInfer/CUTLASS | NVIDIA Blackwell (B100, B200, RTX 5090) |
+
+
 ### WorldEngine
 
 `WorldEngine` computes each new frame from past frames, the controls, and the current prompt, then appends it to the sequence so later frames stay aligned with what has already been generated.

examples/benchmark.py

@@ -56,6 +56,7 @@ def get_warm_engine(model_uri, model_overrides=None):
     engine = WorldEngine(
         model_uri,
         model_config_overrides=model_config_overrides,
+        quant="intw8a8",
         device="cuda",
         load_weights=False
     )
@@ -124,4 +125,4 @@ def target():
         if blocking:
             torch.cuda.synchronize()
 
-    benchmark.pedantic(target, setup=setup, rounds=20)
+    benchmark.pedantic(target, setup=setup, rounds=20)

examples/gen_sample.py

@@ -13,7 +13,7 @@
 
 
 # Create inference engine
-engine = WorldEngine(sys.argv[1], device="cuda")
+engine = WorldEngine(sys.argv[1], quant=None, device="cuda")
 
 
 # Define sequence of controller inputs applied
@@ -39,6 +39,7 @@
 url = random.choice(urls)
 
 seed_frame = cv2.imdecode(np.frombuffer(urllib.request.urlopen(url).read(), np.uint8), cv2.IMREAD_COLOR)
+seed_frame = cv2.resize(seed_frame, (1280, 720))
 seed_frame_x4 = torch.from_numpy(np.repeat(seed_frame[None], 4, axis=0))
 
 

pyproject.toml

@@ -17,10 +17,10 @@ dependencies = [
   "huggingface-hub>=1.7.0",
   "omegaconf",
   "accelerate==1.12.0",
-
   # Triton (platform-specific)
   "triton==3.6.0; sys_platform == 'linux'",
   "triton-windows==3.6.0.post26; sys_platform == 'win32'",
+  "gemlite @ git+https://github.com/dropbox/gemlite.git"
 ]
 
 [tool.setuptools]

src/model/base_model.py

@@ -6,7 +6,6 @@
 from torch import nn
 import torch
 
-
 MODEL_CONFIG_DEFAULTS = OmegaConf.create(
     {
         "auto_aspect_ratio": True,

src/quantize.py

@@ -2,6 +2,7 @@
 
 import torch
 import torch.nn as nn
+import os
 
 
 QUANTS = [None]  # TODO: enable specific quant based on model config, which should specify compatible quants [None, "w8a8", "fp8"]
@@ -12,6 +13,19 @@
     QUANTS.append("nvfp4")
 except ImportError:
     pass
+try:
+    config_path = "gemlite_config.json"
+    from gemlite.helper import A8W8_INT8_dynamic
+    import gemlite
+
+    if os.path.exists(config_path):
+        gemlite.load_config(config_path)
+    else: 
+        gemlite.set_autotune("max")
+        gemlite.helper.warmup(shapes=[(128,2048), (2048,2048), (2048,8192), (4096,2048), (8192,2048)], batch_sizes=[1], processor=A8W8_INT8_dynamic())
+        gemlite.cache_config(config_path)
+except ImportError:
+    A8W8_INT8_dynamic = None
 
 
 @torch.library.custom_op("world_engine::fp4_linear", mutates_args=())
@@ -183,11 +197,83 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
         return result.reshape(x.shape[:-1] + (-1,))
 
+def _per_token_quant_int8(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Local per-token symmetric int8 quantization matching SGLang's W8A8 flow:
+      scale = absmax / 127
+      x_q   = round(x / scale)
+    Returns:
+      x_q:    [..., K] int8
+      scales: [..., 1] float32
+    """
+    x_fp = x.float().nan_to_num()
+    scales = (x_fp.abs().amax(dim=-1, keepdim=True).clamp_min(1e-10) / 127.0).float()
+    x_q = torch.round(x_fp / scales).clamp(-127, 127).to(torch.int8)
+    return x_q, scales
+
+
+@torch.library.custom_op("world_engine::w8a8_int8_linear", mutates_args=())
+def w8a8_int8_linear(
+    a: torch.Tensor,
+    b_int8_T: torch.Tensor,
+    b_scale: torch.Tensor,
+    bias: torch.Tensor,
+) -> torch.Tensor:
+    if sgl_int8_scaled_mm is None:
+        raise ImportError("sgl-kernel is required for quant='w8a8'")
+
+    assert a.ndim == 2, "expected [M, K] input"
+    x_q, x_scale = _per_token_quant_int8(a.contiguous())
+
+    bias_arg = None if bias.numel() == 0 else bias
+    return sgl_int8_scaled_mm(
+        x_q,                 # [M, K] row-major int8
+        b_int8_T,            # [K, N] column-major int8 view
+        x_scale,             # [M, 1] float32
+        b_scale,             # [N, 1] float32
+        out_dtype=a.dtype,
+        bias=bias_arg,
+    )
+
+
+@w8a8_int8_linear.register_fake
+def _w8a8_int8_linear_fake(
+    a: torch.Tensor,
+    b_int8_T: torch.Tensor,
+    b_scale: torch.Tensor,
+    bias: torch.Tensor,
+) -> torch.Tensor:
+    return torch.empty(
+        (a.shape[0], b_int8_T.shape[1]),
+        device=a.device,
+        dtype=a.dtype,
+    )
+
+class INT8W8A8GemLite(nn.Module):
+    __constants__ = ("in_features", "out_features")
+
+    def __init__(self, lin: nn.Linear):
+        super().__init__()
+        if A8W8_INT8_dynamic is None:
+            raise ImportError("Install gemlite for quant='w8a8_gemlite'")
+
+        self.in_features = lin.in_features
+        self.out_features = lin.out_features
+
+        # Minimal wrapper: assumes the layer is already on the target CUDA device.
+        self.impl = A8W8_INT8_dynamic(
+            device=str(lin.weight.device),
+            dtype=lin.weight.dtype,
+        ).from_linear(lin)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.impl(x).type_as(x)
+
 
 def quantize_model(model: nn.Module, quant: str):
     if quant is None:
         return model
-
+    
     def eligible(m: nn.Module) -> bool:
         w = getattr(m, "weight", None)
         if not isinstance(m, nn.Linear):
@@ -198,7 +284,8 @@ def eligible(m: nn.Module) -> bool:
         return (o % 32 == 0) and (k % 32 == 0)
 
     new_linear = {
-        "w8a8": FP8W8A8Linear,
+        "intw8a8": INT8W8A8GemLite,
+        "fp8w8a8": FP8W8A8Linear,
         "nvfp4": FP4Linear,
         "fp8": FP8Linear,
     }[quant]
@@ -207,4 +294,4 @@ def eligible(m: nn.Module) -> bool:
         setattr(model, name, new_linear(child)) if eligible(child) else quantize_model(
             child, quant
         )
-    return model
+    return model

src/world_engine.py

@@ -45,8 +45,7 @@ def __init__(
     ):
         """
         model_uri: HF URI or local folder containing model.safetensors and config.yaml
-        quant: None | w8a8 | nvfp4
-
+        quant: None | intw8a8 | fp8w8a8 | nvfp4
         model_config_overrides: Dict to override model config values
         - auto_aspect_ratio: set to False to work in ae raw space, otherwise in/out are 720p or 360p
         """