Uploaded using `kernel-builder`.

Browse files

Files changed (5) hide show

build/torch214-cxx11-rocm72-x86_64-linux/__init__.py +185 -0
build/torch214-cxx11-rocm72-x86_64-linux/_gdn_rocm_8d6df22.abi3.so +3 -0
build/torch214-cxx11-rocm72-x86_64-linux/_ops.py +9 -0
build/torch214-cxx11-rocm72-x86_64-linux/layers.py +248 -0
build/torch214-cxx11-rocm72-x86_64-linux/metadata.json +35 -0

build/torch214-cxx11-rocm72-x86_64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,185 @@

+# SPDX-License-Identifier: AGPL-3.0-only
+"""Atlas Gated DeltaNet kernels for AMD Strix Halo (gfx1151, RDNA3.5).
+The kernels are the ones the Atlas engine ran for its MLPerf Inference v6.1
+Strix Halo submission (Qwen3.6-27B, Atlas @ eabfa8f). The first four op
+schemas match the GB10/SM121 CUDA build of ``Atlas-Inference/gdn``; ``layers``
+runs the serve chain the submission dispatched:
+  prefill: causal_conv1d_update_prefill -> l2_norm (q, k) -> gdn_prefill
+  decode:  causal_conv1d_update_l2norm_f32 -> gdn_decode_f32
+"""
+from typing import Optional
+import torch
+from ._ops import ops
+from . import layers
+__all__ = [
+    "gdn_decode",
+    "gdn_prefill",
+    "gdn_prefill_fla",
+    "causal_conv1d_fwd",
+    "causal_conv1d_update",
+    "gdn_decode_f32",
+    "causal_conv1d_update_prefill",
+    "causal_conv1d_update_l2norm_f32",
+    "l2_norm",
+    "layers",
+]
+def gdn_decode(
+    h_state: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    gate: torch.Tensor,
+    beta: torch.Tensor,
+    output: torch.Tensor,
+) -> None:
+    """Single-token GDN decode (in-place update of ``h_state`` and ``output``).
+    h_state : (B, num_v_heads, 128, 128)  float32, in-place updated
+    query   : (B, num_k_heads, 128)       float32
+    key     : (B, num_k_heads, 128)       float32
+    value   : (B, num_v_heads, 128)       float32
+    gate    : (B, num_v_heads)            float32  (exp(g_t) decay)
+    beta    : (B, num_v_heads)            float32  (sigmoid(b_t))
+    output  : (B, num_v_heads, 128)       bfloat16, in-place written
+    """
+    ops.gdn_decode(h_state, query, key, value, gate, beta, output)
+def gdn_prefill(
+    h_state: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    gate: torch.Tensor,
+    beta: torch.Tensor,
+    output: torch.Tensor,
+) -> None:
+    """Multi-token GDN prefill.
+    h_state : (B, num_v_heads, 128, 128)        float32, in-place updated
+    query   : (B, seq_len, num_k_heads, 128)    bfloat16
+    key     : (B, seq_len, num_k_heads, 128)    bfloat16
+    value   : (B, seq_len, num_v_heads, 128)    bfloat16
+    gate    : (B, seq_len, num_v_heads)         float32
+    beta    : (B, seq_len, num_v_heads)         float32
+    output  : (B, seq_len, num_v_heads, 128)    bfloat16
+    """
+    ops.gdn_prefill(h_state, query, key, value, gate, beta, output)
+def gdn_prefill_fla(
+    h_state: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    gate: torch.Tensor,
+    beta: torch.Tensor,
+    output: torch.Tensor,
+) -> None:
+    """Multi-token GDN prefill, FLA-chunked (64-token chunks, three kernels).
+    The prefill Atlas main serves on gfx1151. Same arguments and layouts as
+    ``gdn_prefill``; allocates its chunk scratch per call.
+    """
+    ops.gdn_prefill_fla(h_state, query, key, value, gate, beta, output)
+def causal_conv1d_fwd(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: Optional[torch.Tensor],
+    out: torch.Tensor,
+) -> None:
+    """Depthwise causal Conv1d + SiLU.
+    x      : (B, D, L) bfloat16
+    weight : (D, d_conv) bfloat16, d_conv <= 8
+    bias   : (D,) float32 or None
+    out    : (B, D, L) bfloat16
+    """
+    ops.causal_conv1d_fwd(x, weight, bias, out)
+def causal_conv1d_update(
+    conv_state: torch.Tensor,
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: Optional[torch.Tensor],
+    out: torch.Tensor,
+) -> None:
+    """Single-step causal Conv1d + SiLU (decode).
+    conv_state : (B, D, d_conv) float32, in-place updated (rolled left, last slot = x)
+    x          : (B, D) bfloat16
+    weight     : (D, d_conv) bfloat16
+    bias       : (D,) float32 or None
+    out        : (B, D) bfloat16
+    """
+    ops.causal_conv1d_update(conv_state, x, weight, bias, out)
+def gdn_decode_f32(
+    h_state: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    gate: torch.Tensor,
+    beta: torch.Tensor,
+    output: torch.Tensor,
+) -> None:
+    """``gdn_decode`` with a float32 ``output`` (B, num_v_heads, 128): the
+    decode kernel Atlas serves (``gated_delta_rule_decode_f32``)."""
+    ops.gdn_decode_f32(h_state, query, key, value, gate, beta, output)
+def causal_conv1d_update_prefill(
+    conv_state: torch.Tensor,
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: Optional[torch.Tensor],
+    out: torch.Tensor,
+) -> None:
+    """Multi-token causal Conv1d + SiLU scanned through a float32 window.
+    conv_state : (B, D, 4) float32, in-place updated (the last 4 inputs)
+    x          : (B, S, D) bfloat16, token-major
+    weight     : (D, 4) bfloat16
+    bias       : (D,) float32 or None
+    out        : (B, S, D) bfloat16
+    """
+    ops.causal_conv1d_update_prefill(conv_state, x, weight, bias, out)
+def causal_conv1d_update_l2norm_f32(
+    conv_state: torch.Tensor,
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: Optional[torch.Tensor],
+    out: torch.Tensor,
+    qk_channels: int,
+    head_dim: int,
+    eps: float,
+) -> None:
+    """Single-step causal Conv1d + SiLU, then per-head L2 norm of the first
+    ``qk_channels`` channels (q and k); float32 ``out`` (B, D).
+    conv_state : (B, D, d_conv) float32, in-place updated
+    x          : (B, D) bfloat16
+    head_dim   : must be 128; ``qk_channels`` a multiple of 256
+    """
+    ops.causal_conv1d_update_l2norm_f32(conv_state, x, weight, bias, out, qk_channels, head_dim, eps)
+def l2_norm(data: torch.Tensor, num_heads: int, head_dim: int, eps: float) -> None:
+    """In-place L2 norm of heads ``[0, num_heads)`` of every row of ``data``
+    (N, stride) bfloat16: ``x * rsqrt(sum(x^2) + eps)``."""
+    ops.l2_norm(data, num_heads, head_dim, eps)

build/torch214-cxx11-rocm72-x86_64-linux/_gdn_rocm_8d6df22.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:774be8a3f9c680fdab5f936a4494d9851b30798f6707b84b047899286fd04f66
+size 530576

build/torch214-cxx11-rocm72-x86_64-linux/_ops.py ADDED Viewed

	@@ -0,0 +1,9 @@

+import torch
+from . import _gdn_rocm_8d6df22
+ops = torch.ops._gdn_rocm_8d6df22
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_gdn_rocm_8d6df22::{op_name}"

build/torch214-cxx11-rocm72-x86_64-linux/layers.py ADDED Viewed

	@@ -0,0 +1,248 @@

+# SPDX-License-Identifier: AGPL-3.0-only
+"""Pure, stateless ``kernels`` layers that map the Atlas Gated DeltaNet
+kernels onto Hugging Face ``transformers``. This build targets AMD Strix Halo
+(gfx1151); the layer contract is identical to the GB10/SM121 CUDA build.
+When a model is ``kernelize()``-d, ``kernels`` binds one of these classes'
+``forward`` onto the host module instance (``MethodType(layer.forward,
+module)``), so ``self`` here is the *adopting* GatedDeltaNet module: it exposes
+the host's submodules, parameters (``A_log``, ``dt_bias``), config scalars
+(``num_v_heads``, ``key_dim`` ...), gated ``norm``, and ``out_proj``. We replace
+only the two compute cores (causal conv1d + gated-delta-rule); everything else
+is reused verbatim from the host.
+Two host architectures share the exact same GDN core and differ only in their
+input-projection layout, so the shared core lives in the module-level
+``_gdn_run`` and each ``forward`` only does its own projection preamble:
+  * ``GatedDeltaNet``        -> ``Qwen3NextGatedDeltaNet`` (Qwen3-Next-80B):
+      fused ``in_proj_qkvz`` + ``in_proj_ba``, split via the host's
+      ``fix_query_key_value_ordering``.
+  * ``Qwen3_5GatedDeltaNet`` -> ``Qwen3_5GatedDeltaNet`` (Qwen3.6-27B dense) and
+      ``Qwen3_5MoeGatedDeltaNet`` (Qwen3.6-35B-A3B): already-split
+      ``in_proj_qkv`` / ``in_proj_z`` / ``in_proj_b`` / ``in_proj_a``,
+      no ordering fixup.
+``kernels`` forbids extra class members and a custom ``__init__`` on a layer
+(``_validate_layer``), which is why all helpers are module-level functions, not
+methods. ``_validate_layer`` also requires the layer ``forward`` signature to
+match the host's argument count exactly, so ``forward`` takes the same
+``**kwargs`` (``Unpack[TransformersKwargs]``) the host GDN layers carry in
+transformers >= 5.10; the kernel path ignores those kwargs.
+On Strix Halo (gfx1151), as on the DGX Spark, the upstream ``fla`` /
+``causal_conv1d`` fast paths have no build, so ``transformers`` silently falls
+back to a slow pure-torch implementation. These kernels fill exactly that gap.
+The conv / norm / recurrence chain is the one the Atlas Strix Halo serve
+dispatches for a single sequence:
+  prefill: causal_conv1d_update_prefill -> l2_norm (q, k) -> gdn_prefill_fla
+           (FLA chunked, Atlas main's gfx1151 default; ATLAS_GDN_FLA_GFX=0
+           selects the split4 recurrence the MLPerf v6.1 submission ran)
+  decode:  causal_conv1d_update_l2norm_f32 -> gdn_decode_f32 (fp32 into the norm)
+Convention (pinned against transformers' torch reference):
+  * q,k are L2-normalized before the recurrence (the kernel applies 1/sqrt(d))
+  * gate = exp(g),  g = -A_log.exp() * softplus(a + dt_bias)   (per token)
+  * beta = sigmoid(b)
+  * the causal conv1d kernels apply SiLU internally
+  * recurrent state h: fp32 [B, num_v_heads, head_k_dim, head_v_dim]
+"""
+import os
+import torch
+import torch.nn.functional as F
+from torch import nn
+from ._ops import ops
+def _l2norm(x, eps: float = 1e-6):
+    # Matches transformers' fla-aligned l2norm (eps inside the rsqrt).
+    return x * torch.rsqrt((x * x).sum(dim=-1, keepdim=True) + eps)
+def _apply_mask_to_padding_states(hidden_states, attention_mask):
+    if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1:
+        dtype = hidden_states.dtype
+        hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
+    return hidden_states
+def _state(states):
+    # transformers >= 5.17: dict keyed by state_idx; earlier: the tensor itself.
+    return states[0] if isinstance(states, dict) else states
+def _gdn_run(host, hidden_states, cache_params, mixed_qkv, z, b, a):
+    """Shared GDN compute core (conv1d + gated delta rule + gated RMSNorm).
+    Architecture-agnostic: the caller supplies the already-projected tensors.
+    Parameters
+    ----------
+    host        : the adopting GatedDeltaNet module (source of submodules/dims)
+    mixed_qkv   : [B, conv_dim, S] bf16  (concatenated Q|K|V, pre-conv)
+    z           : [B, S, num_v_heads, head_v_dim]  gate for the output RMSNorm
+    b, a        : [B, S, num_v_heads]  raw beta / decay projections
+    """
+    batch_size, seq_len, _ = hidden_states.shape
+    layer = cache_params.layers[host.layer_idx] if cache_params is not None else None
+    has_prev = cache_params is not None and cache_params.has_previous_state(host.layer_idx)
+    # transformers >= 5.17 keys per-layer states by state_idx and can keep the
+    # full conv history for speculative rollback (record_past).
+    keyed = layer is not None and isinstance(layer.conv_states, dict)
+    record_past = bool(getattr(layer, "record_past", False))
+    use_precomputed_states = has_prev and seq_len == 1
+    conv_w = host.conv1d.weight.squeeze(1).contiguous()  # [conv_dim, K], bf16
+    K = host.conv_kernel_size
+    qk_channels = 2 * host.key_dim
+    # --- causal conv1d (+SiLU), then L2 norm of q/k: the Atlas serve chain ---
+    if use_precomputed_states and not record_past:
+        # decode: roll the fp32 window, conv + SiLU + q/k L2 norm in one kernel,
+        # fp32 out (causal_conv1d_update_l2norm_f32).
+        conv_state = _state(layer.conv_states)
+        cs = conv_state.to(torch.float32).contiguous()
+        x_step = mixed_qkv[:, :, 0].contiguous()  # [B, conv_dim]
+        conv_out = torch.empty(x_step.shape, device=x_step.device, dtype=torch.float32)
+        ops.causal_conv1d_update_l2norm_f32(
+            cs, x_step, conv_w, None, conv_out, qk_channels, host.head_k_dim, 1e-6
+        )
+        conv_state.copy_(cs.to(conv_state.dtype))  # persist rolled window
+        mixed_qkv = conv_out.unsqueeze(1)  # [B, 1, conv_dim] fp32, q/k normalized
+    else:
+        # prefill (or a record_past step): the fp32 window is seeded from the
+        # cached left context, then causal_conv1d_update_prefill scans the new
+        # tokens with it in registers.
+        if keyed:
+            # Cached left context + new tokens (just the new tokens, zero-padded
+            # to K, on a fresh prefill); the cache itself is updated here.
+            full = cache_params.update_conv_state(mixed_qkv, host.layer_idx, conv_kernel_size=K)
+            left = full[..., : full.shape[-1] - seq_len]
+        else:
+            if cache_params is not None:
+                padded = F.pad(mixed_qkv, (K - mixed_qkv.shape[-1], 0))
+                cache_params.update_conv_state(padded, host.layer_idx)
+            left = mixed_qkv[..., :0]
+        left = left[..., -K:]
+        cs = F.pad(left, (K - left.shape[-1], 0)).to(torch.float32).contiguous()  # [B, conv_dim, K]
+        x = mixed_qkv.transpose(1, 2).contiguous()  # [B, S, conv_dim], token-major as Atlas lays it
+        conv_out = torch.empty_like(x)
+        ops.causal_conv1d_update_prefill(cs, x, conv_w, None, conv_out)
+        # q,k L2 norm in place (l2_norm_bf16), heads [0, 2*num_k_heads) of each token row
+        ops.l2_norm(conv_out.view(-1, conv_out.shape[-1]), 2 * host.num_k_heads, host.head_k_dim, 1e-6)
+        mixed_qkv = conv_out  # [B, S, conv_dim] bf16, q/k normalized
+    query, key, value = torch.split(
+        mixed_qkv, [host.key_dim, host.key_dim, host.value_dim], dim=-1
+    )
+    # q/k stay at num_k_heads: the kernels map v-head -> k-head themselves.
+    query = query.reshape(batch_size, -1, host.num_k_heads, host.head_k_dim)
+    key = key.reshape(batch_size, -1, host.num_k_heads, host.head_k_dim)
+    value = value.reshape(batch_size, -1, host.num_v_heads, host.head_v_dim)
+    beta = b.sigmoid()
+    g = -host.A_log.float().exp() * F.softplus(a.float() + host.dt_bias)
+    if not use_precomputed_states:
+        # --- gated delta rule prefill (gated_delta_rule_prefill_split4) ---
+        qn, kn, vv = (t.contiguous() for t in (query, key, value))
+        gate = g.exp().float().contiguous()  # [B, S, VH]
+        betaf = beta.float().contiguous()
+        if has_prev and keyed:
+            # continuation chunk: start from the cached state (h is in/out)
+            h = _state(layer.recurrent_states).to(torch.float32).clone().contiguous()
+        else:
+            h = torch.zeros(
+                batch_size, host.num_v_heads, host.head_k_dim, host.head_v_dim,
+                device=hidden_states.device, dtype=torch.float32,
+            )
+        core_attn_out = torch.empty(
+            batch_size, seq_len, host.num_v_heads, host.head_v_dim,
+            device=hidden_states.device, dtype=torch.bfloat16,
+        )
+        prefill = ops.gdn_prefill if os.environ.get("ATLAS_GDN_FLA_GFX") == "0" else ops.gdn_prefill_fla
+        prefill(h, qn, kn, vv, gate, betaf, core_attn_out)
+        if cache_params is not None:
+            cache_params.update_recurrent_state(h, host.layer_idx)
+    else:
+        # --- single-token recurrence, fp32 end to end (gated_delta_rule_decode_f32) ---
+        qn = query[:, 0].to(torch.float32).contiguous()  # [B, QH, KD]
+        kn = key[:, 0].to(torch.float32).contiguous()
+        vv = value[:, 0].to(torch.float32).contiguous()  # [B, VH, VD]
+        gate = g[:, 0].exp().float().contiguous()  # [B, VH]
+        betaf = beta[:, 0].float().contiguous()
+        h = _state(layer.recurrent_states).to(torch.float32).contiguous()
+        out_t = torch.empty(
+            batch_size, host.num_v_heads, host.head_v_dim,
+            device=hidden_states.device, dtype=torch.float32,
+        )
+        ops.gdn_decode_f32(h, qn, kn, vv, gate, betaf, out_t)
+        cache_params.update_recurrent_state(h, host.layer_idx)
+        # fp32 into the gated norm, as Atlas feeds gated_rms_norm_f32
+        core_attn_out = out_t.unsqueeze(1)  # [B, 1, VH, VD]
+    # --- gated RMSNorm + output projection (reused from host) ---
+    z_shape_og = z.shape
+    core_attn_out = core_attn_out.reshape(-1, core_attn_out.shape[-1])
+    z = z.reshape(-1, z.shape[-1])
+    core_attn_out = host.norm(core_attn_out, z)
+    core_attn_out = core_attn_out.reshape(z_shape_og)
+    core_attn_out = core_attn_out.reshape(core_attn_out.shape[0], core_attn_out.shape[1], -1)
+    return host.out_proj(core_attn_out.to(hidden_states.dtype))
+class GatedDeltaNet(nn.Module):
+    """Drop-in for ``Qwen3NextGatedDeltaNet.forward`` (Qwen3-Next-80B).
+    Fused QKVZ / BA projections, split via the host's
+    ``fix_query_key_value_ordering``.
+    """
+    # Pure recurrent/conv kernels: no autograd, not torch.compile-traceable.
+    has_backward: bool = False
+    can_torch_compile: bool = False
+    def forward(self, hidden_states, cache_params=None, attention_mask=None, **kwargs):
+        hidden_states = _apply_mask_to_padding_states(hidden_states, attention_mask)
+        projected_states_qkvz = self.in_proj_qkvz(hidden_states)
+        projected_states_ba = self.in_proj_ba(hidden_states)
+        query, key, value, z, b, a = self.fix_query_key_value_ordering(
+            projected_states_qkvz, projected_states_ba
+        )
+        query, key, value = (x.reshape(x.shape[0], x.shape[1], -1) for x in (query, key, value))
+        mixed_qkv = torch.cat((query, key, value), dim=-1).transpose(1, 2)  # [B, conv_dim, S]
+        return _gdn_run(self, hidden_states, cache_params, mixed_qkv, z, b, a)
+class Qwen3_5GatedDeltaNet(nn.Module):
+    """Drop-in for the Qwen3.5/3.6 GDN layer.
+    Targets both ``Qwen3_5GatedDeltaNet`` (Qwen3.6-27B dense) and
+    ``Qwen3_5MoeGatedDeltaNet`` (Qwen3.6-35B-A3B); their GDN cores are identical.
+    Already-split projections: ``in_proj_qkv`` (Q|K|V), ``in_proj_z`` (gate),
+    ``in_proj_b`` (beta), ``in_proj_a`` (decay). No ordering fixup.
+    """
+    has_backward: bool = False
+    can_torch_compile: bool = False
+    def forward(self, hidden_states, cache_params=None, attention_mask=None, **kwargs):
+        hidden_states = _apply_mask_to_padding_states(hidden_states, attention_mask)
+        batch_size, seq_len, _ = hidden_states.shape
+        mixed_qkv = self.in_proj_qkv(hidden_states).transpose(1, 2)  # [B, conv_dim, S]
+        z = self.in_proj_z(hidden_states).reshape(batch_size, seq_len, -1, self.head_v_dim)
+        b = self.in_proj_b(hidden_states)  # [B, S, num_v_heads]
+        a = self.in_proj_a(hidden_states)  # [B, S, num_v_heads]
+        return _gdn_run(self, hidden_states, cache_params, mixed_qkv, z, b, a)
+# Qwen3.6-35B-A3B (MoE) shares the dense layer's GDN core verbatim. Expose the
+# host class name so a single LayerRepository entry resolves for either model.
+Qwen3_5MoeGatedDeltaNet = Qwen3_5GatedDeltaNet

build/torch214-cxx11-rocm72-x86_64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,35 @@

+{
+  "name": "gdn",
+  "id": "_gdn_rocm_8d6df22",
+  "version": 1,
+  "kernels-minver": "0.14.0",
+  "license": "AGPL-3.0-only",
+  "python-depends": [],
+  "kernel-depends": [],
+  "backend": {
+    "type": "rocm",
+    "archs": [
+      "gfx1151"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "LiP481kYPHgVyWn/hX7QlwghehLrO6jOKeUkE8w0HNU=",
+      "_gdn_rocm_8d6df22.abi3.so": "d0voo/nGgP2rX5NqRJTZhRsweY9nB7hLBHiZKG/QT2Y=",
+      "_ops.py": "8RvxjaM9m2Nc3SoC6Nmkjrj2ERAl+o8Oi1pIpPO1bh0=",
+      "layers.py": "y5VWO3DlzuCO1r7wuJoYjbMkZImfPQXqON5MouNWnWQ="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.18.0-dev0",
+      "commit": "f13568d6c568d18b2da481edf66b83714300e441",
+      "dirty": false
+    },
+    "kernel": {
+      "commit": "8d6df2295686c12b0fa37fdd71eba3f65273fae0",
+      "dirty": false
+    }
+  }
+}