feat(qwen4_exp): load-time per-tensor FP8 dense projections (W8A8 via _scaled_mm)

python/freetoken/engine/config.py

@@ -74,7 +74,7 @@ class EngineConfig:
     # ratio default above. A runtime cache rebuild sets this (num_swa_pages) to pin the window
     # regardless of the full anchor; the ratio is the startup default and the fallback.
     swa_num_pages_override: int | None = None
-    distributed_timeout: float = 60.0
+    distributed_timeout: float = 1800.0  # ranks reach the first collective minutes apart on a 100+ GiB offload load
     use_dummy_weight: bool = False
     use_pynccl: bool = True
     max_seq_len_override: int | None = None

python/freetoken/layers/fp8_dynamic.py

@@ -0,0 +1,154 @@
+"""Load-time per-tensor FP8 dense linear (W8A8 through cuBLASLt ``torch._scaled_mm``).
+
+The weight is e4m3 ``[out_local, in_local]`` with one fp32 ``weight_scale`` (shape ``()``),
+produced by the model's weight reader from the bf16 checkpoint tensor (after TP sharding).
+The activation is quantized per call with a dynamic per-tensor scale: one fused Triton launch
+(amax pass, then cast) at decode sizes, a torch reduction plus a cast kernel above that. No
+host sync anywhere (the scales stay on the device), so the decode path is CUDA-graph safe;
+the branch between the two paths is on the tensor *shape*, never on its values.
+
+Measured on an RTX 6000 Ada (sm_89, torch 2.11.0+cu130) at the qwen4_exp TP=2 shapes, per
+decode step per rank over 48 layers: bf16 cuBLAS 3.2-3.4 ms, this path 1.9 ms at M=1/8/16.
+Requires sm_89+ (``_scaled_mm``'s floor; Ampere has no FP8 tensor cores).
+"""
+
+from __future__ import annotations
+
+from typing import List
+
+import torch
+import triton
+import triton.language as tl
+from freetoken.distributed import DistributedCommunicator, get_tp_info
+from freetoken.utils import div_even
+
+from .base import BaseOP
+
+FP8 = torch.float8_e4m3fn
+E4M3_MAX = 448.0
+_FUSED_MAX_ELEMENTS = (
+    65536  # one program handles the whole tensor below this (decode sizes)
+)
+_SCALE_FLOOR = 1e-12  # an all-zero activation (graph warmup buffers) must not give 1/0
+
+
+@triton.jit
+def _quant_fused_kernel(x_ptr, out_ptr, scale_ptr, n, BLOCK: tl.constexpr):
+    """One program: max|x| over the tensor, then the cast. Decode-sized inputs only."""
+    acc = tl.zeros([BLOCK], dtype=tl.float32)
+    for start in range(0, n, BLOCK):
+        offs = start + tl.arange(0, BLOCK)
+        v = tl.load(x_ptr + offs, mask=offs < n, other=0.0).to(tl.float32)
+        acc = tl.maximum(acc, tl.abs(v))
+    amax = tl.maximum(tl.max(acc, axis=0), 1e-12)
+    tl.store(scale_ptr, amax / 448.0)
+    inv = 448.0 / amax
+    for start in range(0, n, BLOCK):
+        offs = start + tl.arange(0, BLOCK)
+        mask = offs < n
+        v = tl.load(x_ptr + offs, mask=mask, other=0.0).to(tl.float32) * inv
+        v = tl.minimum(tl.maximum(v, -448.0), 448.0)
+        tl.store(out_ptr + offs, v.to(tl.float8e4nv), mask=mask)
+
+
+@triton.jit
+def _quant_cast_kernel(x_ptr, out_ptr, scale_ptr, n, BLOCK: tl.constexpr):
+    """Cast under a scale already on the device (prefill sizes; the amax is a torch reduction)."""
+    offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < n
+    inv = 1.0 / tl.load(scale_ptr)
+    v = tl.load(x_ptr + offs, mask=mask, other=0.0).to(tl.float32) * inv
+    v = tl.minimum(tl.maximum(v, -448.0), 448.0)
+    tl.store(out_ptr + offs, v.to(tl.float8e4nv), mask=mask)
+
+
+def quant_per_tensor(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """``(x_fp8, scale)`` with ``x ~= x_fp8 * scale``; ``x`` contiguous, ``scale`` fp32 ``()``."""
+    n = x.numel()
+    out = torch.empty_like(x, dtype=FP8)
+    if n <= _FUSED_MAX_ELEMENTS:
+        scale = torch.empty((), dtype=torch.float32, device=x.device)
+        _quant_fused_kernel[(1,)](x, out, scale, n, BLOCK=4096, num_warps=8)
+        return out, scale
+    amax = torch.linalg.vector_norm(x, ord=float("inf")).float()
+    scale = amax.clamp_min_(_SCALE_FLOOR).div_(E4M3_MAX)
+    _quant_cast_kernel[(triton.cdiv(n, 4096),)](
+        x, out, scale, n, BLOCK=4096, num_warps=4
+    )
+    return out, scale
+
+
+def fp8_dynamic_linear(
+    x: torch.Tensor, weight: torch.Tensor, weight_scale: torch.Tensor
+) -> torch.Tensor:
+    """``x @ (weight * weight_scale)^T`` in W8A8; ``weight`` [N, K] e4m3 row-major, whose ``.t()``
+    is the column-major operand cuBLASLt wants (a stride change, never a copy)."""
+    *lead, k = x.shape
+    x2 = x.reshape(-1, k).contiguous()
+    x8, scale = quant_per_tensor(x2)
+    y = torch._scaled_mm(
+        x8, weight.t(), scale_a=scale, scale_b=weight_scale, out_dtype=x.dtype
+    )
+    return y.reshape(*lead, weight.shape[0])
+
+
+class Fp8DynamicLinear(BaseOP):
+    """Per-tensor FP8 linear over the local shard; ``all_reduce`` adds the TP sum (row-parallel)."""
+
+    def __init__(self, local_isize: int, local_osize: int, *, all_reduce: bool = False):
+        assert local_isize % 16 == 0 and local_osize % 16 == 0, (
+            local_isize,
+            local_osize,
+        )
+        self.local_input_size = local_isize
+        self.local_output_size = local_osize
+        self.weight = torch.empty(local_osize, local_isize, dtype=FP8)
+        self.weight_scale = torch.empty((), dtype=torch.float32)
+        self._comm = (
+            DistributedCommunicator() if all_reduce and get_tp_info().size > 1 else None
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        y = fp8_dynamic_linear(x, self.weight, self.weight_scale)
+        if self._comm is not None:
+            y = self._comm.all_reduce(y)
+        return y
+
+
+class Fp8DynamicColMerged(Fp8DynamicLinear):
+    """Drop-in for ``LinearColParallelMerged``: one weight concatenating several projections
+    along the output dim; the caller splits the output by the local sizes as before."""
+
+    def __init__(
+        self,
+        input_size: int,
+        output_sizes: List[int],
+        local_output_sizes: List[int] | None = None,
+    ):
+        tp = get_tp_info()
+        if local_output_sizes is None:
+            local_output_sizes = [div_even(size, tp.size) for size in output_sizes]
+        self.output_sizes = list(output_sizes)
+        self.local_output_sizes = list(local_output_sizes)
+        super().__init__(input_size, sum(local_output_sizes))
+
+
+class Fp8DynamicRowParallel(Fp8DynamicLinear):
+    """Drop-in for ``LinearOProj`` / ``LinearRowParallel``: the input dim is sharded, the
+    all-reduce runs after the local GEMM (each rank scales its own shard)."""
+
+    def __init__(self, input_size: int, output_size: int):
+        super().__init__(
+            div_even(input_size, get_tp_info().size), output_size, all_reduce=True
+        )
+
+
+__all__ = [
+    "FP8",
+    "E4M3_MAX",
+    "Fp8DynamicColMerged",
+    "Fp8DynamicLinear",
+    "Fp8DynamicRowParallel",
+    "fp8_dynamic_linear",
+    "quant_per_tensor",
+]

python/freetoken/layers/linear.py

@@ -59,10 +59,14 @@ def __init__(
         input_size: int,
         output_sizes: List[int],
         has_bias: bool,
+        local_output_sizes: List[int] | None = None,
     ):
-        # check that all output sizes are divisible by tp_size
+        # check that all output sizes are divisible by tp_size (a caller that replicates
+        # GQA kv heads across ranks passes the per-rank sizes explicitly)
         tp_info = get_tp_info()
-        tp_output_sizes = [div_even(size, tp_info.size) for size in output_sizes]
+        if local_output_sizes is None:
+            local_output_sizes = [div_even(size, tp_info.size) for size in output_sizes]
+        tp_output_sizes = local_output_sizes
         output_size = sum(output_sizes)
         tp_output_size = sum(tp_output_sizes)
         super().__init__(input_size, output_size, input_size, tp_output_size, has_bias)

python/freetoken/models/config.py

@@ -18,6 +18,13 @@ def vision_load_enabled() -> bool:
     return os.getenv("FREETOKEN_LOAD_VISION", "0").strip().lower() in _VISION_TRUE
 
 
+def fp8_dense_enabled() -> bool:
+    """Load-time FP8 for a model's bf16 attention / GDN projections (opt-in, default OFF):
+    per-tensor e4m3 weights, per-tensor dynamic activation scale, cuBLASLt W8A8 GEMMs
+    (``torch._scaled_mm``, sm_89+). ``FREETOKEN_FP8_DENSE=1``."""
+    return os.getenv("FREETOKEN_FP8_DENSE", "0").strip().lower() in _VISION_TRUE
+
+
 def detect_expert_quant(hf_config: Any) -> str:
     """Routed-expert quantization from a checkpoint's ``quantization_config``: ``"nvfp4"`` for
     a ModelOpt FP4 build (``quant_algo: NVFP4``) OR an llm-compressor NVFP4 export

python/freetoken/models/nvfp4_banks.py

@@ -9,7 +9,8 @@
 
 import safetensors
 import torch
-from freetoken.utils import download_hf_weight
+from freetoken.distributed import get_tp_info
+from freetoken.utils import div_even, download_hf_weight
 from tqdm import tqdm
 
 LayerToBank = Callable[[int, object], int | None]
@@ -78,6 +79,41 @@ def _alloc_nvfp4_host_banks(num_layers: int, E: int, H: int, I: int):
     }, num_layers)
 
 
+def _tp_slice(inter: int) -> tuple[int, int]:
+    """``(i_local, i_lo)``: this rank's slice of the intermediate axis. TP shards every expert
+    along I (gate/up rows, down columns), the ``stream_moe_expert_sources`` convention, so the
+    routed output is a partial sum the MoE layer all-reduces."""
+    tp = get_tp_info()
+    i_local = div_even(inter, tp.size)
+    assert i_local % 16 == 0, f"NVFP4 TP shard {i_local} must cover whole 16-wide scale blocks"
+    return i_local, tp.rank * i_local
+
+
+class _Placer:
+    """Writes one checkpoint expert tensor into its bank slot (this rank's I slice only)."""
+
+    def __init__(self, banks: dict, inter: int):
+        self.b = banks
+        self.i_local, self.i_lo = _tp_slice(inter)
+
+    def put(self, layer: int, expert: int, role: str, kind: str, tensor, global_scale=None):
+        n, lo = self.i_local, self.i_lo
+        b = self.b
+        if role == "down":  # [H, I/2] codes, [H, I/16] scales, [H] global
+            if kind == "weight":
+                b["down_packed"][layer][expert] = tensor[:, lo // 2 : (lo + n) // 2]
+            else:
+                b["down_scale"][layer][expert] = tensor[:, lo // 16 : (lo + n) // 16]
+                b["down_global"][layer][expert] = global_scale
+            return
+        rows = slice(0, n) if role == "gate" else slice(n, 2 * n)  # gate | up on the row axis
+        if kind == "weight":
+            b["gate_up_packed"][layer][expert, rows] = tensor[lo : lo + n]
+        else:
+            b["gate_up_scale"][layer][expert, rows] = tensor[lo : lo + n]
+            b["gate_up_global"][layer][expert, rows] = global_scale  # per-tensor scalar
+
+
 def load_nvfp4_expert_source_banks(
     model_path: str,
     config,
@@ -149,13 +185,9 @@ def load_nvfp4_expert_source_banks(
                 globals_map[key] = _ingest_global(spec, f.get_tensor(name))
         drop_page_cache(path)
 
-    _hb = _alloc_nvfp4_host_banks(num_layers, E, H, I)  # unpinned; pinned after fill
-    gate_up_packed = [b.tensor for b in _hb["gate_up_packed"]]
-    gate_up_scale = [b.tensor for b in _hb["gate_up_scale"]]
-    gate_up_global = [b.tensor for b in _hb["gate_up_global"]]
-    down_packed = [b.tensor for b in _hb["down_packed"]]
-    down_scale = [b.tensor for b in _hb["down_scale"]]
-    down_global = [b.tensor for b in _hb["down_global"]]
+    _hb = _alloc_nvfp4_host_banks(num_layers, E, H, _tp_slice(I)[0])  # unpinned; pinned after fill
+    banks = {name: [b.tensor for b in layers] for name, layers in _hb.items()}
+    place = _Placer(banks, I)
 
     from freetoken.moe.host_banks import LayerCompletionTracker, PinPipeline
 
@@ -170,30 +202,11 @@ def _load(sink) -> int:
                     expert = int(match.group("expert"))
                     proj = match.group("proj")
                     role = spec.proj_to_role[proj]
+                    if role not in ("gate", "up", "down"):
+                        raise ValueError(f"{spec.desc}: unknown projection role {role!r}")
                     kind = _canon_kind(spec, match.group("kind"))
-                    tensor = f.get_tensor(name)
-                    if kind == "weight":
-                        if role == "gate":
-                            gate_up_packed[bank_layer_id][expert, :I] = tensor
-                        elif role == "up":
-                            gate_up_packed[bank_layer_id][expert, I:] = tensor
-                        elif role == "down":
-                            down_packed[bank_layer_id][expert] = tensor
-                        else:
-                            raise ValueError(f"{spec.desc}: unknown projection role {role!r}")
-                    else:
-                        global_scale = globals_map[(layer, expert, proj)]
-                        if role == "gate":
-                            gate_up_scale[bank_layer_id][expert, :I] = tensor
-                            gate_up_global[bank_layer_id][expert, :I] = global_scale
-                        elif role == "up":
-                            gate_up_scale[bank_layer_id][expert, I:] = tensor
-                            gate_up_global[bank_layer_id][expert, I:] = global_scale
-                        elif role == "down":
-                            down_scale[bank_layer_id][expert] = tensor
-                            down_global[bank_layer_id][expert] = global_scale
-                        else:
-                            raise ValueError(f"{spec.desc}: unknown projection role {role!r}")
+                    place.put(bank_layer_id, expert, role, kind, f.get_tensor(name),
+                              None if kind == "weight" else globals_map[(layer, expert, proj)])
                     tracker.note(bank_layer_id)
                     placed += 1
             drop_page_cache(path)
@@ -207,14 +220,7 @@ def _load(sink) -> int:
 
     expected = num_layers * E * 6
     assert placed == expected, f"{spec.desc}: loaded {placed} expert tensors, expected {expected}"
-    return {
-        "gate_up_packed": gate_up_packed,
-        "gate_up_scale": gate_up_scale,
-        "gate_up_global": gate_up_global,
-        "down_packed": down_packed,
-        "down_scale": down_scale,
-        "down_global": down_global,
-    }
+    return banks
 
 
 def load_nvfp4_expert_source_banks_parallel(
@@ -273,13 +279,9 @@ def load_nvfp4_expert_source_banks_parallel(
                 )
         drop_page_cache(path)
 
-    _hb = _alloc_nvfp4_host_banks(num_layers, E, H, I)  # unpinned; pinned after fill
-    gate_up_packed = [b.tensor for b in _hb["gate_up_packed"]]
-    gate_up_scale = [b.tensor for b in _hb["gate_up_scale"]]
-    gate_up_global = [b.tensor for b in _hb["gate_up_global"]]
-    down_packed = [b.tensor for b in _hb["down_packed"]]
-    down_scale = [b.tensor for b in _hb["down_scale"]]
-    down_global = [b.tensor for b in _hb["down_global"]]
+    _hb = _alloc_nvfp4_host_banks(num_layers, E, H, _tp_slice(I)[0])  # unpinned; pinned after fill
+    banks = {name: [b.tensor for b in layers] for name, layers in _hb.items()}
+    place = _Placer(banks, I)
 
     from freetoken.moe.host_banks import LayerCompletionTracker, PinPipeline
 
@@ -296,24 +298,8 @@ def _load(sink) -> int:
             proj = match.group("proj")
             role = spec.proj_to_role[proj]
             kind = _canon_kind(spec, match.group("kind"))
-            if kind == "weight":
-                if role == "gate":
-                    gate_up_packed[bank_layer_id][expert, :I] = tensor
-                elif role == "up":
-                    gate_up_packed[bank_layer_id][expert, I:] = tensor
-                else:
-                    down_packed[bank_layer_id][expert] = tensor
-            else:
-                g = globals_map[(layer, expert, proj)]
-                if role == "gate":
-                    gate_up_scale[bank_layer_id][expert, :I] = tensor
-                    gate_up_global[bank_layer_id][expert, :I] = g
-                elif role == "up":
-                    gate_up_scale[bank_layer_id][expert, I:] = tensor
-                    gate_up_global[bank_layer_id][expert, I:] = g
-                else:
-                    down_scale[bank_layer_id][expert] = tensor
-                    down_global[bank_layer_id][expert] = g
+            place.put(bank_layer_id, expert, role, kind, tensor,
+                      None if kind == "weight" else globals_map[(layer, expert, proj)])
             tracker.note(bank_layer_id)
             placed += 1
         return placed
@@ -326,14 +312,7 @@ def _load(sink) -> int:
 
     expected = num_layers * E * 6
     assert placed == expected, f"{spec.desc}: loaded {placed} expert tensors, expected {expected}"
-    return {
-        "gate_up_packed": gate_up_packed,
-        "gate_up_scale": gate_up_scale,
-        "gate_up_global": gate_up_global,
-        "down_packed": down_packed,
-        "down_scale": down_scale,
-        "down_global": down_global,
-    }
+    return banks
 
 
 __all__ = [