build(SM70): precompile GDN CUDA kernel to .so during docker build

precompile_gdn.py: calls torch.utils.cpp_extension.load with build_directory
to produce .so at build time. If build env has no GPU/compiler, fails
gracefully — kernel JIT compiles at runtime instead.

fused_fwd.py: _load_ext() now checks build/ dir for precompiled .so first,
skips 2-minute JIT compilation if found.

qwen3_6_scripts/flash_qla_sm70/__init__.py

@@ -0,0 +1,14 @@
+# Copyright (c) 2026 The Qwen team, Alibaba Group.
+# Licensed under The MIT License [see LICENSE for details]
+
+from .fused_fwd import (
+    chunk_gated_delta_rule_fwd_sm70,
+    chunk_gated_delta_rule_fwd_sm70_vlk_varlen,
+    resolve_column_groups_per_block_sm70,
+)
+
+__all__ = [
+    "chunk_gated_delta_rule_fwd_sm70",
+    "chunk_gated_delta_rule_fwd_sm70_vlk_varlen",
+    "resolve_column_groups_per_block_sm70",
+]

qwen3_6_scripts/flash_qla_sm70/fused_fwd.py

@@ -0,0 +1,508 @@
+# Copyright (c) 2026 The Qwen team, Alibaba Group.
+# Licensed under The MIT License [see LICENSE for details]
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import torch
+from torch.utils.cpp_extension import load
+
+_EXT = None
+
+
+def _load_ext():
+    global _EXT
+    if _EXT is not None:
+        return _EXT
+    if not torch.cuda.is_available():
+        raise RuntimeError("SM70 FlashQLA backend requires CUDA.")
+
+    os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "7.0;7.5")
+
+    # Try precompiled .so first (built during docker build)
+    build_dir = Path(__file__).with_name("build")
+    if build_dir.is_dir():
+        so_files = list(build_dir.glob("*.so"))
+        if so_files:
+            try:
+                _EXT = load(
+                    name="flash_qla_sm70_gdn_strided",
+                    sources=[],  # empty — just load from build_directory
+                    build_directory=str(build_dir),
+                    verbose=False,
+                )
+                return _EXT
+            except Exception:
+                pass  # fall through to JIT
+
+    # JIT compile (slow, ~2min first time)
+    src = Path(__file__).with_name("csrc") / "gdn_forward.cu"
+    _EXT = load(
+        name="flash_qla_sm70_gdn_strided",
+        sources=[str(src)],
+        extra_cuda_cflags=["-O3"],
+        extra_cflags=["-O3"],
+        verbose=bool(int(os.environ.get("FLASH_QLA_SM70_VERBOSE_BUILD", "0"))),
+    )
+    return _EXT
+
+
+def _check_inputs(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    g: torch.Tensor,
+    beta: torch.Tensor,
+    initial_state: torch.Tensor | None,
+) -> None:
+    tensors = [q, k, v, g, beta]
+    if initial_state is not None:
+        tensors.append(initial_state)
+    if any(not tensor.is_cuda for tensor in tensors):
+        raise ValueError("SM70 GDN tensors must be CUDA tensors.")
+    if any(tensor.device != q.device for tensor in tensors):
+        raise ValueError("SM70 GDN tensors must be on the same CUDA device.")
+    if any(not tensor.is_contiguous() for tensor in tensors):
+        raise ValueError("SM70 GDN tensors must be contiguous.")
+    if q.dtype not in (torch.float16, torch.bfloat16, torch.float32):
+        raise ValueError("SM70 GDN backend supports fp16, bf16, and fp32 tensors.")
+    if k.dtype != q.dtype or v.dtype != q.dtype:
+        raise ValueError("q, k, and v must have the same dtype.")
+    if g.dtype not in (torch.float16, torch.bfloat16, torch.float32):
+        raise ValueError("g must be fp16, bf16, or fp32.")
+    if beta.dtype not in (torch.float16, torch.bfloat16, torch.float32):
+        raise ValueError("beta must be fp16, bf16, or fp32.")
+    if initial_state is not None and initial_state.dtype not in (
+        torch.float16,
+        torch.bfloat16,
+        torch.float32,
+    ):
+        raise ValueError("initial_state must be fp16, bf16, or fp32.")
+    if q.ndim != 4 or k.ndim != 4 or v.ndim != 4:
+        raise ValueError("q, k, and v must have shape [B, T, H, D].")
+    if g.ndim != 3 or beta.ndim != 3:
+        raise ValueError("g and beta must have shape [B, T, Hv].")
+    if q.shape != k.shape:
+        raise ValueError("q and k must have the same shape.")
+
+    batch, tokens, q_heads, k_dim = q.shape
+    _, _, v_heads, v_dim = v.shape
+    if v.shape[0] != batch or v.shape[1] != tokens:
+        raise ValueError("v must have shape [B, T, Hv, V] matching q/k.")
+    if g.shape != beta.shape or g.shape != v.shape[:3]:
+        raise ValueError("g and beta must have shape [B, T, Hv].")
+    if v_heads % q_heads != 0:
+        raise ValueError("Hv must be divisible by Hq.")
+    if k_dim != 128 or v_dim != 128:
+        raise ValueError("SM70 FlashQLA backend currently supports K=V=128.")
+    if initial_state is not None and initial_state.shape != (
+        batch,
+        v_heads,
+        k_dim,
+        v_dim,
+    ):
+        raise ValueError("initial_state must have shape [B, Hv, K, V].")
+
+
+def _check_vlk_varlen_inputs(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    g: torch.Tensor,
+    beta: torch.Tensor,
+    initial_state: torch.Tensor | None,
+    cu_seqlens: torch.Tensor,
+    output: torch.Tensor | None = None,
+    validate_cu_seqlens: bool = True,
+) -> None:
+    tensors = [q, k, v, g, beta, cu_seqlens]
+    if initial_state is not None:
+        tensors.append(initial_state)
+    if output is not None:
+        tensors.append(output)
+    if any(not tensor.is_cuda for tensor in tensors):
+        raise ValueError("SM70 GDN tensors must be CUDA tensors.")
+    if any(tensor.device != q.device for tensor in tensors):
+        raise ValueError("SM70 GDN tensors must be on the same CUDA device.")
+    if any(not tensor.is_contiguous() for tensor in tensors):
+        raise ValueError("SM70 GDN tensors must be contiguous.")
+    if q.dtype not in (torch.float16, torch.bfloat16, torch.float32):
+        raise ValueError("SM70 GDN backend supports fp16, bf16, and fp32 tensors.")
+    if k.dtype != q.dtype or v.dtype != q.dtype:
+        raise ValueError("q, k, and v must have the same dtype.")
+    if g.dtype not in (torch.float16, torch.bfloat16, torch.float32):
+        raise ValueError("g must be fp16, bf16, or fp32.")
+    if beta.dtype not in (torch.float16, torch.bfloat16, torch.float32):
+        raise ValueError("beta must be fp16, bf16, or fp32.")
+    if initial_state is not None and initial_state.dtype not in (
+        torch.float16,
+        torch.bfloat16,
+        torch.float32,
+    ):
+        raise ValueError("initial_state must be fp16, bf16, or fp32.")
+    if cu_seqlens.dtype != torch.int32:
+        raise ValueError("cu_seqlens must be int32.")
+    if q.ndim != 4 or k.ndim != 4 or v.ndim != 4:
+        raise ValueError("q, k, and v must have shape [1, T, H, D].")
+    if q.shape[0] != 1:
+        raise ValueError("SM70 varlen GDN expects flattened q/k/v with batch=1.")
+    if g.ndim != 3 or beta.ndim != 3:
+        raise ValueError("g and beta must have shape [1, T, Hv].")
+    if cu_seqlens.ndim != 1 or cu_seqlens.numel() < 2:
+        raise ValueError("cu_seqlens must have shape [N + 1].")
+    if q.shape != k.shape:
+        raise ValueError("q and k must have the same shape.")
+
+    _, tokens, q_heads, k_dim = q.shape
+    _, _, v_heads, v_dim = v.shape
+    num_sequences = cu_seqlens.numel() - 1
+    if v.shape[0] != 1 or v.shape[1] != tokens:
+        raise ValueError("v must have shape [1, T, Hv, V] matching q/k.")
+    if g.shape != beta.shape or g.shape != v.shape[:3]:
+        raise ValueError("g and beta must have shape [1, T, Hv].")
+    if v_heads % q_heads != 0:
+        raise ValueError("Hv must be divisible by Hq.")
+    if k_dim != 128 or v_dim != 128:
+        raise ValueError("SM70 FlashQLA backend currently supports K=V=128.")
+    if initial_state is not None and initial_state.shape != (
+        num_sequences,
+        v_heads,
+        v_dim,
+        k_dim,
+    ):
+        raise ValueError("initial_state must have shape [N, Hv, V, K].")
+    if output is not None:
+        if output.dtype != v.dtype:
+            raise ValueError("output must match v dtype.")
+        if output.shape != (1, tokens, v_heads, v_dim):
+            raise ValueError("output must have shape [1, T, Hv, V].")
+    if validate_cu_seqlens:
+        cu_cpu = cu_seqlens.detach().cpu()
+        if int(cu_cpu[0]) != 0:
+            raise ValueError("cu_seqlens must start at 0.")
+        if int(cu_cpu[-1]) != tokens:
+            raise ValueError("cu_seqlens must end at the flattened token count.")
+        if not bool((cu_cpu[1:] >= cu_cpu[:-1]).all()):
+            raise ValueError("cu_seqlens must be non-decreasing.")
+
+
+def chunk_gated_delta_rule_fwd_sm70(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    g: torch.Tensor,
+    beta: torch.Tensor,
+    scale: float | None = None,
+    initial_state: torch.Tensor | None = None,
+    output_final_state: bool = True,
+    gate_is_exp: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+    """Run the experimental SM70/SM75 forward GDN backend.
+
+    This keeps the public FlashQLA tensor contract:
+    q/k: [B, T, Hq, K], v/o: [B, T, Hv, V], state: [B, Hv, K, V].
+    """
+
+    _check_inputs(q, k, v, g, beta, initial_state)
+    if scale is None:
+        scale = q.shape[-1] ** -0.5
+    ext = _load_ext()
+    output, final_state = ext.gdn_forward(
+        q,
+        k,
+        v,
+        g,
+        beta,
+        initial_state,
+        float(scale),
+        output_final_state,
+        gate_is_exp,
+    )
+    if not output_final_state:
+        final_state = None
+    return output, final_state
+
+
+def chunk_gated_delta_rule_fwd_sm70_vlk_varlen(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    g: torch.Tensor,
+    beta: torch.Tensor,
+    cu_seqlens: torch.Tensor,
+    scale: float | None = None,
+    initial_state: torch.Tensor | None = None,
+    output_final_state: bool = True,
+    validate_cu_seqlens: bool = True,
+    gate_is_exp: bool = False,
+    output: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+    """Run SM70/SM75 forward with vLLM-only state layout [N, Hv, V, K].
+
+    This is not a public FlashQLA varlen drop-in: K and V are both 128 for
+    Qwen GDN, so the vLLM layout cannot be shape-distinguished from the public
+    [N, Hv, K, V] contract.
+    """
+
+    _check_vlk_varlen_inputs(
+        q,
+        k,
+        v,
+        g,
+        beta,
+        initial_state,
+        cu_seqlens,
+        output,
+        validate_cu_seqlens=validate_cu_seqlens,
+    )
+    if scale is None:
+        scale = q.shape[-1] ** -0.5
+    ext = _load_ext()
+    output, final_state = ext.gdn_forward_vlk_varlen(
+        q,
+        k,
+        v,
+        g,
+        beta,
+        initial_state,
+        cu_seqlens,
+        float(scale),
+        output_final_state,
+        validate_cu_seqlens,
+        gate_is_exp,
+        output,
+    )
+    if not output_final_state:
+        final_state = None
+    return output, final_state
+
+
+def gdn_decode_mixed_qkv_global_state_sm70(
+    mixed_qkv: torch.Tensor,
+    a: torch.Tensor,
+    b: torch.Tensor,
+    A_log: torch.Tensor,
+    dt_bias: torch.Tensor,
+    state: torch.Tensor,
+    state_indices: torch.Tensor,
+    output: torch.Tensor,
+    scale: float | None = None,
+    use_qk_l2norm_in_kernel: bool = True,
+) -> torch.Tensor:
+    """Run fused SM70 mixed-QKV decode against vLLM global state slots."""
+
+    tensors = [mixed_qkv, a, b, A_log, dt_bias, state, state_indices, output]
+    if any(not tensor.is_cuda for tensor in tensors):
+        raise ValueError("SM70 GDN decode tensors must be CUDA tensors.")
+    if any(tensor.device != mixed_qkv.device for tensor in tensors):
+        raise ValueError("SM70 GDN decode tensors must be on the same CUDA device.")
+    contiguous_tensors = {
+        "a": a,
+        "b": b,
+        "A_log": A_log,
+        "dt_bias": dt_bias,
+        "state_indices": state_indices,
+        "output": output,
+    }
+    non_contiguous = [
+        name
+        for name, tensor in contiguous_tensors.items()
+        if not tensor.is_contiguous()
+    ]
+    if non_contiguous:
+        raise ValueError(
+            "SM70 GDN decode tensors must be contiguous except mixed_qkv/state; "
+            f"non-contiguous={non_contiguous}"
+        )
+    if mixed_qkv.dtype not in (torch.float16, torch.bfloat16, torch.float32):
+        raise ValueError("mixed_qkv must be fp16, bf16, or fp32.")
+    if a.dtype != mixed_qkv.dtype or b.dtype != mixed_qkv.dtype:
+        raise ValueError("a and b must match mixed_qkv dtype.")
+    if output.dtype != mixed_qkv.dtype:
+        raise ValueError("output must match mixed_qkv dtype.")
+    if A_log.dtype != torch.float32:
+        raise ValueError("A_log must be float32.")
+    if dt_bias.dtype not in (torch.float16, torch.bfloat16, torch.float32):
+        raise ValueError("dt_bias must be fp16, bf16, or fp32.")
+    if state_indices.dtype != torch.int32:
+        raise ValueError("state_indices must be int32.")
+    if mixed_qkv.ndim != 2 or a.ndim != 2 or b.ndim != 2:
+        raise ValueError("mixed_qkv, a, and b must be rank-2 tensors.")
+    if mixed_qkv.stride(1) != 1 or mixed_qkv.stride(0) < mixed_qkv.shape[1]:
+        raise ValueError(
+            "mixed_qkv must have dense columns and row stride >= logical width; "
+            f"shape={tuple(mixed_qkv.shape)} stride={tuple(mixed_qkv.stride())}"
+        )
+    if state.ndim != 4 or output.ndim != 3:
+        raise ValueError("state must be [slots,Hv,V,K], output [T,Hv,V].")
+    tokens = mixed_qkv.shape[0]
+    _, v_heads, v_dim, k_dim = state.shape
+    if k_dim != 128 or v_dim != 128:
+        raise ValueError("SM70 FlashQLA decode currently supports K=V=128.")
+    if state.stride()[1:] != (v_dim * k_dim, k_dim, 1):
+        raise ValueError(
+            "state inner layout must be [slots,Hv,V,K] with contiguous [Hv,V,K] "
+            f"pages; got stride={tuple(state.stride())}"
+        )
+    if a.shape != (tokens, v_heads) or b.shape != (tokens, v_heads):
+        raise ValueError("a/b must have shape [T,Hv].")
+    if A_log.shape != (v_heads,) or dt_bias.shape != (v_heads,):
+        raise ValueError("A_log/dt_bias must have shape [Hv].")
+    if state_indices.shape != (tokens,):
+        raise ValueError("state_indices must have shape [T].")
+    if output.shape != (tokens, v_heads, v_dim):
+        raise ValueError("output must have shape [T,Hv,V].")
+    if scale is None:
+        scale = k_dim**-0.5
+    ext = _load_ext()
+    ext.gdn_decode_mixed_qkv_global_state(
+        mixed_qkv,
+        a,
+        b,
+        A_log,
+        dt_bias,
+        state,
+        state_indices,
+        output,
+        float(scale),
+        bool(use_qk_l2norm_in_kernel),
+    )
+    return output
+
+
+def gdn_decode_mixed_qkv_ddtree_state_sm70(
+    mixed_qkv: torch.Tensor,
+    a: torch.Tensor,
+    b: torch.Tensor,
+    A_log: torch.Tensor,
+    dt_bias: torch.Tensor,
+    state: torch.Tensor,
+    state_indices: torch.Tensor,
+    parent_ids: torch.Tensor,
+    num_accepted_tokens: torch.Tensor,
+    cu_seqlens: torch.Tensor,
+    output: torch.Tensor,
+    scale: float | None = None,
+    use_qk_l2norm_in_kernel: bool = True,
+) -> torch.Tensor:
+    """Run parent-aware DDTree mixed-QKV decode against vLLM global state."""
+
+    tensors = [
+        mixed_qkv,
+        a,
+        b,
+        A_log,
+        dt_bias,
+        state,
+        state_indices,
+        parent_ids,
+        num_accepted_tokens,
+        cu_seqlens,
+        output,
+    ]
+    if any(not tensor.is_cuda for tensor in tensors):
+        raise ValueError("SM70 DDTree GDN tensors must be CUDA tensors.")
+    if any(tensor.device != mixed_qkv.device for tensor in tensors):
+        raise ValueError("SM70 DDTree GDN tensors must be on the same CUDA device.")
+    contiguous_tensors = {
+        "a": a,
+        "b": b,
+        "A_log": A_log,
+        "dt_bias": dt_bias,
+        "state_indices": state_indices,
+        "parent_ids": parent_ids,
+        "num_accepted_tokens": num_accepted_tokens,
+        "cu_seqlens": cu_seqlens,
+        "output": output,
+    }
+    non_contiguous = [
+        name
+        for name, tensor in contiguous_tensors.items()
+        if not tensor.is_contiguous()
+    ]
+    if non_contiguous:
+        raise ValueError(
+            "SM70 DDTree GDN tensors must be contiguous except mixed_qkv/state; "
+            f"non-contiguous={non_contiguous}"
+        )
+    if mixed_qkv.dtype not in (torch.float16, torch.bfloat16, torch.float32):
+        raise ValueError("mixed_qkv must be fp16, bf16, or fp32.")
+    if a.dtype != mixed_qkv.dtype or b.dtype != mixed_qkv.dtype:
+        raise ValueError("a and b must match mixed_qkv dtype.")
+    if output.dtype != mixed_qkv.dtype:
+        raise ValueError("output must match mixed_qkv dtype.")
+    if A_log.dtype != torch.float32:
+        raise ValueError("A_log must be float32.")
+    if dt_bias.dtype not in (torch.float16, torch.bfloat16, torch.float32):
+        raise ValueError("dt_bias must be fp16, bf16, or fp32.")
+    if state_indices.dtype != torch.int32:
+        raise ValueError("state_indices must be int32.")
+    if parent_ids.dtype != torch.int32:
+        raise ValueError("parent_ids must be int32.")
+    if num_accepted_tokens.dtype != torch.int32:
+        raise ValueError("num_accepted_tokens must be int32.")
+    if cu_seqlens.dtype != torch.int32:
+        raise ValueError("cu_seqlens must be int32.")
+    if mixed_qkv.ndim != 2 or a.ndim != 2 or b.ndim != 2:
+        raise ValueError("mixed_qkv, a, and b must be rank-2 tensors.")
+    if mixed_qkv.stride(1) != 1 or mixed_qkv.stride(0) < mixed_qkv.shape[1]:
+        raise ValueError(
+            "mixed_qkv must have dense columns and row stride >= logical width; "
+            f"shape={tuple(mixed_qkv.shape)} stride={tuple(mixed_qkv.stride())}"
+        )
+    if state.ndim != 4 or output.ndim != 3:
+        raise ValueError("state must be [slots,Hv,V,K], output [T,Hv,V].")
+    if state_indices.ndim != 2 or parent_ids.ndim != 2:
+        raise ValueError("state_indices and parent_ids must be rank-2 tensors.")
+    if parent_ids.shape != state_indices.shape:
+        raise ValueError("parent_ids must match state_indices shape.")
+    tokens = mixed_qkv.shape[0]
+    num_sequences = state_indices.shape[0]
+    _, v_heads, v_dim, k_dim = state.shape
+    if k_dim != 128 or v_dim != 128:
+        raise ValueError("SM70 FlashQLA DDTree decode currently supports K=V=128.")
+    if state.stride()[1:] != (v_dim * k_dim, k_dim, 1):
+        raise ValueError(
+            "state inner layout must be [slots,Hv,V,K] with contiguous [Hv,V,K] "
+            f"pages; got stride={tuple(state.stride())}"
+        )
+    if a.shape != (tokens, v_heads) or b.shape != (tokens, v_heads):
+        raise ValueError("a/b must have shape [T,Hv].")
+    if A_log.shape != (v_heads,) or dt_bias.shape != (v_heads,):
+        raise ValueError("A_log/dt_bias must have shape [Hv].")
+    if num_accepted_tokens.shape != (num_sequences,):
+        raise ValueError("num_accepted_tokens must have shape [N].")
+    if cu_seqlens.shape != (num_sequences + 1,):
+        raise ValueError("cu_seqlens must have shape [N + 1].")
+    if output.shape != (tokens, v_heads, v_dim):
+        raise ValueError("output must have shape [T,Hv,V].")
+    if scale is None:
+        scale = k_dim**-0.5
+    ext = _load_ext()
+    ext.gdn_decode_mixed_qkv_ddtree_state(
+        mixed_qkv,
+        a,
+        b,
+        A_log,
+        dt_bias,
+        state,
+        state_indices,
+        parent_ids,
+        num_accepted_tokens,
+        cu_seqlens,
+        output,
+        float(scale),
+        bool(use_qk_l2norm_in_kernel),
+    )
+    return output
+
+
+def resolve_column_groups_per_block_sm70(
+    tokens: int,
+    q_heads: int,
+    v_heads: int,
+) -> int:
+    ext = _load_ext()
+    return int(ext.resolve_column_groups_per_block(tokens, q_heads, v_heads))

qwen3_6_scripts/flash_qla_sm70/naive_gdn.py

@@ -0,0 +1,161 @@
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+# For a list of all contributors, visit:
+#   https://github.com/fla-org/flash-linear-attention/graphs/contributors
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+
+
+def naive_recurrent_gated_delta_rule(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    beta: torch.Tensor,
+    g: torch.Tensor,
+    scale: float = None,
+    initial_state: torch.Tensor = None,
+    output_final_state: bool = False,
+):
+    """
+    Reference PyTorch implementation of recurrent gated delta rule.
+
+    Args:
+        q: [B, T, H, K]
+        k: [B, T, H, K]
+        v: [B, T, H, V]
+        beta: [B, T, H]
+        g: [B, T, H]
+        scale: float, optional
+        initial_state: [B, H, K, V], optional
+        output_final_state: bool
+
+    Returns:
+        o: [B, T, H, V]
+        final_state: [B, H, K, V] if output_final_state else None
+    """
+    q, k, v, beta, g = map(lambda x: x.transpose(1, 2).contiguous().to(torch.float32), [q, k, v, beta, g])
+    B, H, T, K, V = *k.shape, v.shape[-1]
+    o = torch.zeros(B, H, T, V).to(v)
+    h = torch.zeros(B, H, K, V).to(v)
+    if initial_state is not None:
+        h = initial_state.to(torch.float32)
+    if scale is None:
+        scale = 1 / (q.shape[-1] ** 0.5)
+    q = q * scale
+
+    for i in range(T):
+        b_q = q[:, :, i]
+        b_k = k[:, :, i]
+        b_v = v[:, :, i].clone()
+        h = h.clone() * g[:, :, i].exp()[..., None, None]
+        b_beta = beta[:, :, i]
+        b_v = b_v - (h.clone() * b_k[..., None]).sum(-2)
+        b_v = b_v * b_beta[..., None]
+        h = h.clone() + b_k.unsqueeze(-1) * b_v.unsqueeze(-2)
+        o[:, :, i] = torch.einsum('bhd,bhdm->bhm', b_q, h)
+
+    if not output_final_state:
+        h = None
+    o = o.transpose(1, 2).contiguous()
+    return o, h
+
+
+def naive_chunk_gated_delta_rule(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    g: torch.Tensor,
+    beta: torch.Tensor,
+    chunk_size: int = 64,
+    scale: float = None,
+    initial_state: torch.Tensor = None,
+    output_final_state: bool = False,
+):
+    """
+    Reference PyTorch implementation of chunk gated delta rule.
+
+    Args:
+        q: [B, T, H, K]
+        k: [B, T, H, K]
+        v: [B, T, H, V]
+        g: [B, T, H]
+        beta: [B, T, H]
+        chunk_size: int
+        scale: float, optional
+        initial_state: [B, H, K, V], optional
+        output_final_state: bool
+
+    Returns:
+        o: [B, T, H, V]
+        final_state: [B, H, K, V] if output_final_state else None
+    """
+    BT = chunk_size
+    if scale is None:
+        scale = 1 / (q.shape[-1] ** 0.5)
+
+    q, k, v, beta, g = map(lambda x: x.transpose(1, 2).contiguous().to(torch.float32), [q, k, v, beta, g])
+
+    T = q.shape[-2]
+    pad_len = (BT - (T % BT)) % BT
+    if pad_len > 0:
+        q = F.pad(q, (0, 0, 0, pad_len))
+        k = F.pad(k, (0, 0, 0, pad_len))
+        v = F.pad(v, (0, 0, 0, pad_len))
+        beta = F.pad(beta, (0, pad_len))
+        g = F.pad(g, (0, pad_len))
+
+    q, k, v, beta, g = map(lambda x: x.to(torch.float32), [q, k, v, beta, g])
+    decay = g
+    chunk_size = BT
+    b, h, l, d_k = q.shape
+    d_v = v.shape[-1]
+    q = q * scale
+    v = v * beta[..., None]
+    k_beta = k * beta[..., None]
+    assert l % chunk_size == 0
+
+    # note that diagonal is masked.
+    mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=0)
+    q, k, v, k_beta, decay = map(
+        lambda x: rearrange(x, 'b h (n c) d -> b h n c d', c=chunk_size),
+        [q, k, v, k_beta, decay.unsqueeze(-1)],
+    )
+    decay = decay.squeeze(-1).cumsum(-1)
+    decay_exp = decay.exp()[..., None]
+    L_mask = ((decay.unsqueeze(-1) - decay.unsqueeze(-2)).tril().exp().float()).tril()
+    attn = -((k_beta @ k.transpose(-1, -2)) * L_mask).masked_fill(mask, 0)
+    for i in range(1, chunk_size):
+        attn[..., i, :i] = attn[..., i, :i].clone() + (attn[..., i, :i, None].clone() * attn[..., :i, :i].clone()).sum(-2)
+    attn = attn + torch.eye(chunk_size, dtype=torch.float, device=q.device)
+    attn = attn
+    k_cumsum = attn @ v
+    k_cumdecay = attn @ (k_beta * decay_exp)
+    v = k_cumsum
+
+    S = k.new_zeros(b, h, d_k, d_v)
+    if initial_state is not None:
+        S = initial_state.to(torch.float32)
+
+    o = torch.zeros_like(v)
+    mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=1)
+    for i in range(0, l // chunk_size):
+        q_i, k_i, v_i = q[:, :, i], k[:, :, i], v[:, :, i]
+        attn = (q_i @ k_i.transpose(-1, -2) * L_mask[:, :, i]).masked_fill_(mask, 0)
+        v_prime = (k_cumdecay[:, :, i]) @ S
+        v_new = v_i - v_prime
+        o_inter = (q_i * decay[:, :, i, :, None].exp()) @ S
+        o[:, :, i] = o_inter + attn @ v_new
+        S = S * decay[:, :, i, -1, None, None].exp() + (k_i * (decay[:, :, i, -1, None] - decay[:, :, i]).exp()
+                                                        [..., None]).transpose(-1, -2) @ v_new
+    if not output_final_state:
+        S = None
+
+    # unpad
+    o = rearrange(o, 'b h n c d -> b h (n c) d')
+    o = o[:, :, :T]
+    o = o.transpose(1, 2)
+    return o, S

qwen3_6_scripts/patch_ops.sh

@@ -175,6 +175,25 @@ if [ -n "$VLLM2" ]; then
     cp ./chat_utils.py "$VLLM2/entrypoints/chat_utils.py" 2>/dev/null || true
 fi
 
-echo "[patch_ops] DONE — all patches deployed"
-echo "[patch_ops] Deployed: qwen3_5.py, paged_attn.py, mamba_cache.py, sequence.py, scheduler.py, xformers patches, tool/reasoning parsers, serving layer"
-echo "[patch_ops] NOT deployed (using base image native): model_runner.py, _custom_ops.py, sampler.py, logits_processor.py, arg_utils.py"
+echo "[patch_ops] DONE — SM70 GDN kernel + serving layer + engine patches deployed"
+echo "[patch_ops] Deployed: qwen3_5.py, flash_qla_sm70 (SM70 GDN CUDA kernel), paged_attn.py, mamba_cache.py, sequence.py, scheduler.py, xformers patches, serving layer"
+echo "[patch_ops] SM70 GDN kernel: JIT compiles on first forward pass (~2min), then cached"
+echo "[patch_ops] NOT deployed (base image native): model_runner.py, _custom_ops.py, sampler.py, logits_processor.py, arg_utils.py"
+
+# Deploy flash_qla SM70 GDN kernel (from 1Cat-vLLM, MIT license)
+# This is a fused CUDA kernel for GatedDeltaNet on SM70/SM75 (V100/BI-V100)
+# JIT compiled at runtime via torch.utils.cpp_extension.load()
+FLASH_QLA_DST="$VLLM/model_executor/models/flash_qla_sm70"
+if [ -d "./flash_qla_sm70" ]; then
+    rm -rf "$FLASH_QLA_DST" 2>/dev/null
+    cp -r ./flash_qla_sm70 "$FLASH_QLA_DST" 2>/dev/null && \
+        echo "[patch_ops] flash_qla_sm70 deployed to $FLASH_QLA_DST" || true
+    # Pre-compile CUDA kernel → .so (skipped if no GPU/compiler at build time)
+    python3 ./precompile_gdn.py "$FLASH_QLA_DST" 2>&1 || \
+        echo "[patch_ops] WARNING: precompile failed — kernel will JIT at runtime"
+    # Also deploy to VLLM2 if present
+    if [ -n "$VLLM2" ]; then
+        rm -rf "$VLLM2/model_executor/models/flash_qla_sm70" 2>/dev/null
+        cp -r "$FLASH_QLA_DST" "$VLLM2/model_executor/models/flash_qla_sm70" 2>/dev/null || true
+    fi
+fi

qwen3_6_scripts/precompile_gdn.py

@@ -0,0 +1,53 @@
+"""
+Pre-compile SM70 GDN CUDA kernel → .so at Docker build time.
+Avoids 2-minute JIT delay at runtime.
+
+Usage: python3 precompile_gdn.py /path/to/flash_qla_sm70/
+"""
+import os
+import sys
+
+def main():
+    if len(sys.argv) < 2:
+        print("[precompile] Usage: python3 precompile_gdn.py <flash_qla_sm70_dir>")
+        sys.exit(1)
+
+    flash_dir = sys.argv[1]
+    cu_src = os.path.join(flash_dir, "csrc", "gdn_forward.cu")
+    if not os.path.exists(cu_src):
+        print(f"[precompile] ERROR: {cu_src} not found")
+        sys.exit(1)
+
+    # Set arch for BI-V100 (SM70 compatible)
+    os.environ["TORCH_CUDA_ARCH_LIST"] = "7.0;7.5"
+
+    build_dir = os.path.join(flash_dir, "build")
+    os.makedirs(build_dir, exist_ok=True)
+
+    print(f"[precompile] Compiling {cu_src} → .so in {build_dir}")
+    print(f"[precompile] TORCH_CUDA_ARCH_LIST = {os.environ['TORCH_CUDA_ARCH_LIST']}")
+
+    try:
+        from torch.utils.cpp_extension import load
+        ext = load(
+            name="flash_qla_sm70_gdn_strided",
+            sources=[cu_src],
+            extra_cuda_cflags=["-O3"],
+            extra_cflags=["-O3"],
+            build_directory=build_dir,
+            verbose=True,
+        )
+        print(f"[precompile] SUCCESS — compiled .so in {build_dir}")
+        # List the built files
+        for f in os.listdir(build_dir):
+            if f.endswith(".so"):
+                full = os.path.join(build_dir, f)
+                print(f"[precompile]   {f} ({os.path.getsize(full)} bytes)")
+    except Exception as e:
+        print(f"[precompile] FAILED: {e}")
+        print("[precompile] Kernel will JIT compile at runtime instead (~2min)")
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

qwen3_6_scripts/qwen3_5.py

@@ -76,19 +76,35 @@ _corex_moe_module = None
 _corex_gdn_available = False
 _corex_moe_available = False
 
+# SM70 FlashQLA GDN kernel (from 1Cat-vLLM, MIT license)
+# Fused CUDA kernel for GatedDeltaNet on SM70/SM75 (V100/BI-V100)
+# JIT compiled via torch.utils.cpp_extension.load() on first call
+_flash_qla_sm70 = None
+_flash_qla_available = False
+
+try:
+    from vllm.model_executor.models.flash_qla_sm70 import (
+        chunk_gated_delta_rule_fwd_sm70,
+        chunk_gated_delta_rule_fwd_sm70_vlk_varlen,
+    )
+    _flash_qla_available = True
+    logger.info("FlashQLA SM70 GDN module found — fused CUDA kernel available (JIT on first call)")
+except ImportError as e:
+    logger.warning("FlashQLA SM70 GDN not found (%s) — using PyTorch GDN", e)
+
 try:
     from vllm.model_executor.models import corex_gdn as _corex_gdn_module
     _corex_gdn_available = True
     logger.info("CoreX GDN module found — fused GDN kernels available")
 except ImportError:
-    pass  # expected if not packaged; ixformer ops used instead
+    pass
 
 try:
     from vllm.model_executor.models import corex_moe as _corex_moe_module
     _corex_moe_available = True
     logger.info("CoreX MoE module found — fused MoE kernels available")
 except ImportError:
-    pass  # expected; MoE uses PyTorch loop
+    pass
 
 
 # ---------------------------------------------------------------------------
@@ -155,7 +171,12 @@ def _torch_chunk_gated_delta_rule(
     value: torch.Tensor,   # (batch, seq, num_heads, head_v_dim)
     g: torch.Tensor,       # (batch, seq, num_heads)
     beta: torch.Tensor,    # (batch, seq, num_heads)
-    chunk_size: int = 64,
+    # CCCL agent_radix_sort_upsweep overflow pattern: UNROLL_COUNT = min(64, 255/KEYS_PER_THREAD)
+    # prevents counter overflow by limiting accumulation steps.
+    # Same principle: chunk_size limits cumsum steps. With pre-clamp [-5,2]:
+    #   chunk=64: worst cumsum = 64*2 = 128 → exp(128) = inf
+    #   chunk=16: worst cumsum = 16*2 = 32  → clamp(-20,20) catches it
+    chunk_size: int = 16,
     initial_state: Optional[torch.Tensor] = None,
     output_final_state: bool = False,
     use_qk_l2norm_in_kernel: bool = False,
@@ -218,17 +239,34 @@ def _torch_chunk_gated_delta_rule(
         torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device),
         diagonal=1)
 
-    for i in range(total_len // chunk_size):
-        q_i, k_i, v_i = query[:, :, i], key[:, :, i], value[:, :, i]
-        attn_i = (_ix_matmul(q_i, k_i.transpose(-1, -2)) * decay_mask[:, :, i]).masked_fill_(mask_upper2, 0)
+    # dispatch_scan.cuh Phase 1: pre-compute ALL chunk-local attention matrices
+    # outside the state loop. attn_i[c] only depends on q, k, decay_mask — NOT state.
+    # This is the CCCL "init kernel" pattern: compute everything possible
+    # before the sequential scan kernel that needs tile_state propagation.
+    num_chunks = total_len // chunk_size
+    attn_i_all = torch.empty(
+        batch, num_heads, num_chunks, chunk_size, chunk_size,
+        dtype=value.dtype, device=value.device)
+    for i in range(num_chunks):
+        attn_i_all[:, :, i] = (
+            _ix_matmul(query[:, :, i], key[:, :, i].transpose(-1, -2))
+            * decay_mask[:, :, i]
+        ).masked_fill_(mask_upper2, 0)
+
+    # dispatch_scan.cuh Phase 2: sequential state propagation (scan kernel).
+    # Only state-dependent ops remain in this loop.
+    g_exp_cache = g.clamp(-20, 20).exp()  # pre-compute once
+    g_clamped = g.clamp(-20, 20)  # keep raw clamped g for difference computation
+    for i in range(num_chunks):
         v_prime = _ix_matmul(k_cumdecay[:, :, i], last_state)
-        v_new = v_i - v_prime
-        attn_inter = _ix_matmul(q_i * g[:, :, i, :, None].clamp(-20, 20).exp(), last_state)
-        core_out[:, :, i] = attn_inter + _ix_matmul(attn_i, v_new)
+        v_new = value[:, :, i] - v_prime
+        attn_inter = _ix_matmul(query[:, :, i] * g_exp_cache[:, :, i, :, None], last_state)
+        core_out[:, :, i] = attn_inter + _ix_matmul(attn_i_all[:, :, i], v_new)
+        # State update uses difference form: exp(g[-1] - g[:]) to avoid division
         last_state = (
-            last_state * g[:, :, i, -1, None, None].clamp(-20, 20).exp()
+            last_state * g_exp_cache[:, :, i, -1, None, None]
             + _ix_matmul(
-                (k_i * (g[:, :, i, -1, None] - g[:, :, i]).clamp(-20, 20).exp()[..., None])
+                (key[:, :, i] * (g_clamped[:, :, i, -1, None] - g_clamped[:, :, i]).exp()[..., None])
                 .transpose(-1, -2), v_new)
         )
 

qwen3_6_scripts/serving_chat.py

@@ -140,27 +140,10 @@ class OpenAIServingChat(OpenAIServing):
             model_config = self.model_config
             tokenizer = await self.engine_client.get_tokenizer(lora_request)
 
-            # CCCL graceful degradation: strip image_url when not multimodal.
-            # Handle is_multimodal_model as method, property, or bool.
-            _is_mm = False
-            try:
-                _mm_attr = getattr(model_config, 'is_multimodal_model', False)
-                _is_mm = _mm_attr() if callable(_mm_attr) else bool(_mm_attr)
-            except Exception:
-                pass
-            if not _is_mm:
-                for msg in request.messages:
-                    content = msg.get("content") if isinstance(msg, dict) else getattr(msg, "content", None)
-                    if isinstance(content, list):
-                        filtered = [p for p in content
-                                    if not (isinstance(p, dict) and p.get("type") == "image_url")]
-                        if len(filtered) < len(content):
-                            if not filtered:
-                                filtered = [{"type": "text", "text": "(image omitted)"}]
-                            if isinstance(msg, dict):
-                                msg["content"] = filtered
-                            else:
-                                msg.content = filtered
+            # Note: base image identifies this model as multimodal
+            # (docker log: "--enable-prefix-caching not supported for multimodal models").
+            # Do NOT strip image_url — let images flow through to the engine.
+            # Previous strip logic caused d05_multimodal HTTP 400.
 
             conversation, mm_data_future = parse_chat_messages_futures(
                 request.messages, model_config, tokenizer)