cpu_rewrite.patch

--- modeling_phi4flash_original.py	2026-04-03 06:10:53
+++ modeling_phi4flash_cpu.py	2026-04-03 06:10:41
@@ -41,7 +41,7 @@
     add_code_sample_docstrings,
     add_start_docstrings,
     add_start_docstrings_to_model_forward,
-    is_flash_attn_greater_or_equal_2_10,
+    is_flash_attn_greater_or_equal_2_10,  # may not exist without flash_attn
     logging,
     replace_return_docstrings,
 )
@@ -88,6 +88,115 @@
     selective_scan_cuda = None
     HAS_SELECTIVE_SCAN = False
 
+
+def _selective_scan_cpu(u, delta, A, B, C, D=None, z=None, delta_bias=None, delta_softplus=False, return_last_state=False):
+    """Pure PyTorch selective scan (forward only, for CPU inference).
+    Implements the S6 selective scan algorithm without CUDA kernels.
+    Uses float64 accumulation for numerical stability across 32+ layers.
+    """
+    orig_dtype = u.dtype
+    batch, dim, seqlen = u.shape
+
+    # Upcast to float64 for stable accumulation
+    u = u.double()
+    delta = delta.double()
+    A = A.double()
+
+    if B.dim() == 4:
+        B = B.squeeze(1)
+    if C.dim() == 4:
+        C = C.squeeze(1)
+    B = B.double()
+    C = C.double()
+
+    dstate = A.shape[1]
+
+    if delta_bias is not None:
+        delta = delta + delta_bias.double().unsqueeze(0).unsqueeze(-1)
+    if delta_softplus:
+        delta = F.softplus(delta)
+
+    # Discretize
+    deltaA = torch.exp(torch.einsum('bdl,dn->bdnl', delta, A))  # (B, D, N, L)
+    deltaB_u = torch.einsum('bdl,bnl->bdnl', delta, B) * u.unsqueeze(2)  # (B, D, N, L)
+
+    # Scan (recurrent)
+    x = torch.zeros(batch, dim, dstate, device=u.device, dtype=torch.float64)
+    ys = []
+    for i in range(seqlen):
+        x = deltaA[:, :, :, i] * x + deltaB_u[:, :, :, i]
+        y = torch.einsum('bdn,bn->bd', x, C[:, :, i])
+        ys.append(y)
+    y = torch.stack(ys, dim=-1)  # (B, D, L)
+
+    # Back to original dtype
+    y = y.to(orig_dtype)
+    x = x.to(orig_dtype)
+
+    if D is not None:
+        D_orig = D.to(orig_dtype)
+        u_orig = u.to(orig_dtype)
+        y = y + u_orig * D_orig.unsqueeze(0).unsqueeze(-1)
+
+    if z is not None:
+        y = y * F.silu(z.to(orig_dtype))
+
+    if return_last_state:
+        return y, x
+    return y
+
+
+def _sdpa_attention(q, k, v, causal=True, window_size=(-1, -1), softmax_scale=None, dropout_p=0.0):
+    """Scaled dot-product attention (replaces flash_attn_func for CPU).
+    Handles GQA via reshape instead of repeat_interleave for exact flash_attn match.
+    Args: q, k, v: (B, S, H, D) tensors (H may differ between q and k/v for GQA)
+    Returns: out: (B, S, H_q, D)
+    """
+    B, S_q, H_q, D = q.shape
+    _, S_k, H_kv, _ = k.shape
+
+    orig_dtype = q.dtype
+    q = q.float()
+    k = k.float()
+    v = v.float()
+
+    scale = softmax_scale if softmax_scale is not None else D ** -0.5
+
+    if H_q == H_kv:
+        # Standard attention — no GQA
+        # (B, S, H, D) -> (B, H, S, D)
+        q_t = q.transpose(1, 2)
+        k_t = k.transpose(1, 2)
+        v_t = v.transpose(1, 2)
+    else:
+        # GQA: reshape q into groups matching kv heads
+        # flash_attn broadcasts kv heads across groups internally
+        # We match this by reshaping q from (B, S, H_q, D) to (B*H_kv, S, n_groups, D)
+        # then kv from (B, S, H_kv, D) to (B*H_kv, S, 1, D) and broadcast
+        n_groups = H_q // H_kv
+        # Reshape: (B, S, H_kv, n_groups, D) -> (B, H_kv, n_groups, S, D) -> (B*H_kv, n_groups, S, D)
+        q_t = q.view(B, S_q, H_kv, n_groups, D).permute(0, 2, 3, 1, 4).reshape(B * H_kv, n_groups, S_q, D)
+        k_t = k.transpose(1, 2).reshape(B * H_kv, 1, S_k, D).expand(-1, n_groups, -1, -1).reshape(B * H_kv, n_groups, S_k, D)
+        v_t = v.transpose(1, 2).reshape(B * H_kv, 1, S_k, D).expand(-1, n_groups, -1, -1).reshape(B * H_kv, n_groups, S_k, D)
+
+    # Causal mask (bottom-right aligned for autoregressive generation)
+    if causal:
+        mask = torch.ones(S_q, S_k, dtype=torch.bool, device=q.device).tril(diagonal=S_k - S_q)
+        mask = mask.unsqueeze(0).unsqueeze(0)
+        out = F.scaled_dot_product_attention(q_t, k_t, v_t, attn_mask=mask, dropout_p=0.0, scale=scale)
+    else:
+        out = F.scaled_dot_product_attention(q_t, k_t, v_t, is_causal=False, dropout_p=0.0, scale=scale)
+
+    if H_q == H_kv:
+        out = out.transpose(1, 2)  # (B, S, H, D)
+    else:
+        # Reshape back: (B*H_kv, n_groups, S, D) -> (B, H_kv, n_groups, S, D) -> (B, S, H_kv*n_groups, D)
+        out = out.view(B, H_kv, n_groups, S_q, D).permute(0, 3, 1, 2, 4).reshape(B, S_q, H_q, D)
+
+    return out.to(orig_dtype)
+
+
+
 _CHECKPOINT_FOR_DOC = "microsoft/Phi-4-mini-flash-reasoning"
 _CONFIG_FOR_DOC = "Phi4FlashConfig"
 
@@ -379,35 +488,10 @@
 
 
 
-swiglu_fwd_codestring = """
-template <typename T> T swiglu_fwd(T x, T y) {
-    return float(x) * float(y) / (1.0f + ::exp(-float(x)));
-}
-"""
-swiglu_bwd_codestring = """
-template <typename T> T swiglu_bwd(T x, T y, T g, T& dx, T& dy) {
-    float x_sigmoid = 1.0f / (1.0f + ::exp(-float(x)));
-    dx = x_sigmoid * (1 + float(x) * (1.0f - x_sigmoid)) * float(g) * float(y);
-    dy = float(x) * x_sigmoid * float(g);
-}
-"""
-swiglu_fwd = torch.cuda.jiterator._create_jit_fn(swiglu_fwd_codestring)
-swiglu_bwd = torch.cuda.jiterator._create_multi_output_jit_fn(swiglu_bwd_codestring, num_outputs=2)
 
 
-class SwiGLUFunction(torch.autograd.Function):
-
-    @staticmethod
-    def forward(ctx, x, y):
-        ctx.save_for_backward(x, y)
-        return swiglu_fwd(x, y)
-
-    @staticmethod
-    def backward(ctx, dout):
-        x, y = ctx.saved_tensors
-        return swiglu_bwd(x, y, dout)
-
-swiglu = SwiGLUFunction.apply
+def swiglu(x, y):
+    return x * torch.sigmoid(x) * y
 
 
 # Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm with Llama->SambaY
@@ -670,27 +754,32 @@
         use_sliding_windows=False,
     ):
         """
-        Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token
-        first unpad the input, then computes the attention scores and pad the final attention scores.
-
-        Args:
-            query_states (`torch.Tensor`):
-                Input query states to be passed to Flash Attention API
-            key_states (`torch.Tensor`):
-                Input key states to be passed to Flash Attention API
-            value_states (`torch.Tensor`):
-                Input value states to be passed to Flash Attention API
-            attention_mask (`torch.Tensor`):
-                The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
-                position of padding tokens and 1 for the position of non-padding tokens.
-            dropout (`float`):
-                Attention dropout
-            softmax_scale (`float`, *optional*):
-                The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
-            use_sliding_windows (`bool`, *optional*):
-                Whether to activate sliding window attention.
+        Flash Attention forward with CPU fallback.
+        On CPU: skips unpadding, passes directly to inner_cross_attn (which uses SDPA).
+        On GPU: uses flash_attn with unpadding for efficiency.
         """
         causal = self.is_causal
+
+        # CPU path: skip unpadding, call inner_cross_attn directly
+        if not query_states.is_cuda:
+            window_size = (-1, -1)
+            if use_sliding_windows:
+                window_size = (
+                    self.config.sliding_window[self.layer_idx] - 1,
+                    self.config.sliding_window[self.layer_idx] - 1,
+                )
+            attn_output = self.inner_cross_attn(
+                query_states,
+                key_states,
+                value_states,
+                dropout_p=dropout,
+                softmax_scale=softmax_scale,
+                causal=causal,
+                window_size=window_size,
+            )
+            return attn_output
+
+        # GPU path: use flash_attn with unpadding
         # Contains at least one padding token in the sequence
         if attention_mask is not None:
             batch_size = query_states.shape[0]
@@ -827,7 +916,7 @@
         self.expand = expand
         self.d_inner = int(self.expand * self.d_model)
         self.dt_rank = math.ceil(self.d_model / 16) if dt_rank == "auto" else dt_rank
-        self.use_fast_path = use_fast_path
+        self.use_fast_path = False  # Disabled for CPU — fused kernels need CUDA
         self.layer_idx = layer_idx
 
         self.yoco_cross = yoco_cross
@@ -1166,7 +1255,7 @@
     supports_gradient_checkpointing = True
     _no_split_modules = ["SambaYDecoderLayer"]
     _skip_keys_device_placement = "past_key_values"
-    _supports_flash_attn_2 = True
+    _supports_flash_attn_2 = HAS_FLASH_ATTN
     _supports_sdpa = False
     _supports_cache_class = True
 
@@ -1819,7 +1908,10 @@
     last_state has shape (batch, dim, dstate). Note that the gradient of the last state is
     not considered in the backward pass.
     """
-    return SelectiveScanFn.apply(u, delta, A, B, C, D, z, delta_bias, delta_softplus, return_last_state)
+    if selective_scan_cuda is not None and u.is_cuda:
+        return SelectiveScanFn.apply(u, delta, A, B, C, D, z, delta_bias, delta_softplus, return_last_state)
+    else:
+        return _selective_scan_cpu(u, delta, A, B, C, D, z, delta_bias, delta_softplus, return_last_state)
 
 
 class MambaInnerFn(torch.autograd.Function):
@@ -1833,7 +1925,7 @@
         """
              xz: (batch, dim, seqlen)
         """
-        assert causal_conv1d_cuda is not None, "causal_conv1d_cuda is not available. Please install causal-conv1d."
+        assert causal_conv1d_cuda is not None, "causal_conv1d_cuda required for fused path"
         assert checkpoint_lvl in [0, 1]
         L = xz.shape[-1]
         delta_rank = delta_proj_weight.shape[1]
@@ -1908,7 +2000,7 @@
     @custom_bwd(device_type="cuda")
     def backward(ctx, dout):
         # dout: (batch, seqlen, dim)
-        assert causal_conv1d_cuda is not None, "causal_conv1d_cuda is not available. Please install causal-conv1d."
+        assert causal_conv1d_cuda is not None, "causal_conv1d_cuda required for fused path"
         (xz, conv1d_weight, conv1d_bias, x_dbl, x_proj_weight, delta_proj_weight, out_proj_weight,
          conv1d_out, delta, A, B, C, D, delta_bias, scan_intermediates, out) = ctx.saved_tensors
         L = xz.shape[-1]
@@ -2014,8 +2106,7 @@
         fa_og = True,
     ):
         super().__init__()
-        assert flash_attn_varlen_func is not None, "FlashAttention is not installed"
-        assert flash_attn_func is not None, "FlashAttention is not installed"
+        # flash_attn assertions removed for CPU compatibility
         self.head_dim = head_dim
         self.fa_og = fa_og # turning it to false needs customized flash attention https://github.com/xiayuqing0622/flex_head_fa
         self.lambda_init = lambda_init_fn(depth)
@@ -2056,12 +2147,15 @@
             out: (total, H, D) if cu_seqlens is not None and max_seqlen is not None,
                 else (B, S, H, D).
         """
-        q = q.to(torch.bfloat16)
-        k = k.to(torch.bfloat16)
-        v = v.to(torch.bfloat16)
-
-        assert q.dtype in [torch.float16, torch.bfloat16]
-        assert q.is_cuda and k.is_cuda and v.is_cuda
+        if q.is_cuda:
+            q = q.to(torch.bfloat16)
+            k = k.to(torch.bfloat16)
+            v = v.to(torch.bfloat16)
+        else:
+            q = q.float()
+            k = k.float()
+            v = v.float()
+        # assert q.is_cuda removed for CPU compatibility
         #causal = self.causal if causal is None else causal
         unpadded = cu_seqlens_q is not None
         q1, q2 = split_heads(q)
@@ -2078,36 +2172,45 @@
             "window_size": window_size,
         }
 
-        if unpadded:
-            assert cu_seqlens_q.dtype == torch.int32
-            assert max_seqlen_q is not None
-            assert isinstance(max_seqlen_q, int)
-            assert cu_seqlens_k is not None
-            assert cu_seqlens_k.dtype == torch.int32
-            assert max_seqlen_k is not None
-            assert isinstance(max_seqlen_k, int)
+        if HAS_FLASH_ATTN and q.is_cuda:
+            # Use flash attention on GPU
+            if unpadded:
+                kwargs.update({
+                    "cu_seqlens_q": cu_seqlens_q, "max_seqlen_q": max_seqlen_q,
+                    "cu_seqlens_k": cu_seqlens_k, "max_seqlen_k": max_seqlen_k,
+                })
+                attn_func = flash_attn_varlen_func
+            else:
+                attn_func = flash_attn_func
 
-            kwargs.update({
-                "cu_seqlens_q": cu_seqlens_q,
-                "max_seqlen_q": max_seqlen_q,
-                "cu_seqlens_k": cu_seqlens_k,
-                "max_seqlen_k": max_seqlen_k,
-            })
-            attn_func = flash_attn_varlen_func
+            if self.fa_og:
+                attn11 = attn_func(q1, k1, v1, **kwargs)
+                attn12 = attn_func(q1, k1, v2, **kwargs)
+                attn1 = torch.cat([attn11, attn12], dim=-1)
+                attn21 = attn_func(q2, k2, v1, **kwargs)
+                attn22 = attn_func(q2, k2, v2, **kwargs)
+                attn2 = torch.cat([attn21, attn22], dim=-1)
+            else:
+                attn1 = attn_func(q1, k1, v, **kwargs)
+                attn2 = attn_func(q2, k2, v, **kwargs)
         else:
-            attn_func = flash_attn_func
+            # CPU fallback using PyTorch scaled_dot_product_attention
+            sdpa_kwargs = {
+                "causal": kwargs.get("causal", False),
+                "softmax_scale": kwargs.get("softmax_scale", None),
+                "dropout_p": kwargs.get("dropout_p", 0.0),
+            }
+            if self.fa_og:
+                attn11 = _sdpa_attention(q1, k1, v1, **sdpa_kwargs)
+                attn12 = _sdpa_attention(q1, k1, v2, **sdpa_kwargs)
+                attn1 = torch.cat([attn11, attn12], dim=-1)
+                attn21 = _sdpa_attention(q2, k2, v1, **sdpa_kwargs)
+                attn22 = _sdpa_attention(q2, k2, v2, **sdpa_kwargs)
+                attn2 = torch.cat([attn21, attn22], dim=-1)
+            else:
+                attn1 = _sdpa_attention(q1, k1, v, **sdpa_kwargs)
+                attn2 = _sdpa_attention(q2, k2, v, **sdpa_kwargs)
 
-        if self.fa_og:
-            attn11 = attn_func(q1, k1, v1, **kwargs)
-            attn12 = attn_func(q1, k1, v2, **kwargs)
-            attn1 = torch.cat([attn11, attn12], dim=-1)
-            attn21 = attn_func(q2, k2, v1, **kwargs)
-            attn22 = attn_func(q2, k2, v2, **kwargs)
-            attn2 = torch.cat([attn21, attn22], dim=-1)
-        else:
-            attn1 = attn_func(q1, k1, v, **kwargs)
-            attn2 = attn_func(q2, k2, v, **kwargs)
-
         lambda_1 = torch.exp(torch.sum(self.lambda_q1 * self.lambda_k1, dim=-1).float()).type_as(q)
         lambda_2 = torch.exp(torch.sum(self.lambda_q2 * self.lambda_k2, dim=-1).float()).type_as(q)
         lambda_full = lambda_1 - lambda_2 + self.lambda_init