chore: import upstream snapshot with attribution

python/sglang/jit_kernel/triton/gdn_fused_proj.py

@@ -0,0 +1,407 @@
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+# =============================================================================
+# Fused kernel — reads INTERLEAVED input format
+# Used by Qwen3-Next whose checkpoint stores fused in_proj_qkvz weights
+# in per-head-group interleaved layout:
+#   [g0_q, g0_k, g0_v, g0_z, g1_q, g1_k, g1_v, g1_z, ...]
+# =============================================================================
+
+
+@triton.jit
+def fused_qkvzba_split_reshape_cat_kernel(
+    mixed_qkv,
+    z,
+    b,
+    a,
+    mixed_qkvz,
+    mixed_ba,
+    NUM_HEADS_QK: tl.constexpr,
+    NUM_HEADS_V: tl.constexpr,
+    HEAD_QK: tl.constexpr,
+    HEAD_V: tl.constexpr,
+):
+    i_bs, i_qk = tl.program_id(0), tl.program_id(1)
+    QKVZ_DIM_T: tl.constexpr = HEAD_QK * 2 + NUM_HEADS_V // NUM_HEADS_QK * HEAD_V * 2
+    BA_DIM_T: tl.constexpr = NUM_HEADS_V // NUM_HEADS_QK * 2
+    QKV_DIM_T: tl.constexpr = HEAD_QK * 2 + NUM_HEADS_V // NUM_HEADS_QK * HEAD_V
+    q_end: tl.constexpr = HEAD_QK
+    blk_q_ptr = (
+        mixed_qkvz
+        + i_bs * NUM_HEADS_QK * QKVZ_DIM_T
+        + i_qk * QKVZ_DIM_T
+        + tl.arange(0, q_end)
+    )
+    k_end: tl.constexpr = q_end + HEAD_QK
+    blk_k_ptr = (
+        mixed_qkvz
+        + i_bs * NUM_HEADS_QK * QKVZ_DIM_T
+        + i_qk * QKVZ_DIM_T
+        + tl.arange(q_end, k_end)
+    )
+    v_end: tl.constexpr = k_end + NUM_HEADS_V // NUM_HEADS_QK * HEAD_V
+    blk_v_ptr = (
+        mixed_qkvz
+        + i_bs * NUM_HEADS_QK * QKVZ_DIM_T
+        + i_qk * QKVZ_DIM_T
+        + tl.arange(k_end, v_end)
+    )
+    z_end: tl.constexpr = v_end + NUM_HEADS_V // NUM_HEADS_QK * HEAD_V
+    blk_z_ptr = (
+        mixed_qkvz
+        + i_bs * NUM_HEADS_QK * QKVZ_DIM_T
+        + i_qk * QKVZ_DIM_T
+        + tl.arange(v_end, z_end)
+    )
+    blk_q_st_ptr = (
+        mixed_qkv
+        + i_bs * NUM_HEADS_QK * QKV_DIM_T
+        + i_qk * HEAD_QK
+        + tl.arange(0, HEAD_QK)
+    )
+    blk_k_st_ptr = (
+        mixed_qkv
+        + i_bs * NUM_HEADS_QK * QKV_DIM_T
+        + NUM_HEADS_QK * HEAD_QK
+        + i_qk * HEAD_QK
+        + tl.arange(0, HEAD_QK)
+    )
+    blk_v_st_ptr = (
+        mixed_qkv
+        + i_bs * NUM_HEADS_QK * QKV_DIM_T
+        + NUM_HEADS_QK * HEAD_QK * 2
+        + i_qk * HEAD_V * NUM_HEADS_V // NUM_HEADS_QK
+        + tl.arange(0, HEAD_V * NUM_HEADS_V // NUM_HEADS_QK)
+    )
+    blk_z_st_ptr = (
+        z
+        + i_bs * NUM_HEADS_V * HEAD_V
+        + i_qk * HEAD_V * NUM_HEADS_V // NUM_HEADS_QK
+        + tl.arange(0, HEAD_V * NUM_HEADS_V // NUM_HEADS_QK)
+    )
+    tl.store(blk_q_st_ptr, tl.load(blk_q_ptr))
+    tl.store(blk_k_st_ptr, tl.load(blk_k_ptr))
+    tl.store(blk_v_st_ptr, tl.load(blk_v_ptr))
+    tl.store(blk_z_st_ptr, tl.load(blk_z_ptr))
+    b_end: tl.constexpr = NUM_HEADS_V // NUM_HEADS_QK
+    a_end: tl.constexpr = b_end + NUM_HEADS_V // NUM_HEADS_QK
+    for i in tl.static_range(b_end):
+        blk_b_ptr = mixed_ba + i_bs * NUM_HEADS_QK * BA_DIM_T + i_qk * BA_DIM_T + i
+        blk_b_st_ptr = b + i_bs * NUM_HEADS_V + i_qk * NUM_HEADS_V // NUM_HEADS_QK + i
+        tl.store(blk_b_st_ptr, tl.load(blk_b_ptr))
+    for i in tl.static_range(b_end, a_end):
+        blk_a_ptr = mixed_ba + i_bs * NUM_HEADS_QK * BA_DIM_T + i_qk * BA_DIM_T + i
+        blk_a_st_ptr = (
+            a + i_bs * NUM_HEADS_V + i_qk * NUM_HEADS_V // NUM_HEADS_QK + (i - b_end)
+        )
+        tl.store(blk_a_st_ptr, tl.load(blk_a_ptr))
+
+
+def fused_qkvzba_split_reshape_cat(
+    mixed_qkvz,
+    mixed_ba,
+    num_heads_qk,
+    num_heads_v,
+    head_qk,
+    head_v,
+):
+    batch, seq_len = mixed_qkvz.shape[0], 1
+    qkv_dim_t = num_heads_qk * head_qk * 2 + num_heads_v * head_v
+    mixed_qkv = torch.empty(
+        [batch * seq_len, qkv_dim_t],
+        dtype=mixed_qkvz.dtype,
+        device=mixed_qkvz.device,
+    )
+    z = torch.empty(
+        [batch * seq_len, num_heads_v, head_v],
+        dtype=mixed_qkvz.dtype,
+        device=mixed_qkvz.device,
+    )
+    b = torch.empty(
+        [batch * seq_len, num_heads_v],
+        dtype=mixed_ba.dtype,
+        device=mixed_ba.device,
+    )
+    a = torch.empty_like(b)
+    grid = (batch * seq_len, num_heads_qk)
+    fused_qkvzba_split_reshape_cat_kernel[grid](
+        mixed_qkv,
+        z,
+        b,
+        a,
+        mixed_qkvz,
+        mixed_ba,
+        num_heads_qk,
+        num_heads_v,
+        head_qk,
+        head_v,
+        num_warps=1,
+        num_stages=3,
+    )
+    return mixed_qkv, z, b, a
+
+
+# =============================================================================
+# Fused kernel — reads CONTIGUOUS input format
+# Used by Qwen3.5 whose checkpoint stores in_proj_qkv and in_proj_z separately.
+# After MergedColumnParallelLinear loads them, the matmul output is contiguous:
+#   mixed_qkvz: [all_q | all_k | all_v | all_z]
+#   mixed_ba:   [all_b | all_a]
+#
+# Output format is identical to the interleaved kernel (same downstream consumer).
+# =============================================================================
+
+
+@triton.jit
+def fused_qkvzba_split_reshape_cat_contiguous_kernel(
+    mixed_qkv,
+    z,
+    b,
+    a,
+    mixed_qkvz,
+    mixed_ba,
+    NUM_HEADS_QK: tl.constexpr,
+    NUM_HEADS_V: tl.constexpr,
+    HEAD_QK: tl.constexpr,
+    HEAD_V: tl.constexpr,
+):
+    i_bs, i_qk = tl.program_id(0), tl.program_id(1)
+
+    V_PER_GROUP: tl.constexpr = NUM_HEADS_V // NUM_HEADS_QK
+
+    # ── Input dimensions (contiguous layout) ──
+    TOTAL_Q: tl.constexpr = NUM_HEADS_QK * HEAD_QK
+    TOTAL_K: tl.constexpr = NUM_HEADS_QK * HEAD_QK
+    TOTAL_V: tl.constexpr = NUM_HEADS_V * HEAD_V
+    TOTAL_QKVZ: tl.constexpr = TOTAL_Q + TOTAL_K + TOTAL_V + TOTAL_V
+    TOTAL_BA: tl.constexpr = NUM_HEADS_V * 2
+
+    # ── Output dimensions ──
+    QKV_DIM_T: tl.constexpr = TOTAL_Q + TOTAL_K + TOTAL_V
+
+    # ── Read from contiguous input ──
+    # q for head group i_qk: in the all_q region, offset i_qk * HEAD_QK
+    blk_q_ptr = mixed_qkvz + i_bs * TOTAL_QKVZ + i_qk * HEAD_QK + tl.arange(0, HEAD_QK)
+    # k for head group i_qk: in the all_k region
+    blk_k_ptr = (
+        mixed_qkvz
+        + i_bs * TOTAL_QKVZ
+        + TOTAL_Q
+        + i_qk * HEAD_QK
+        + tl.arange(0, HEAD_QK)
+    )
+    # v for head group i_qk: in the all_v region
+    blk_v_ptr = (
+        mixed_qkvz
+        + i_bs * TOTAL_QKVZ
+        + TOTAL_Q
+        + TOTAL_K
+        + i_qk * V_PER_GROUP * HEAD_V
+        + tl.arange(0, V_PER_GROUP * HEAD_V)
+    )
+    # z for head group i_qk: in the all_z region
+    blk_z_ptr = (
+        mixed_qkvz
+        + i_bs * TOTAL_QKVZ
+        + TOTAL_Q
+        + TOTAL_K
+        + TOTAL_V
+        + i_qk * V_PER_GROUP * HEAD_V
+        + tl.arange(0, V_PER_GROUP * HEAD_V)
+    )
+
+    # ── Write to output (identical layout to the interleaved kernel) ──
+    blk_q_st_ptr = mixed_qkv + i_bs * QKV_DIM_T + i_qk * HEAD_QK + tl.arange(0, HEAD_QK)
+    blk_k_st_ptr = (
+        mixed_qkv
+        + i_bs * QKV_DIM_T
+        + NUM_HEADS_QK * HEAD_QK
+        + i_qk * HEAD_QK
+        + tl.arange(0, HEAD_QK)
+    )
+    blk_v_st_ptr = (
+        mixed_qkv
+        + i_bs * QKV_DIM_T
+        + NUM_HEADS_QK * HEAD_QK * 2
+        + i_qk * V_PER_GROUP * HEAD_V
+        + tl.arange(0, V_PER_GROUP * HEAD_V)
+    )
+    blk_z_st_ptr = (
+        z
+        + i_bs * NUM_HEADS_V * HEAD_V
+        + i_qk * V_PER_GROUP * HEAD_V
+        + tl.arange(0, V_PER_GROUP * HEAD_V)
+    )
+
+    tl.store(blk_q_st_ptr, tl.load(blk_q_ptr))
+    tl.store(blk_k_st_ptr, tl.load(blk_k_ptr))
+    tl.store(blk_v_st_ptr, tl.load(blk_v_ptr))
+    tl.store(blk_z_st_ptr, tl.load(blk_z_ptr))
+
+    # ── b and a from contiguous [all_b | all_a] ──
+    for i in tl.static_range(V_PER_GROUP):
+        blk_b_ptr = mixed_ba + i_bs * TOTAL_BA + i_qk * V_PER_GROUP + i
+        blk_b_st_ptr = b + i_bs * NUM_HEADS_V + i_qk * V_PER_GROUP + i
+        tl.store(blk_b_st_ptr, tl.load(blk_b_ptr))
+
+    for i in tl.static_range(V_PER_GROUP):
+        blk_a_ptr = mixed_ba + i_bs * TOTAL_BA + NUM_HEADS_V + i_qk * V_PER_GROUP + i
+        blk_a_st_ptr = a + i_bs * NUM_HEADS_V + i_qk * V_PER_GROUP + i
+        tl.store(blk_a_st_ptr, tl.load(blk_a_ptr))
+
+
+def fused_qkvzba_split_reshape_cat_contiguous(
+    mixed_qkvz,
+    mixed_ba,
+    num_heads_qk,
+    num_heads_v,
+    head_qk,
+    head_v,
+):
+    """Fused split/reshape/cat for CONTIGUOUS input format (Qwen3.5).
+
+    Input layout:
+        mixed_qkvz: [all_q | all_k | all_v | all_z]
+        mixed_ba:   [all_b | all_a]
+
+    Output layout (same as fused_qkvzba_split_reshape_cat):
+        mixed_qkv: [all_q | all_k | all_v]  (z stripped)
+        z: [num_v_heads, head_v]
+        b: [num_v_heads]
+        a: [num_v_heads]
+    """
+    batch, seq_len = mixed_qkvz.shape[0], 1
+    qkv_dim_t = num_heads_qk * head_qk * 2 + num_heads_v * head_v
+    mixed_qkv = torch.empty(
+        [batch * seq_len, qkv_dim_t],
+        dtype=mixed_qkvz.dtype,
+        device=mixed_qkvz.device,
+    )
+    z = torch.empty(
+        [batch * seq_len, num_heads_v, head_v],
+        dtype=mixed_qkvz.dtype,
+        device=mixed_qkvz.device,
+    )
+    b = torch.empty(
+        [batch * seq_len, num_heads_v],
+        dtype=mixed_ba.dtype,
+        device=mixed_ba.device,
+    )
+    a = torch.empty_like(b)
+    grid = (batch * seq_len, num_heads_qk)
+    fused_qkvzba_split_reshape_cat_contiguous_kernel[grid](
+        mixed_qkv,
+        z,
+        b,
+        a,
+        mixed_qkvz,
+        mixed_ba,
+        num_heads_qk,
+        num_heads_v,
+        head_qk,
+        head_v,
+        num_warps=1,
+        num_stages=3,
+    )
+    return mixed_qkv, z, b, a
+
+
+@triton.jit
+def fused_qkv_split_gdn_prefill_kernel(
+    q,
+    k,
+    v,
+    mixed_qkv,
+    MIXED_QKV_STRIDE_T: tl.constexpr,
+    MIXED_QKV_STRIDE_D: tl.constexpr,
+    NUM_Q_HEADS: tl.constexpr,
+    NUM_K_HEADS: tl.constexpr,
+    NUM_V_HEADS: tl.constexpr,
+    HEAD_Q: tl.constexpr,
+    HEAD_K: tl.constexpr,
+    HEAD_V: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    i_t = tl.program_id(0)
+    offsets = tl.arange(0, BLOCK_SIZE)
+
+    q_dim: tl.constexpr = NUM_Q_HEADS * HEAD_Q
+    k_dim: tl.constexpr = NUM_K_HEADS * HEAD_K
+    v_dim: tl.constexpr = NUM_V_HEADS * HEAD_V
+    qk_dim: tl.constexpr = q_dim + k_dim
+    qkv_dim: tl.constexpr = qk_dim + v_dim
+
+    mask = offsets < qkv_dim
+    values = tl.load(
+        mixed_qkv + i_t * MIXED_QKV_STRIDE_T + offsets * MIXED_QKV_STRIDE_D,
+        mask=mask,
+    )
+
+    q_mask = offsets < q_dim
+    tl.store(q + i_t * q_dim + offsets, values, mask=q_mask)
+
+    k_offsets = offsets - q_dim
+    k_mask = (offsets >= q_dim) & (offsets < qk_dim)
+    tl.store(k + i_t * k_dim + k_offsets, values, mask=k_mask)
+
+    v_offsets = offsets - qk_dim
+    v_mask = (offsets >= qk_dim) & (offsets < qkv_dim)
+    tl.store(v + i_t * v_dim + v_offsets, values, mask=v_mask)
+
+
+def fused_qkv_split_gdn_prefill(
+    mixed_qkv: torch.Tensor,
+    num_q_heads: int,
+    num_k_heads: int,
+    num_v_heads: int,
+    head_q: int,
+    head_k: int,
+    head_v: int,
+):
+    """Split packed post-conv GDN QKV into contiguous FLA prefill tensors.
+
+    `mixed_qkv` is laid out per token as `[all_q | all_k | all_v]`. The FLA
+    chunk kernels consume separate contiguous `[1, T, H, D]` tensors, so this
+    fused split replaces three independent `aten::copy_` kernels from the
+    generic FLA input guard. `mixed_qkv` may be a strided `[T, qkv_dim]` view.
+    """
+    seq_len = mixed_qkv.shape[0]
+    q = torch.empty(
+        (1, seq_len, num_q_heads, head_q),
+        dtype=mixed_qkv.dtype,
+        device=mixed_qkv.device,
+    )
+    k = torch.empty(
+        (1, seq_len, num_k_heads, head_k),
+        dtype=mixed_qkv.dtype,
+        device=mixed_qkv.device,
+    )
+    v = torch.empty(
+        (1, seq_len, num_v_heads, head_v),
+        dtype=mixed_qkv.dtype,
+        device=mixed_qkv.device,
+    )
+
+    qkv_dim = num_q_heads * head_q + num_k_heads * head_k + num_v_heads * head_v
+    fused_qkv_split_gdn_prefill_kernel[(seq_len,)](
+        q,
+        k,
+        v,
+        mixed_qkv,
+        mixed_qkv.stride(0),
+        mixed_qkv.stride(1),
+        num_q_heads,
+        num_k_heads,
+        num_v_heads,
+        head_q,
+        head_k,
+        head_v,
+        BLOCK_SIZE=triton.next_power_of_2(qkv_dim),
+        num_warps=8,
+        num_stages=3,
+    )
+    return q, k, v

python/sglang/jit_kernel/triton/hash_topk.py

@@ -0,0 +1,99 @@
+"""HIP fallback for ``hash_topk``: ``csrc/deepseek_v4/hash_topk.cuh`` uses
+CUDA-only primitives, so on ROCm we dispatch to this Triton implementation.
+"""
+
+from __future__ import annotations
+
+from typing import Tuple
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _hash_topk_triton_kernel(
+    router_logits_ptr,
+    input_ids_ptr,
+    tid2eid_ptr,
+    topk_weights_ptr,
+    topk_ids_ptr,
+    num_routed_experts: tl.constexpr,
+    topk_routed: tl.constexpr,
+    topk_fused: tl.constexpr,
+    routed_scaling_factor,
+    BLOCK_K: tl.constexpr,
+):
+    token_pos = tl.program_id(0)
+    token_id = tl.load(input_ids_ptr + token_pos).to(tl.int64)
+
+    k_off = tl.arange(0, BLOCK_K)
+    routed_mask = k_off < topk_routed
+    fused_mask = k_off < topk_fused
+    is_shared = k_off >= topk_routed
+
+    expert_id = tl.load(
+        tid2eid_ptr + token_id * topk_routed + k_off,
+        mask=routed_mask,
+        other=0,
+    ).to(tl.int32)
+    logit = tl.load(
+        router_logits_ptr + token_pos * num_routed_experts + expert_id,
+        mask=routed_mask,
+        other=0.0,
+    ).to(tl.float32)
+
+    softplus = tl.maximum(logit, 0.0) + tl.log(1.0 + tl.exp(-tl.abs(logit)))
+    weight = tl.sqrt(softplus)
+    weight = tl.where(routed_mask, weight, 0.0)
+    routed_sum = tl.sum(weight, axis=0)
+
+    shared_weight = 1.0 / routed_scaling_factor
+    final_weight = tl.where(is_shared, shared_weight, weight / routed_sum)
+    shared_id = num_routed_experts + (k_off - topk_routed)
+    final_id = tl.where(is_shared, shared_id, expert_id).to(tl.int32)
+
+    out_off = token_pos * topk_fused + k_off
+    tl.store(topk_weights_ptr + out_off, final_weight, mask=fused_mask)
+    tl.store(topk_ids_ptr + out_off, final_id, mask=fused_mask)
+
+
+def hash_topk_triton(
+    router_logits: torch.Tensor,
+    input_ids: torch.Tensor,
+    tid2eid: torch.Tensor,
+    num_fused_shared_experts: int,
+    routed_scaling_factor: float,
+    scoring_func: str,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    assert scoring_func == "sqrtsoftplus"
+
+    num_tokens = router_logits.size(0)
+    num_routed_experts = router_logits.size(1)
+    topk_routed = tid2eid.size(1)
+    topk_fused = topk_routed + num_fused_shared_experts
+
+    topk_weights = torch.empty(
+        (num_tokens, topk_fused), dtype=torch.float32, device=router_logits.device
+    )
+    topk_ids = torch.empty(
+        (num_tokens, topk_fused), dtype=torch.int32, device=router_logits.device
+    )
+    if num_tokens == 0:
+        return topk_weights, topk_ids
+
+    block_k = max(triton.next_power_of_2(topk_fused), 1)
+    _hash_topk_triton_kernel[(num_tokens,)](
+        router_logits,
+        input_ids,
+        tid2eid,
+        topk_weights,
+        topk_ids,
+        num_routed_experts=num_routed_experts,
+        topk_routed=topk_routed,
+        topk_fused=topk_fused,
+        routed_scaling_factor=float(routed_scaling_factor),
+        BLOCK_K=block_k,
+        num_warps=1,
+    )
+    return topk_weights, topk_ids

python/sglang/jit_kernel/triton/sigmoid_gate_mul.py

@@ -0,0 +1,87 @@
+"""Fused sigmoid-gate-multiply Triton kernels.
+
+Two variants:
+- ``sigmoid_gate_mul``: element-wise ``x * sigmoid(gate)`` when x and gate
+  have identical shapes.
+- ``sigmoid_gate_mul_broadcast``: broadcast ``x * sigmoid(gate)`` when gate
+  is ``(N, 1)`` and x is ``(N, D)``.
+"""
+
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.utils import is_hip
+
+_is_hip = is_hip()
+
+
+@triton.jit
+def _sigmoid_gate_mul_kernel(
+    x_ptr,
+    gate_ptr,
+    out_ptr,
+    n_elements,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < n_elements
+    x = tl.load(x_ptr + offsets, mask=mask).to(tl.float32)
+    g = tl.load(gate_ptr + offsets, mask=mask).to(tl.float32)
+    out = x * tl.sigmoid(g)
+    tl.store(out_ptr + offsets, out.to(x_ptr.dtype.element_ty), mask=mask)
+
+
+def sigmoid_gate_mul(x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor:
+    """Compute ``x * sigmoid(gate)`` in a single fused kernel (same-shape)."""
+    out = torch.empty_like(x)
+    n = x.numel()
+    grid = lambda meta: (triton.cdiv(n, meta["BLOCK_SIZE"]),)
+    _sigmoid_gate_mul_kernel[grid](x, gate, out, n, BLOCK_SIZE=1024)
+    return out
+
+
+@triton.jit
+def _sigmoid_gate_mul_broadcast_kernel(
+    out_ptr,
+    gate_ptr,
+    x_ptr,
+    hidden_dim: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    row = tl.program_id(0)
+    g = tl.load(gate_ptr + row).to(tl.float32)
+    g = tl.sigmoid(g)
+
+    offs = tl.arange(0, BLOCK_SIZE)
+    mask = offs < hidden_dim
+    x = tl.load(x_ptr + row * hidden_dim + offs, mask=mask).to(tl.float32)
+    out = x * g
+    tl.store(
+        out_ptr + row * hidden_dim + offs,
+        out.to(x_ptr.dtype.element_ty),
+        mask=mask,
+    )
+
+
+def sigmoid_gate_mul_broadcast(x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor:
+    """Compute ``x * sigmoid(gate)`` where gate is (N, 1) and x is (N, D)."""
+    bs, hidden_dim = x.shape
+    out = torch.empty_like(x)
+    BLOCK_SIZE = triton.next_power_of_2(hidden_dim)
+    max_warps = 16 if _is_hip else 32
+    num_warps = max(
+        min(triton.next_power_of_2(triton.cdiv(hidden_dim, 8 * 32)), max_warps), 4
+    )
+    _sigmoid_gate_mul_broadcast_kernel[(bs,)](
+        out,
+        gate,
+        x,
+        hidden_dim=hidden_dim,
+        BLOCK_SIZE=BLOCK_SIZE,
+        num_warps=num_warps,
+    )
+    return out