Llama phase split quantisation support

src/chop/nn/quantized/modules/llama/mlp.py

@@ -1,8 +1,19 @@
+"""Llama MLP quantization with phase-aware dispatch.
+
+Default policy keeps decode full precision (`fp_only`) for compatibility.
+`quantized` decode is opt-in and consumes decode-phase sub-config directly.
+"""
+
 import torch
 from torch import nn, Tensor
 
 from chop.nn.quantizers.SNN.LSQ import LSQInteger
 from chop.nn.quantized.functional.silu import silu_minifloat
+from chop.nn.quantized.modules.phase_context import get_runtime_phase
+from chop.nn.quantized.modules.phase_config import (
+    get_phase_subconfig,
+    normalize_phase_q_config,
+)
 
 from transformers.models.llama.modeling_llama import LlamaMLP, ACT2FN
 
@@ -51,13 +62,21 @@ class LlamaMLPMXFP(LlamaMLP):
     def __init__(self, config, layer_idx=None, q_config: dict = None):
         super().__init__(config)
         self.layer_idx = layer_idx
-        self.q_config = q_config or {}
-        self.bypass = self.q_config.get("bypass", False)
+        self.phase_q_config = normalize_phase_q_config(q_config)
+        self.decode_policy = self.phase_q_config["decode_policy"]
 
     def forward(self, x: Tensor) -> Tensor:
-        if self.bypass:
+        runtime_phase = get_runtime_phase()
+        phase_subconfig, decode_policy = get_phase_subconfig(
+            self.phase_q_config, runtime_phase
+        )
+        bypass = phase_subconfig.get("bypass", False)
+        if runtime_phase == "decode" and decode_policy == "fp_only":
+            bypass = True
+
+        if bypass:
             return super().forward(x)
-        x = silu_minifloat(self.gate_proj(x), self.q_config) * self.up_proj(x)
+        x = silu_minifloat(self.gate_proj(x), phase_subconfig) * self.up_proj(x)
         return self.down_proj(x)
 
 
@@ -67,11 +86,19 @@ class LlamaMLPMXInt(LlamaMLP):
     def __init__(self, config, layer_idx=None, q_config: dict = None):
         super().__init__(config)
         self.layer_idx = layer_idx
-        self.q_config = q_config or {}
-        self.bypass = self.q_config.get("bypass", False)
+        self.phase_q_config = normalize_phase_q_config(q_config)
+        self.decode_policy = self.phase_q_config["decode_policy"]
 
     def forward(self, x: Tensor) -> Tensor:
-        if self.bypass:
+        runtime_phase = get_runtime_phase()
+        phase_subconfig, decode_policy = get_phase_subconfig(
+            self.phase_q_config, runtime_phase
+        )
+        bypass = phase_subconfig.get("bypass", False)
+        if runtime_phase == "decode" and decode_policy == "fp_only":
+            bypass = True
+
+        if bypass:
             return super().forward(x)
-        x = silu_minifloat(self.gate_proj(x), self.q_config) * self.up_proj(x)
+        x = silu_minifloat(self.gate_proj(x), phase_subconfig) * self.up_proj(x)
         return self.down_proj(x)

src/chop/nn/quantized/modules/llama/rms_norm.py

@@ -1,10 +1,21 @@
+"""Llama RMSNorm quantization with phase-aware dispatch.
+
+Default decode behavior remains FP-only for backward compatibility.
+`quantized` decode is opt-in via phase config and policy.
+"""
+
 from functools import partial
 
 import torch
 from torch import Tensor, nn
 
 from chop.nn.quantizers.SNN.LSQ import LSQInteger
 from chop.nn.quantizers._minifloat_mx import MinifloatMeta, minifloat_quantizer_sim
+from chop.nn.quantized.modules.phase_context import get_runtime_phase
+from chop.nn.quantized.modules.phase_config import (
+    get_phase_subconfig,
+    normalize_phase_q_config,
+)
 
 from transformers.models.llama.modeling_llama import LlamaRMSNorm
 
@@ -37,48 +48,59 @@ class LlamaRMSNormMinifloat(LlamaRMSNorm):
     def __init__(self, config=None, layer_idx=None, q_config: dict = None):
         super().__init__(hidden_size=config.hidden_size, eps=config.rms_norm_eps)
         self.layer_idx = layer_idx
-        self.q_config = q_config or {}
         self.variance_epsilon = config.rms_norm_eps
-        self.bypass = self.q_config.get("bypass", False)
-        self.weight_bypass = self.q_config.get("weight_bypass", False)
-        self.data_in_bypass = self.q_config.get("data_in_bypass", False)
-
-        if not self.bypass and not self.weight_bypass:
-            self.w_quantizer = partial(
-                minifloat_quantizer_sim,
-                minifloat_meta=MinifloatMeta(
-                    exp_bits=self.q_config["weight_exponent_width"],
-                    frac_bits=self.q_config["weight_frac_width"],
-                    is_finite=self.q_config.get("weight_is_finite", True),
-                    round_mode=self.q_config.get("weight_round_mode", "rn"),
-                ),
-            )
-        else:
-            self.w_quantizer = None
-
-        if not self.bypass and not self.data_in_bypass:
-            self.x_quantizer = partial(
-                minifloat_quantizer_sim,
-                minifloat_meta=MinifloatMeta(
-                    exp_bits=self.q_config["data_in_exponent_width"],
-                    frac_bits=self.q_config["data_in_frac_width"],
-                    is_finite=self.q_config.get("data_in_is_finite", True),
-                    round_mode=self.q_config.get("data_in_round_mode", "rn"),
-                ),
-            )
-        else:
-            self.x_quantizer = None
+        self.phase_q_config = normalize_phase_q_config(q_config)
+        self.decode_policy = self.phase_q_config["decode_policy"]
+
+    @staticmethod
+    def _build_weight_quantizer(sub_cfg: dict, bypass: bool):
+        """Build weight quantizer for current phase config."""
+
+        if bypass or sub_cfg.get("weight_bypass", False):
+            return None
+        return partial(
+            minifloat_quantizer_sim,
+            minifloat_meta=MinifloatMeta(
+                exp_bits=sub_cfg["weight_exponent_width"],
+                frac_bits=sub_cfg["weight_frac_width"],
+                is_finite=sub_cfg.get("weight_is_finite", True),
+                round_mode=sub_cfg.get("weight_round_mode", "rn"),
+            ),
+        )
+
+    @staticmethod
+    def _build_input_quantizer(sub_cfg: dict, bypass: bool):
+        """Build input quantizer for current phase config."""
+
+        if bypass or sub_cfg.get("data_in_bypass", False):
+            return None
+        return partial(
+            minifloat_quantizer_sim,
+            minifloat_meta=MinifloatMeta(
+                exp_bits=sub_cfg["data_in_exponent_width"],
+                frac_bits=sub_cfg["data_in_frac_width"],
+                is_finite=sub_cfg.get("data_in_is_finite", True),
+                round_mode=sub_cfg.get("data_in_round_mode", "rn"),
+            ),
+        )
 
     def forward(self, hidden_states):
+        runtime_phase = get_runtime_phase()
+        phase_subconfig, decode_policy = get_phase_subconfig(
+            self.phase_q_config, runtime_phase
+        )
+        bypass = phase_subconfig.get("bypass", False)
+        if runtime_phase == "decode" and decode_policy == "fp_only":
+            bypass = True
+
+        w_quantizer = self._build_weight_quantizer(phase_subconfig, bypass)
+        x_quantizer = self._build_input_quantizer(phase_subconfig, bypass)
+
         input_dtype = hidden_states.dtype
-        if self.x_quantizer is not None:
-            hidden_states = self.x_quantizer(hidden_states)
+        if x_quantizer is not None:
+            hidden_states = x_quantizer(hidden_states)
         hidden_states = hidden_states.to(torch.float32)
         variance = hidden_states.pow(2).mean(-1, keepdim=True)
         hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
-        weight = (
-            self.w_quantizer(self.weight)
-            if self.w_quantizer is not None
-            else self.weight
-        )
+        weight = w_quantizer(self.weight) if w_quantizer is not None else self.weight
         return weight * hidden_states.to(input_dtype)

src/chop/nn/quantized/modules/phase_config.py

@@ -0,0 +1,57 @@
+"""Shared phase-config normalization helpers for quantized modules.
+
+Step-1 integration contracts:
+1. Accept both legacy flat configs and phase-structured configs.
+2. Return a stable normalized shape consumed by runtime modules.
+3. Do not decide runtime policy here (that remains module-side logic).
+"""
+
+from __future__ import annotations
+
+from copy import deepcopy
+from typing import Any
+
+
+def normalize_phase_q_config(q_config: dict[str, Any] | None) -> dict[str, Any]:
+    """Normalize config into `{decode_policy, prefill, decode}` shape.
+
+    Compatibility behavior:
+    - Legacy flat config -> copied into both prefill and decode buckets.
+    - Explicit phase config -> missing decode bucket falls back to prefill.
+
+    Why we keep this helper minimal:
+    - It preserves old config semantics without enforcing model-specific
+      policy in the parsing layer.
+    """
+
+    cfg = deepcopy(q_config or {})
+    decode_policy = cfg.get("decode_policy", "fp_only")
+
+    if "prefill" in cfg or "decode" in cfg:
+        prefill = deepcopy(cfg.get("prefill", {}))
+        decode = deepcopy(cfg.get("decode", prefill))
+    else:
+        legacy = deepcopy(cfg)
+        legacy.pop("decode_policy", None)
+        prefill = legacy
+        decode = deepcopy(legacy)
+
+    return {
+        "decode_policy": decode_policy,
+        "prefill": prefill,
+        "decode": decode,
+    }
+
+
+def get_phase_subconfig(
+    normalized: dict[str, Any], phase: str
+) -> tuple[dict[str, Any], str]:
+    """Select phase sub-config from normalized phase config.
+
+    Returns:
+        (sub_config, decode_policy)
+    """
+
+    if phase == "decode":
+        return deepcopy(normalized.get("decode", {})), normalized["decode_policy"]
+    return deepcopy(normalized.get("prefill", {})), normalized["decode_policy"]

src/chop/nn/quantized/modules/phase_context.py

@@ -0,0 +1,141 @@
+"""Shared runtime phase context for quantized decoder-only inference.
+
+This module intentionally centralizes phase state (`prefill` vs `decode`) in a
+single place so quantized modules can stay loosely coupled:
+
+1. Llama decoder-layer pre-hooks detect runtime phase from cache semantics.
+2. Downstream quantized modules (attention/linear/mlp/rms) read the same phase
+   without changing
+   their public `forward(...)` signatures.
+
+Why `ContextVar` is used:
+- It avoids global mutable state bleeding across threads/tasks.
+- It keeps the integration minimally invasive for the existing MASE module API.
+"""
+
+from __future__ import annotations
+
+from contextvars import ContextVar
+from typing import Any, Literal
+
+Phase = Literal["prefill", "decode"]
+DecodePolicy = Literal["fp_only", "quantized"]
+
+
+_ACTIVE_PHASE: ContextVar[Phase] = ContextVar("active_quant_phase", default="prefill")
+_DECODE_POLICY: ContextVar[DecodePolicy] = ContextVar(
+    "active_decode_policy", default="fp_only"
+)
+
+
+def set_runtime_phase(phase: Phase) -> None:
+    """Set runtime phase used by phase-aware quantized modules."""
+
+    if phase not in ("prefill", "decode"):
+        raise ValueError(f"Unsupported phase: {phase}")
+    _ACTIVE_PHASE.set(phase)
+
+
+def get_runtime_phase() -> Phase:
+    """Return current runtime phase.
+
+    Defaults to `prefill` when no explicit phase has been written yet.
+    """
+
+    return _ACTIVE_PHASE.get()
+
+
+def set_runtime_decode_policy(policy: DecodePolicy) -> None:
+    """Set decode policy in runtime phase context.
+
+    Supported policies:
+    - `fp_only`: decode path is forced to full precision.
+    - `quantized`: decode path may consume decode-phase quantized configs.
+
+    Why we validate here:
+    - This context is shared by all quantized modules during forward.
+    - Central validation prevents silent fallback to unintended behavior.
+    """
+
+    if policy not in ("fp_only", "quantized"):
+        raise ValueError(f"Unsupported decode policy: {policy}")
+    _DECODE_POLICY.set(policy)
+
+
+def get_runtime_decode_policy() -> DecodePolicy:
+    """Return decode policy from runtime phase context.
+
+    Defaults to `fp_only` to preserve backward compatibility.
+    """
+
+    return _DECODE_POLICY.get()
+
+
+def infer_runtime_phase_from_hidden_and_cache(
+    hidden_states: Any, past_cache: Any
+) -> Phase:
+    """Infer runtime phase from hidden-state shape and cache state.
+
+    Rules are intentionally unchanged from step-1 attention-local logic:
+    1) `past_cache is None` -> `prefill`
+    2) `past_cache.get_seq_length() > 0` -> `decode`
+    3) Fallback: if current query length is 1 -> `decode`, else `prefill`
+
+    Why this helper exists:
+    - The same heuristic must be shared by all Llama decoder layers.
+    - We keep it in the shared phase module so hooks and tests stay consistent.
+    """
+
+    if past_cache is None:
+        return "prefill"
+
+    past_len = 0
+    get_seq_length = getattr(past_cache, "get_seq_length", None)
+    if callable(get_seq_length):
+        try:
+            past_len = int(get_seq_length())
+        except Exception:  # pragma: no cover - defensive for custom Cache impls
+            past_len = 0
+
+    if past_len > 0:
+        return "decode"
+
+    q_len = 0
+    shape = getattr(hidden_states, "shape", None)
+    if shape is not None and len(shape) >= 2:
+        q_len = int(shape[-2])
+    if q_len == 1:
+        return "decode"
+    return "prefill"
+
+
+def extract_decoder_layer_past_cache(
+    args: tuple[Any, ...],
+    kwargs: dict[str, Any],
+) -> Any:
+    """Extract past-cache object from decoder-layer call arguments.
+
+    Compatibility:
+    - New HF naming: `past_key_values`
+    - Legacy alias: `past_key_value`
+    - Positional call fallback: 4th argument for `LlamaDecoderLayer.forward`
+    """
+
+    if "past_key_values" in kwargs:
+        return kwargs["past_key_values"]
+    if "past_key_value" in kwargs:
+        return kwargs["past_key_value"]
+    if len(args) >= 4:
+        return args[3]
+    return None
+
+
+def infer_runtime_phase_from_decoder_layer_inputs(
+    args: tuple[Any, ...],
+    kwargs: dict[str, Any],
+) -> Phase:
+    """Infer runtime phase from decoder-layer forward inputs."""
+
+    hidden_states = kwargs.get("hidden_states", args[0] if args else None)
+    past_cache = extract_decoder_layer_past_cache(args, kwargs)
+    return infer_runtime_phase_from_hidden_and_cache(hidden_states, past_cache)