From 8cf73ad39c0c84a5424372eb8f92091e7cc149d4 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 10 Aug 2026 01:06:32 +0000 Subject: [PATCH] feat(SM70): add 1Cat-vLLM FlashQLA fused GDN CUDA kernel for BI-V100 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Source: github.com/1CatAI/1Cat-vLLM (MIT license) flash_qla/ops/gated_delta_rule/chunk/sm70/ Files added: - csrc/gdn_forward.cu (1919 lines) — 4 CUDA kernels for SM70/SM75: gdn_forward, gdn_forward_vlk_varlen, gdn_decode_mixed_qkv_global_state, gdn_decode_mixed_qkv_ddtree_state - fused_fwd.py — Python wrapper, JIT compiles via torch.utils.cpp_extension.load() - naive_gdn.py — fla reference PyTorch implementation for fallback - __init__.py — exports chunk_gated_delta_rule_fwd_sm70 Build: JIT compiled at runtime (TORCH_CUDA_ARCH_LIST=7.0;7.5 -O3) Deploy: patch_ops.sh copies flash_qla_sm70/ to vllm models dir qwen3_5.py updated to try import flash_qla_sm70 before PyTorch fallback --- qwen3_6_scripts/flash_qla_sm70/__init__.py | 14 + .../flash_qla_sm70/csrc/gdn_forward.cu | 1919 +++++++++++++++++ qwen3_6_scripts/flash_qla_sm70/fused_fwd.py | 490 +++++ qwen3_6_scripts/flash_qla_sm70/naive_gdn.py | 161 ++ qwen3_6_scripts/patch_ops.sh | 22 +- qwen3_6_scripts/qwen3_5.py | 20 +- 6 files changed, 2621 insertions(+), 5 deletions(-) create mode 100644 qwen3_6_scripts/flash_qla_sm70/__init__.py create mode 100644 qwen3_6_scripts/flash_qla_sm70/csrc/gdn_forward.cu create mode 100644 qwen3_6_scripts/flash_qla_sm70/fused_fwd.py create mode 100644 qwen3_6_scripts/flash_qla_sm70/naive_gdn.py diff --git a/qwen3_6_scripts/flash_qla_sm70/__init__.py b/qwen3_6_scripts/flash_qla_sm70/__init__.py new file mode 100644 index 00000000..7deafb2f --- /dev/null +++ b/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", +] diff --git a/qwen3_6_scripts/flash_qla_sm70/csrc/gdn_forward.cu b/qwen3_6_scripts/flash_qla_sm70/csrc/gdn_forward.cu new file mode 100644 index 00000000..f111d8b5 --- /dev/null +++ b/qwen3_6_scripts/flash_qla_sm70/csrc/gdn_forward.cu @@ -0,0 +1,1919 @@ +#include + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +namespace { + +void check_cuda(cudaError_t status, const char* context) { + if (status != cudaSuccess) { + throw std::runtime_error(std::string(context) + ": " + + cudaGetErrorString(status)); + } +} + +template +__device__ __forceinline__ float load_as_float(const T* ptr, int64_t index) { + return static_cast(ptr[index]); +} + +template <> +__device__ __forceinline__ float load_as_float(const at::Half* ptr, + int64_t index) { + return __half2float(reinterpret_cast(ptr)[index]); +} + +template <> +__device__ __forceinline__ float load_as_float( + const at::BFloat16* ptr, int64_t index) { + return __bfloat162float(reinterpret_cast(ptr)[index]); +} + +template +__device__ __forceinline__ void store_from_float(T* ptr, + int64_t index, + float value) { + ptr[index] = static_cast(value); +} + +template <> +__device__ __forceinline__ void store_from_float(at::Half* ptr, + int64_t index, + float value) { + reinterpret_cast<__half*>(ptr)[index] = __float2half(value); +} + +template <> +__device__ __forceinline__ void store_from_float( + at::BFloat16* ptr, + int64_t index, + float value) { + reinterpret_cast<__nv_bfloat16*>(ptr)[index] = __float2bfloat16(value); +} + +__device__ __forceinline__ float subgroup_sum(float value, int width) { + constexpr unsigned mask = 0xffffffffU; + for (int offset = width / 2; offset > 0; offset >>= 1) { + value += __shfl_down_sync(mask, value, offset, width); + } + return value; +} + +__device__ __forceinline__ float subgroup_broadcast(float value, int width) { + return __shfl_sync(0xffffffffU, value, 0, width); +} + +template +__global__ void gdn_forward_kernel(const scalar_t* __restrict__ q, + const scalar_t* __restrict__ k, + const scalar_t* __restrict__ v, + const gate_t* __restrict__ gate, + const beta_t* __restrict__ beta, + const state_t* __restrict__ initial_state, + scalar_t* __restrict__ output, + float* __restrict__ final_state, + int batch, + int tokens, + int q_heads, + int v_heads, + float scale, + bool gate_is_exp) { + static_assert(K % WIDTH == 0); + constexpr int subgroups_per_warp = 32 / WIDTH; + constexpr int rows_per_lane = K / WIDTH; + + const int hv = blockIdx.x; + const int b = blockIdx.y; + const int subgroup = threadIdx.x / WIDTH; + const int lane = threadIdx.x % WIDTH; + const int group_base = + (blockIdx.z * blockDim.y + threadIdx.y) * subgroups_per_warp + subgroup; + const int col_base = group_base * COLS; + const int hq = hv / (v_heads / q_heads); + + float state_shard[COLS][rows_per_lane]; + +#pragma unroll + for (int c = 0; c < COLS; ++c) { + const int col = col_base + c; +#pragma unroll + for (int r = 0; r < rows_per_lane; ++r) { + const int row = r * WIDTH + lane; + float value = 0.0F; + if (col < V) { + const int64_t state_index = + (((static_cast(b) * v_heads + hv) * K + row) * V) + col; + value = initial_state == nullptr ? 0.0F + : load_as_float(initial_state, state_index); + } + state_shard[c][r] = value; + } + } + + for (int t = 0; t < tokens; ++t) { + const int64_t gate_index = + ((static_cast(b) * tokens + t) * v_heads + hv); + float gate_value = 0.0F; + float beta_value = 0.0F; + if (threadIdx.x == 0) { + const float gate_raw = load_as_float(gate, gate_index); + gate_value = GateIsExp ? gate_raw : __expf(gate_raw); + beta_value = load_as_float(beta, gate_index); + } + gate_value = __shfl_sync(0xffffffffU, gate_value, 0); + beta_value = __shfl_sync(0xffffffffU, beta_value, 0); + + float k_reg[rows_per_lane]; + float q_reg[rows_per_lane]; + float kv_partial[COLS]; +#pragma unroll + for (int c = 0; c < COLS; ++c) { + kv_partial[c] = 0.0F; + } + +#pragma unroll + for (int r = 0; r < rows_per_lane; ++r) { + const int row = r * WIDTH + lane; + const int64_t qk_index = + (((static_cast(b) * tokens + t) * q_heads + hq) * K) + row; + const float q_value = load_as_float(q, qk_index); + const float k_value = load_as_float(k, qk_index); + q_reg[r] = q_value; + k_reg[r] = k_value; +#pragma unroll + for (int c = 0; c < COLS; ++c) { + kv_partial[c] += state_shard[c][r] * k_value; + } + } + + float delta[COLS]; +#pragma unroll + for (int c = 0; c < COLS; ++c) { + const int col = col_base + c; + const float kv_col = subgroup_sum(kv_partial[c], WIDTH); + float delta_value = 0.0F; + if (lane == 0 && col < V) { + const int64_t v_index = + (((static_cast(b) * tokens + t) * v_heads + hv) * V) + col; + delta_value = + (load_as_float(v, v_index) - gate_value * kv_col) * beta_value; + } + delta[c] = subgroup_broadcast(delta_value, WIDTH); + } + + float attn_partial[COLS]; +#pragma unroll + for (int c = 0; c < COLS; ++c) { + attn_partial[c] = 0.0F; + } + +#pragma unroll + for (int r = 0; r < rows_per_lane; ++r) { +#pragma unroll + for (int c = 0; c < COLS; ++c) { + const float new_state = + fmaf(k_reg[r], delta[c], gate_value * state_shard[c][r]); + state_shard[c][r] = new_state; + attn_partial[c] += new_state * q_reg[r]; + } + } + +#pragma unroll + for (int c = 0; c < COLS; ++c) { + attn_partial[c] = subgroup_sum(attn_partial[c], WIDTH); + } + + if (lane == 0) { + const int64_t out_base = + (((static_cast(b) * tokens + t) * v_heads + hv) * V); +#pragma unroll + for (int c = 0; c < COLS; ++c) { + const int col = col_base + c; + if (col < V) { + store_from_float(output, out_base + col, attn_partial[c] * scale); + } + } + } + } + + if (final_state != nullptr) { +#pragma unroll + for (int c = 0; c < COLS; ++c) { + const int col = col_base + c; +#pragma unroll + for (int r = 0; r < rows_per_lane; ++r) { + const int row = r * WIDTH + lane; + if (col < V) { + const int64_t state_index = + (((static_cast(b) * v_heads + hv) * K + row) * V) + col; + final_state[state_index] = state_shard[c][r]; + } + } + } + } +} + +template +__global__ void gdn_forward_vlk_varlen_kernel( + const scalar_t* __restrict__ q, + const scalar_t* __restrict__ k, + const scalar_t* __restrict__ v, + const gate_t* __restrict__ gate, + const beta_t* __restrict__ beta, + const state_t* __restrict__ initial_state, + const int32_t* __restrict__ cu_seqlens, + scalar_t* __restrict__ output, + float* __restrict__ final_state, + int q_heads, + int v_heads, + float scale, + bool gate_is_exp) { + static_assert(K % WIDTH == 0); + constexpr int subgroups_per_warp = 32 / WIDTH; + constexpr int rows_per_lane = K / WIDTH; + + const int hv = blockIdx.x; + const int n = blockIdx.y; + const int subgroup = threadIdx.x / WIDTH; + const int lane = threadIdx.x % WIDTH; + const int group_base = + (blockIdx.z * blockDim.y + threadIdx.y) * subgroups_per_warp + subgroup; + const int col_base = group_base * COLS; + const int hq = hv / (v_heads / q_heads); + const int seq_start = cu_seqlens[n]; + const int seq_end = cu_seqlens[n + 1]; + + float state_shard[COLS][rows_per_lane]; + +#pragma unroll + for (int c = 0; c < COLS; ++c) { + const int col = col_base + c; +#pragma unroll + for (int r = 0; r < rows_per_lane; ++r) { + const int row = r * WIDTH + lane; + float value = 0.0F; + if (col < V) { + const int64_t state_index = + (((static_cast(n) * v_heads + hv) * V + col) * K) + row; + value = initial_state == nullptr ? 0.0F + : load_as_float(initial_state, state_index); + } + state_shard[c][r] = value; + } + } + + for (int t = seq_start; t < seq_end; ++t) { + const int64_t gate_index = (static_cast(t) * v_heads + hv); + float gate_value = 0.0F; + float beta_value = 0.0F; + if (threadIdx.x == 0) { + const float gate_raw = load_as_float(gate, gate_index); + gate_value = GateIsExp ? gate_raw : __expf(gate_raw); + beta_value = load_as_float(beta, gate_index); + } + gate_value = __shfl_sync(0xffffffffU, gate_value, 0); + beta_value = __shfl_sync(0xffffffffU, beta_value, 0); + + float k_reg[rows_per_lane]; + float q_reg[rows_per_lane]; + float kv_partial[COLS]; +#pragma unroll + for (int c = 0; c < COLS; ++c) { + kv_partial[c] = 0.0F; + } + +#pragma unroll + for (int r = 0; r < rows_per_lane; ++r) { + const int row = r * WIDTH + lane; + const int64_t qk_index = + ((static_cast(t) * q_heads + hq) * K) + row; + const float q_value = load_as_float(q, qk_index); + const float k_value = load_as_float(k, qk_index); + q_reg[r] = q_value; + k_reg[r] = k_value; +#pragma unroll + for (int c = 0; c < COLS; ++c) { + kv_partial[c] += state_shard[c][r] * k_value; + } + } + + float delta[COLS]; +#pragma unroll + for (int c = 0; c < COLS; ++c) { + const int col = col_base + c; + const float kv_col = subgroup_sum(kv_partial[c], WIDTH); + float delta_value = 0.0F; + if (lane == 0 && col < V) { + const int64_t v_index = + ((static_cast(t) * v_heads + hv) * V) + col; + delta_value = + (load_as_float(v, v_index) - gate_value * kv_col) * beta_value; + } + delta[c] = subgroup_broadcast(delta_value, WIDTH); + } + + float attn_partial[COLS]; +#pragma unroll + for (int c = 0; c < COLS; ++c) { + attn_partial[c] = 0.0F; + } + +#pragma unroll + for (int r = 0; r < rows_per_lane; ++r) { +#pragma unroll + for (int c = 0; c < COLS; ++c) { + const float new_state = + fmaf(k_reg[r], delta[c], gate_value * state_shard[c][r]); + state_shard[c][r] = new_state; + attn_partial[c] += new_state * q_reg[r]; + } + } + +#pragma unroll + for (int c = 0; c < COLS; ++c) { + attn_partial[c] = subgroup_sum(attn_partial[c], WIDTH); + } + + if (lane == 0) { + const int64_t out_base = (static_cast(t) * v_heads + hv) * V; +#pragma unroll + for (int c = 0; c < COLS; ++c) { + const int col = col_base + c; + if (col < V) { + store_from_float(output, out_base + col, attn_partial[c] * scale); + } + } + } + } + + if (final_state != nullptr) { +#pragma unroll + for (int c = 0; c < COLS; ++c) { + const int col = col_base + c; +#pragma unroll + for (int r = 0; r < rows_per_lane; ++r) { + const int row = r * WIDTH + lane; + if (col < V) { + const int64_t state_index = + (((static_cast(n) * v_heads + hv) * V + col) * K) + row; + final_state[state_index] = state_shard[c][r]; + } + } + } + } +} + +template +__global__ void gdn_decode_mixed_qkv_global_state_kernel( + const scalar_t* __restrict__ mixed_qkv, + const scalar_t* __restrict__ a, + const scalar_t* __restrict__ b, + const float* __restrict__ A_log, + const bias_t* __restrict__ dt_bias, + state_t* __restrict__ state, + const int32_t* __restrict__ state_indices, + scalar_t* __restrict__ output, + int tokens, + int slots, + int64_t state_slot_stride, + int q_heads, + int v_heads, + int qkv_stride, + float scale, + bool use_qk_l2norm) { + static_assert(K % WIDTH == 0); + constexpr int subgroups_per_warp = 32 / WIDTH; + constexpr int rows_per_lane = K / WIDTH; + + const int hv = blockIdx.x; + const int t = blockIdx.y; + const int subgroup = threadIdx.x / WIDTH; + const int lane = threadIdx.x % WIDTH; + const int group_base = + (blockIdx.z * blockDim.y + threadIdx.y) * subgroups_per_warp + subgroup; + const int col_base = group_base * COLS; + const int hq = hv / (v_heads / q_heads); + const int32_t slot = state_indices[t]; + + if (slot < 0 || slot >= slots) { + if (lane == 0) { + const int64_t out_base = (static_cast(t) * v_heads + hv) * V; +#pragma unroll + for (int c = 0; c < COLS; ++c) { + const int col = col_base + c; + if (col < V) { + store_from_float(output, out_base + col, 0.0F); + } + } + } + return; + } + + float state_shard[COLS][rows_per_lane]; +#pragma unroll + for (int c = 0; c < COLS; ++c) { + const int col = col_base + c; +#pragma unroll + for (int r = 0; r < rows_per_lane; ++r) { + const int row = r * WIDTH + lane; + float value = 0.0F; + if (col < V) { + const int64_t state_index = + static_cast(slot) * state_slot_stride + + ((static_cast(hv) * V + col) * K) + row; + value = load_as_float(state, state_index); + } + state_shard[c][r] = value; + } + } + + const int64_t mixed_base = static_cast(t) * qkv_stride; + float k_reg[rows_per_lane]; + float q_reg[rows_per_lane]; + float q_norm = 0.0F; + float k_norm = 0.0F; +#pragma unroll + for (int r = 0; r < rows_per_lane; ++r) { + const int row = r * WIDTH + lane; + const int64_t q_index = mixed_base + hq * K + row; + const int64_t k_index = mixed_base + q_heads * K + hq * K + row; + const float q_value = load_as_float(mixed_qkv, q_index); + const float k_value = load_as_float(mixed_qkv, k_index); + q_reg[r] = q_value; + k_reg[r] = k_value; + q_norm += q_value * q_value; + k_norm += k_value * k_value; + } + + if (use_qk_l2norm) { + q_norm = subgroup_broadcast(subgroup_sum(q_norm, WIDTH), WIDTH); + k_norm = subgroup_broadcast(subgroup_sum(k_norm, WIDTH), WIDTH); + const float q_inv = rsqrtf(q_norm + 1.0e-6F); + const float k_inv = rsqrtf(k_norm + 1.0e-6F); +#pragma unroll + for (int r = 0; r < rows_per_lane; ++r) { + q_reg[r] *= q_inv; + k_reg[r] *= k_inv; + } + } + + const int64_t gate_index = static_cast(t) * v_heads + hv; + float gate_value = 0.0F; + float beta_value = 0.0F; + if (threadIdx.x == 0) { + const float x = load_as_float(a, gate_index) + load_as_float(dt_bias, hv); + const float softplus_x = + x <= 20.0F ? log1pf(__expf(x)) : x; + const float g_value = -__expf(A_log[hv]) * softplus_x; + gate_value = __expf(g_value); + const float b_value = load_as_float(b, gate_index); + beta_value = 1.0F / (1.0F + __expf(-b_value)); + } + gate_value = __shfl_sync(0xffffffffU, gate_value, 0); + beta_value = __shfl_sync(0xffffffffU, beta_value, 0); + + float kv_partial[COLS]; +#pragma unroll + for (int c = 0; c < COLS; ++c) { + kv_partial[c] = 0.0F; + } +#pragma unroll + for (int r = 0; r < rows_per_lane; ++r) { +#pragma unroll + for (int c = 0; c < COLS; ++c) { + kv_partial[c] += state_shard[c][r] * k_reg[r]; + } + } + + float delta[COLS]; +#pragma unroll + for (int c = 0; c < COLS; ++c) { + const int col = col_base + c; + const float kv_col = subgroup_sum(kv_partial[c], WIDTH); + float delta_value = 0.0F; + if (lane == 0 && col < V) { + const int64_t v_index = + mixed_base + 2 * q_heads * K + hv * V + col; + delta_value = + (load_as_float(mixed_qkv, v_index) - gate_value * kv_col) * beta_value; + } + delta[c] = subgroup_broadcast(delta_value, WIDTH); + } + + float attn_partial[COLS]; +#pragma unroll + for (int c = 0; c < COLS; ++c) { + attn_partial[c] = 0.0F; + } +#pragma unroll + for (int r = 0; r < rows_per_lane; ++r) { +#pragma unroll + for (int c = 0; c < COLS; ++c) { + const float new_state = + fmaf(k_reg[r], delta[c], gate_value * state_shard[c][r]); + state_shard[c][r] = new_state; + attn_partial[c] += new_state * q_reg[r]; + } + } +#pragma unroll + for (int c = 0; c < COLS; ++c) { + attn_partial[c] = subgroup_sum(attn_partial[c], WIDTH); + } + + if (lane == 0) { + const int64_t out_base = (static_cast(t) * v_heads + hv) * V; +#pragma unroll + for (int c = 0; c < COLS; ++c) { + const int col = col_base + c; + if (col < V) { + store_from_float(output, out_base + col, attn_partial[c] * scale); + } + } + } + +#pragma unroll + for (int c = 0; c < COLS; ++c) { + const int col = col_base + c; +#pragma unroll + for (int r = 0; r < rows_per_lane; ++r) { + const int row = r * WIDTH + lane; + if (col < V) { + const int64_t state_index = + static_cast(slot) * state_slot_stride + + ((static_cast(hv) * V + col) * K) + row; + store_from_float(state, state_index, state_shard[c][r]); + } + } + } +} + +template +__global__ void gdn_decode_mixed_qkv_ddtree_state_kernel( + const scalar_t* __restrict__ mixed_qkv, + const scalar_t* __restrict__ a, + const scalar_t* __restrict__ b, + const float* __restrict__ A_log, + const bias_t* __restrict__ dt_bias, + state_t* __restrict__ state, + const int32_t* __restrict__ state_indices, + const int32_t* __restrict__ parent_ids, + const int32_t* __restrict__ num_accepted_tokens, + const int32_t* __restrict__ cu_seqlens, + scalar_t* __restrict__ output, + int num_sequences, + int max_state_tokens, + int tokens, + int slots, + int64_t state_slot_stride, + int q_heads, + int v_heads, + int qkv_stride, + float scale, + bool use_qk_l2norm) { + static_assert(K % WIDTH == 0); + constexpr int subgroups_per_warp = 32 / WIDTH; + constexpr int rows_per_lane = K / WIDTH; + + const int hv = blockIdx.x; + const int n = blockIdx.y; + const int subgroup = threadIdx.x / WIDTH; + const int lane = threadIdx.x % WIDTH; + const int group_base = + (blockIdx.z * blockDim.y + threadIdx.y) * subgroups_per_warp + subgroup; + const int col_base = group_base * COLS; + const int hq = hv / (v_heads / q_heads); + if (n >= num_sequences) { + return; + } + + const int seq_start = cu_seqlens[n]; + const int seq_end = cu_seqlens[n + 1]; + const int seq_tokens = seq_end - seq_start; + if (seq_tokens <= 0 || seq_start < 0 || seq_end > tokens) { + return; + } + + int selector = num_accepted_tokens[n] - 1; + selector = selector < 0 ? 0 : selector; + selector = selector >= max_state_tokens ? max_state_tokens - 1 : selector; + int32_t state_idx = state_indices[n * max_state_tokens + selector]; + + float state_shard[COLS][rows_per_lane]; +#pragma unroll + for (int c = 0; c < COLS; ++c) { + const int col = col_base + c; +#pragma unroll + for (int r = 0; r < rows_per_lane; ++r) { + const int row = r * WIDTH + lane; + float value = 0.0F; + if (col < V && state_idx >= 0 && state_idx < slots) { + const int64_t state_index = + static_cast(state_idx) * state_slot_stride + + ((static_cast(hv) * V + col) * K) + row; + value = load_as_float(state, state_index); + } + state_shard[c][r] = value; + } + } + + for (int local_t = 0; local_t < seq_tokens && local_t < max_state_tokens; + ++local_t) { + if (local_t > 0) { + int parent_t = parent_ids[n * max_state_tokens + local_t]; + parent_t = parent_t < 0 ? 0 : parent_t; + parent_t = + parent_t >= max_state_tokens ? max_state_tokens - 1 : parent_t; + const int32_t parent_state_idx = + state_indices[n * max_state_tokens + parent_t]; +#pragma unroll + for (int c = 0; c < COLS; ++c) { + const int col = col_base + c; +#pragma unroll + for (int r = 0; r < rows_per_lane; ++r) { + const int row = r * WIDTH + lane; + float value = 0.0F; + if (col < V && parent_state_idx >= 0 && parent_state_idx < slots) { + const int64_t state_index = + static_cast(parent_state_idx) * state_slot_stride + + ((static_cast(hv) * V + col) * K) + row; + value = load_as_float(state, state_index); + } + state_shard[c][r] = value; + } + } + } + + const int t = seq_start + local_t; + const int64_t mixed_base = static_cast(t) * qkv_stride; + float k_reg[rows_per_lane]; + float q_reg[rows_per_lane]; + float q_norm = 0.0F; + float k_norm = 0.0F; +#pragma unroll + for (int r = 0; r < rows_per_lane; ++r) { + const int row = r * WIDTH + lane; + const int64_t q_index = mixed_base + hq * K + row; + const int64_t k_index = mixed_base + q_heads * K + hq * K + row; + const float q_value = load_as_float(mixed_qkv, q_index); + const float k_value = load_as_float(mixed_qkv, k_index); + q_reg[r] = q_value; + k_reg[r] = k_value; + q_norm += q_value * q_value; + k_norm += k_value * k_value; + } + + if (use_qk_l2norm) { + q_norm = subgroup_broadcast(subgroup_sum(q_norm, WIDTH), WIDTH); + k_norm = subgroup_broadcast(subgroup_sum(k_norm, WIDTH), WIDTH); + const float q_inv = rsqrtf(q_norm + 1.0e-6F); + const float k_inv = rsqrtf(k_norm + 1.0e-6F); +#pragma unroll + for (int r = 0; r < rows_per_lane; ++r) { + q_reg[r] *= q_inv; + k_reg[r] *= k_inv; + } + } + + const int64_t gate_index = static_cast(t) * v_heads + hv; + float gate_value = 0.0F; + float beta_value = 0.0F; + if (threadIdx.x == 0) { + const float x = load_as_float(a, gate_index) + load_as_float(dt_bias, hv); + const float softplus_x = x <= 20.0F ? log1pf(__expf(x)) : x; + const float g_value = -__expf(A_log[hv]) * softplus_x; + gate_value = __expf(g_value); + const float b_value = load_as_float(b, gate_index); + beta_value = 1.0F / (1.0F + __expf(-b_value)); + } + gate_value = __shfl_sync(0xffffffffU, gate_value, 0); + beta_value = __shfl_sync(0xffffffffU, beta_value, 0); + + float kv_partial[COLS]; +#pragma unroll + for (int c = 0; c < COLS; ++c) { + kv_partial[c] = 0.0F; + } +#pragma unroll + for (int r = 0; r < rows_per_lane; ++r) { +#pragma unroll + for (int c = 0; c < COLS; ++c) { + kv_partial[c] += state_shard[c][r] * k_reg[r]; + } + } + + float delta[COLS]; +#pragma unroll + for (int c = 0; c < COLS; ++c) { + const int col = col_base + c; + const float kv_col = subgroup_sum(kv_partial[c], WIDTH); + float delta_value = 0.0F; + if (lane == 0 && col < V) { + const int64_t v_index = mixed_base + 2 * q_heads * K + hv * V + col; + delta_value = + (load_as_float(mixed_qkv, v_index) - gate_value * kv_col) * + beta_value; + } + delta[c] = subgroup_broadcast(delta_value, WIDTH); + } + + float attn_partial[COLS]; +#pragma unroll + for (int c = 0; c < COLS; ++c) { + attn_partial[c] = 0.0F; + } +#pragma unroll + for (int r = 0; r < rows_per_lane; ++r) { +#pragma unroll + for (int c = 0; c < COLS; ++c) { + const float new_state = + fmaf(k_reg[r], delta[c], gate_value * state_shard[c][r]); + state_shard[c][r] = new_state; + attn_partial[c] += new_state * q_reg[r]; + } + } +#pragma unroll + for (int c = 0; c < COLS; ++c) { + attn_partial[c] = subgroup_sum(attn_partial[c], WIDTH); + } + + if (lane == 0) { + const int64_t out_base = (static_cast(t) * v_heads + hv) * V; +#pragma unroll + for (int c = 0; c < COLS; ++c) { + const int col = col_base + c; + if (col < V) { + store_from_float(output, out_base + col, attn_partial[c] * scale); + } + } + } + + const int32_t dst_state_idx = + state_indices[n * max_state_tokens + local_t]; +#pragma unroll + for (int c = 0; c < COLS; ++c) { + const int col = col_base + c; +#pragma unroll + for (int r = 0; r < rows_per_lane; ++r) { + const int row = r * WIDTH + lane; + if (col < V && dst_state_idx >= 0 && dst_state_idx < slots) { + const int64_t state_index = + static_cast(dst_state_idx) * state_slot_stride + + ((static_cast(hv) * V + col) * K) + row; + store_from_float(state, state_index, state_shard[c][r]); + } + } + } + } +} + +template +void launch_gdn_forward_typed(const scalar_t* q, + const scalar_t* k, + const scalar_t* v, + const gate_t* gate, + const beta_t* beta, + const state_t* initial_state, + scalar_t* output, + float* final_state, + int batch, + int tokens, + int q_heads, + int v_heads, + float scale, + bool gate_is_exp, + int column_groups_per_block, + cudaStream_t stream) { + constexpr int cols = V == 128 ? 4 : 1; + constexpr int width = K == 128 ? 16 : 32; + constexpr int groups_per_warp = 32 / width; + const dim3 block(32, column_groups_per_block); + const int groups = (V + cols - 1) / cols; + const int z = (groups + column_groups_per_block * groups_per_warp - 1) / + (column_groups_per_block * groups_per_warp); + const dim3 grid(v_heads, batch, z); + if (gate_is_exp) { + gdn_forward_kernel + <<>>(q, + k, + v, + gate, + beta, + initial_state, + output, + final_state, + batch, + tokens, + q_heads, + v_heads, + scale, + gate_is_exp); + } else { + gdn_forward_kernel + <<>>(q, + k, + v, + gate, + beta, + initial_state, + output, + final_state, + batch, + tokens, + q_heads, + v_heads, + scale, + gate_is_exp); + } +} + +template +void launch_gdn_forward_kv(const scalar_t* q, + const scalar_t* k, + const scalar_t* v, + const gate_t* gate, + const beta_t* beta, + const state_t* initial_state, + scalar_t* output, + float* final_state, + int batch, + int tokens, + int q_heads, + int v_heads, + int k_dim, + int v_dim, + float scale, + bool gate_is_exp, + int column_groups_per_block, + cudaStream_t stream) { + TORCH_CHECK(k_dim == 128 && v_dim == 128, + "SM70 FlashQLA backend currently supports K=V=128"); + launch_gdn_forward_typed( + q, k, v, gate, beta, initial_state, output, final_state, batch, tokens, + q_heads, v_heads, scale, gate_is_exp, column_groups_per_block, stream); +} + +template +void launch_gdn_forward_vlk_varlen_typed(const scalar_t* q, + const scalar_t* k, + const scalar_t* v, + const gate_t* gate, + const beta_t* beta, + const state_t* initial_state, + const int32_t* cu_seqlens, + scalar_t* output, + float* final_state, + int num_sequences, + int tokens, + int q_heads, + int v_heads, + float scale, + bool gate_is_exp, + int column_groups_per_block, + cudaStream_t stream) { + constexpr int cols = V == 128 ? 4 : 1; + constexpr int width = K == 128 ? 16 : 32; + constexpr int groups_per_warp = 32 / width; + const dim3 block(32, column_groups_per_block); + const int groups = (V + cols - 1) / cols; + const int z = (groups + column_groups_per_block * groups_per_warp - 1) / + (column_groups_per_block * groups_per_warp); + const dim3 grid(v_heads, num_sequences, z); + if (gate_is_exp) { + gdn_forward_vlk_varlen_kernel + <<>>(q, + k, + v, + gate, + beta, + initial_state, + cu_seqlens, + output, + final_state, + q_heads, + v_heads, + scale, + gate_is_exp); + } else { + gdn_forward_vlk_varlen_kernel + <<>>(q, + k, + v, + gate, + beta, + initial_state, + cu_seqlens, + output, + final_state, + q_heads, + v_heads, + scale, + gate_is_exp); + } +} + +template +void launch_gdn_forward_vlk_varlen_kv(const scalar_t* q, + const scalar_t* k, + const scalar_t* v, + const gate_t* gate, + const beta_t* beta, + const state_t* initial_state, + const int32_t* cu_seqlens, + scalar_t* output, + float* final_state, + int num_sequences, + int tokens, + int q_heads, + int v_heads, + int k_dim, + int v_dim, + float scale, + bool gate_is_exp, + int column_groups_per_block, + cudaStream_t stream) { + TORCH_CHECK(k_dim == 128 && v_dim == 128, + "SM70 FlashQLA backend currently supports K=V=128"); + launch_gdn_forward_vlk_varlen_typed( + q, k, v, gate, beta, initial_state, cu_seqlens, output, final_state, + num_sequences, tokens, q_heads, v_heads, scale, gate_is_exp, + column_groups_per_block, stream); +} + +template +void launch_gdn_decode_mixed_qkv_global_state_typed( + const scalar_t* mixed_qkv, + const scalar_t* a, + const scalar_t* b, + const float* A_log, + const bias_t* dt_bias, + state_t* state, + const int32_t* state_indices, + scalar_t* output, + int tokens, + int slots, + int64_t state_slot_stride, + int q_heads, + int v_heads, + int qkv_stride, + float scale, + bool use_qk_l2norm, + int column_groups_per_block, + cudaStream_t stream) { + constexpr int cols = V == 128 ? 4 : 1; + constexpr int width = K == 128 ? 16 : 32; + constexpr int groups_per_warp = 32 / width; + const dim3 block(32, column_groups_per_block); + const int groups = (V + cols - 1) / cols; + const int z = (groups + column_groups_per_block * groups_per_warp - 1) / + (column_groups_per_block * groups_per_warp); + const dim3 grid(v_heads, tokens, z); + gdn_decode_mixed_qkv_global_state_kernel + <<>>(mixed_qkv, + a, + b, + A_log, + dt_bias, + state, + state_indices, + output, + tokens, + slots, + state_slot_stride, + q_heads, + v_heads, + qkv_stride, + scale, + use_qk_l2norm); +} + +template +void launch_gdn_decode_mixed_qkv_global_state_kv( + const scalar_t* mixed_qkv, + const scalar_t* a, + const scalar_t* b, + const float* A_log, + const bias_t* dt_bias, + state_t* state, + const int32_t* state_indices, + scalar_t* output, + int tokens, + int slots, + int64_t state_slot_stride, + int q_heads, + int v_heads, + int k_dim, + int v_dim, + int qkv_stride, + float scale, + bool use_qk_l2norm, + int column_groups_per_block, + cudaStream_t stream) { + TORCH_CHECK(k_dim == 128 && v_dim == 128, + "SM70 FlashQLA decode currently supports K=V=128"); + launch_gdn_decode_mixed_qkv_global_state_typed( + mixed_qkv, a, b, A_log, dt_bias, state, state_indices, output, tokens, + slots, state_slot_stride, q_heads, v_heads, qkv_stride, scale, use_qk_l2norm, + column_groups_per_block, stream); +} + +template +void launch_gdn_decode_mixed_qkv_ddtree_state_typed( + const scalar_t* mixed_qkv, + const scalar_t* a, + const scalar_t* b, + const float* A_log, + const bias_t* dt_bias, + state_t* state, + const int32_t* state_indices, + const int32_t* parent_ids, + const int32_t* num_accepted_tokens, + const int32_t* cu_seqlens, + scalar_t* output, + int num_sequences, + int max_state_tokens, + int tokens, + int slots, + int64_t state_slot_stride, + int q_heads, + int v_heads, + int qkv_stride, + float scale, + bool use_qk_l2norm, + int column_groups_per_block, + cudaStream_t stream) { + constexpr int cols = V == 128 ? 4 : 1; + constexpr int width = K == 128 ? 16 : 32; + constexpr int groups_per_warp = 32 / width; + const dim3 block(32, column_groups_per_block); + const int groups = (V + cols - 1) / cols; + const int z = (groups + column_groups_per_block * groups_per_warp - 1) / + (column_groups_per_block * groups_per_warp); + const dim3 grid(v_heads, num_sequences, z); + gdn_decode_mixed_qkv_ddtree_state_kernel + <<>>(mixed_qkv, + a, + b, + A_log, + dt_bias, + state, + state_indices, + parent_ids, + num_accepted_tokens, + cu_seqlens, + output, + num_sequences, + max_state_tokens, + tokens, + slots, + state_slot_stride, + q_heads, + v_heads, + qkv_stride, + scale, + use_qk_l2norm); +} + +template +void launch_gdn_decode_mixed_qkv_ddtree_state_kv( + const scalar_t* mixed_qkv, + const scalar_t* a, + const scalar_t* b, + const float* A_log, + const bias_t* dt_bias, + state_t* state, + const int32_t* state_indices, + const int32_t* parent_ids, + const int32_t* num_accepted_tokens, + const int32_t* cu_seqlens, + scalar_t* output, + int num_sequences, + int max_state_tokens, + int tokens, + int slots, + int64_t state_slot_stride, + int q_heads, + int v_heads, + int k_dim, + int v_dim, + int qkv_stride, + float scale, + bool use_qk_l2norm, + int column_groups_per_block, + cudaStream_t stream) { + TORCH_CHECK(k_dim == 128 && v_dim == 128, + "SM70 FlashQLA DDTree decode currently supports K=V=128"); + launch_gdn_decode_mixed_qkv_ddtree_state_typed( + mixed_qkv, a, b, A_log, dt_bias, state, state_indices, parent_ids, + num_accepted_tokens, cu_seqlens, output, num_sequences, max_state_tokens, + tokens, slots, state_slot_stride, q_heads, v_heads, qkv_stride, scale, + use_qk_l2norm, column_groups_per_block, stream); +} + +void validate_tensor(const torch::Tensor& tensor, + const char* name, + int64_t dims) { + TORCH_CHECK(tensor.is_cuda(), name, " must be a CUDA tensor"); + TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous"); + TORCH_CHECK(tensor.dim() == dims, name, " has wrong rank"); +} + +void validate_mixed_qkv_tensor(const torch::Tensor& tensor) { + TORCH_CHECK(tensor.is_cuda(), "mixed_qkv must be a CUDA tensor"); + TORCH_CHECK(tensor.dim() == 2, "mixed_qkv has wrong rank"); + TORCH_CHECK(tensor.stride(1) == 1, + "mixed_qkv must have dense columns"); + TORCH_CHECK(tensor.stride(0) >= tensor.size(1), + "mixed_qkv row stride must be >= logical width"); +} + +void validate_cuda_rank(const torch::Tensor& tensor, + const char* name, + int64_t dims) { + TORCH_CHECK(tensor.is_cuda(), name, " must be a CUDA tensor"); + TORCH_CHECK(tensor.dim() == dims, name, " has wrong rank"); +} + +void validate_activation_dtype(const torch::Tensor& tensor, const char* name) { + TORCH_CHECK(tensor.scalar_type() == torch::kFloat16 || + tensor.scalar_type() == torch::kBFloat16 || + tensor.scalar_type() == torch::kFloat32, + name, + " must be fp16, bf16, or fp32"); +} + +void validate_same_device(const torch::Tensor& tensor, + const torch::Tensor& reference, + const char* name) { + TORCH_CHECK(tensor.device() == reference.device(), + name, + " must be on the same CUDA device as q"); +} + +int get_column_groups_per_block(int tokens, + int q_heads, + int v_heads) { + const char* raw = std::getenv("FLASH_QLA_SM70_COLUMN_GROUPS_PER_BLOCK"); + if (raw == nullptr || raw[0] == '\0') { + // Real Qwen3.5/Qwen3.6 SM70 shapes are dominated by Hv=8/12/16/24/32 + // after TP. Keep this as a shape heuristic; the env above remains the + // escape hatch for benchmarks and future model-specific tuning. V100 has + // no native bf16, so the default is derived from the fp16 production path. + if (v_heads == 24 || v_heads == 48) { + return 2; + } + if (v_heads == 32) { + if (q_heads == 16) { + return tokens <= 1024 ? 4 : 2; + } + if (q_heads == 8) { + return tokens <= 1024 ? 2 : 1; + } + return 1; + } + if (v_heads >= 64) { + return 1; + } + if (v_heads == 16) { + if (q_heads == 8) { + return 2; + } + return tokens <= 1024 ? 2 : 1; + } + if (v_heads == 12) { + return tokens >= 1024 ? 2 : 1; + } + if (v_heads == 8) { + return tokens <= 1024 ? 1 : 4; + } + return 2; + } + const int value = std::atoi(raw); + TORCH_CHECK(value == 1 || value == 2 || value == 4 || value == 8, + "FLASH_QLA_SM70_COLUMN_GROUPS_PER_BLOCK must be one of 1, 2, 4, 8"); + return value; +} + +int get_ddtree_column_groups_per_block(int tokens, + int q_heads, + int v_heads) { + const char* raw = std::getenv("FLASH_QLA_SM70_DDTREE_COLUMN_GROUPS_PER_BLOCK"); + if (raw != nullptr && raw[0] != '\0') { + const int value = std::atoi(raw); + TORCH_CHECK(value == 1 || value == 2 || value == 4 || value == 8, + "FLASH_QLA_SM70_DDTREE_COLUMN_GROUPS_PER_BLOCK must be one of " + "1, 2, 4, 8"); + return value; + } + return get_column_groups_per_block(tokens, q_heads, v_heads); +} + +} // namespace + +int resolve_column_groups_per_block(int tokens, int q_heads, int v_heads) { + return get_column_groups_per_block(tokens, q_heads, v_heads); +} + +std::vector gdn_forward(torch::Tensor q, + torch::Tensor k, + torch::Tensor v, + torch::Tensor gate, + torch::Tensor beta, + c10::optional initial_state, + double scale, + bool output_final_state, + bool gate_is_exp) { + validate_tensor(q, "q", 4); + validate_tensor(k, "k", 4); + validate_tensor(v, "v", 4); + validate_tensor(gate, "gate", 3); + validate_tensor(beta, "beta", 3); + validate_activation_dtype(q, "q"); + validate_same_device(k, q, "k"); + validate_same_device(v, q, "v"); + validate_same_device(gate, q, "gate"); + validate_same_device(beta, q, "beta"); + + TORCH_CHECK(k.scalar_type() == q.scalar_type(), "k must match q dtype"); + TORCH_CHECK(v.scalar_type() == q.scalar_type(), "v must match q dtype"); + validate_activation_dtype(gate, "gate"); + validate_activation_dtype(beta, "beta"); + TORCH_CHECK(q.sizes() == k.sizes(), "q and k must have the same shape"); + + const int batch = static_cast(q.size(0)); + const int tokens = static_cast(q.size(1)); + const int q_heads = static_cast(q.size(2)); + const int k_dim = static_cast(q.size(3)); + const int v_heads = static_cast(v.size(2)); + const int v_dim = static_cast(v.size(3)); + TORCH_CHECK(v.size(0) == batch && v.size(1) == tokens, + "v must have shape [B, T, Hv, V] matching q/k"); + TORCH_CHECK(gate.size(0) == batch && gate.size(1) == tokens && + gate.size(2) == v_heads, + "gate must have shape [B, T, Hv]"); + TORCH_CHECK(beta.sizes() == gate.sizes(), + "beta must have the same shape as gate"); + TORCH_CHECK(v_heads % q_heads == 0, "Hv must be divisible by Hq"); + TORCH_CHECK(k_dim == 128 && v_dim == 128, + "SM70 FlashQLA backend currently supports K=V=128"); + const at::cuda::OptionalCUDAGuard device_guard(device_of(q)); + + const void* initial_ptr = nullptr; + at::ScalarType state_dtype = torch::kFloat32; + if (initial_state.has_value() && initial_state.value().defined()) { + const auto& h0 = initial_state.value(); + validate_tensor(h0, "initial_state", 4); + validate_same_device(h0, q, "initial_state"); + TORCH_CHECK(h0.scalar_type() == torch::kFloat16 || + h0.scalar_type() == torch::kBFloat16 || + h0.scalar_type() == torch::kFloat32, + "initial_state must be fp16, bf16, or fp32"); + TORCH_CHECK(h0.size(0) == batch && h0.size(1) == v_heads && + h0.size(2) == k_dim && h0.size(3) == v_dim, + "initial_state must have shape [B, Hv, K, V]"); + initial_ptr = h0.data_ptr(); + state_dtype = h0.scalar_type(); + } + + auto output = torch::empty_like(v); + auto final_state = output_final_state + ? torch::empty({batch, v_heads, k_dim, v_dim}, + q.options().dtype(torch::kFloat32)) + : torch::Tensor(); + float* final_state_ptr = + output_final_state ? final_state.data_ptr() : nullptr; + const int column_groups_per_block = + get_column_groups_per_block(tokens, q_heads, v_heads); + const auto stream = at::cuda::getCurrentCUDAStream(q.device().index()).stream(); + + auto dispatch_state = [&](auto q_ptr, auto k_ptr, auto v_ptr, auto gate_ptr, + auto beta_ptr, auto out_ptr) { + using scalar_t = std::remove_pointer_t; + using gate_t = std::remove_pointer_t; + using beta_t = std::remove_pointer_t; + if (initial_ptr == nullptr || state_dtype == torch::kFloat32) { + launch_gdn_forward_kv( + q_ptr, k_ptr, v_ptr, gate_ptr, beta_ptr, + reinterpret_cast(initial_ptr), out_ptr, + final_state_ptr, batch, tokens, q_heads, v_heads, k_dim, + v_dim, static_cast(scale), gate_is_exp, + column_groups_per_block, stream); + } else if (state_dtype == torch::kFloat16) { + launch_gdn_forward_kv( + q_ptr, k_ptr, v_ptr, gate_ptr, beta_ptr, + reinterpret_cast(initial_ptr), out_ptr, + final_state_ptr, batch, tokens, q_heads, v_heads, k_dim, + v_dim, static_cast(scale), gate_is_exp, + column_groups_per_block, stream); + } else { + launch_gdn_forward_kv( + q_ptr, k_ptr, v_ptr, gate_ptr, beta_ptr, + reinterpret_cast(initial_ptr), out_ptr, + final_state_ptr, batch, tokens, q_heads, v_heads, k_dim, + v_dim, static_cast(scale), gate_is_exp, + column_groups_per_block, stream); + } + }; + + auto dispatch_beta = [&](auto q_ptr, auto k_ptr, auto v_ptr, auto gate_ptr, + auto out_ptr) { + if (beta.scalar_type() == torch::kFloat16) { + dispatch_state(q_ptr, k_ptr, v_ptr, gate_ptr, beta.data_ptr(), + out_ptr); + } else if (beta.scalar_type() == torch::kBFloat16) { + dispatch_state(q_ptr, k_ptr, v_ptr, gate_ptr, + beta.data_ptr(), out_ptr); + } else { + dispatch_state(q_ptr, k_ptr, v_ptr, gate_ptr, beta.data_ptr(), + out_ptr); + } + }; + + auto dispatch_gate = [&](auto q_ptr, auto k_ptr, auto v_ptr, auto out_ptr) { + if (gate.scalar_type() == torch::kFloat16) { + dispatch_beta(q_ptr, k_ptr, v_ptr, gate.data_ptr(), out_ptr); + } else if (gate.scalar_type() == torch::kBFloat16) { + dispatch_beta(q_ptr, k_ptr, v_ptr, gate.data_ptr(), out_ptr); + } else { + dispatch_beta(q_ptr, k_ptr, v_ptr, gate.data_ptr(), out_ptr); + } + }; + + if (q.scalar_type() == torch::kFloat16) { + dispatch_gate(q.data_ptr(), k.data_ptr(), + v.data_ptr(), output.data_ptr()); + } else if (q.scalar_type() == torch::kBFloat16) { + dispatch_gate(q.data_ptr(), k.data_ptr(), + v.data_ptr(), + output.data_ptr()); + } else { + dispatch_gate(q.data_ptr(), k.data_ptr(), v.data_ptr(), + output.data_ptr()); + } + + check_cuda(cudaGetLastError(), "gdn_forward launch"); + return {output, final_state}; +} + +std::vector gdn_forward_vlk_varlen( + torch::Tensor q, + torch::Tensor k, + torch::Tensor v, + torch::Tensor gate, + torch::Tensor beta, + c10::optional initial_state, + torch::Tensor cu_seqlens, + double scale, + bool output_final_state, + bool validate_cu_seqlens, + bool gate_is_exp, + c10::optional output_arg) { + validate_tensor(q, "q", 4); + validate_tensor(k, "k", 4); + validate_tensor(v, "v", 4); + validate_tensor(gate, "gate", 3); + validate_tensor(beta, "beta", 3); + validate_tensor(cu_seqlens, "cu_seqlens", 1); + validate_activation_dtype(q, "q"); + validate_same_device(k, q, "k"); + validate_same_device(v, q, "v"); + validate_same_device(gate, q, "gate"); + validate_same_device(beta, q, "beta"); + validate_same_device(cu_seqlens, q, "cu_seqlens"); + + TORCH_CHECK(k.scalar_type() == q.scalar_type(), "k must match q dtype"); + TORCH_CHECK(v.scalar_type() == q.scalar_type(), "v must match q dtype"); + validate_activation_dtype(gate, "gate"); + validate_activation_dtype(beta, "beta"); + TORCH_CHECK(cu_seqlens.scalar_type() == torch::kInt32, + "cu_seqlens must be int32"); + if (validate_cu_seqlens) { + TORCH_CHECK(cu_seqlens[0].item() == 0, + "cu_seqlens must start at 0"); + } + TORCH_CHECK(q.sizes() == k.sizes(), "q and k must have the same shape"); + + const int batch = static_cast(q.size(0)); + const int tokens = static_cast(q.size(1)); + const int q_heads = static_cast(q.size(2)); + const int k_dim = static_cast(q.size(3)); + const int v_heads = static_cast(v.size(2)); + const int v_dim = static_cast(v.size(3)); + const int num_sequences = static_cast(cu_seqlens.size(0) - 1); + if (validate_cu_seqlens) { + TORCH_CHECK(cu_seqlens[num_sequences].item() == tokens, + "cu_seqlens must end at the flattened token count"); + } + TORCH_CHECK(batch == 1, + "gdn_forward_vlk_varlen expects flattened q/k/v with batch=1"); + TORCH_CHECK(num_sequences > 0, "cu_seqlens must contain at least one sequence"); + TORCH_CHECK(v.size(0) == batch && v.size(1) == tokens, + "v must have shape [1, T, Hv, V] matching q/k"); + TORCH_CHECK(gate.size(0) == batch && gate.size(1) == tokens && + gate.size(2) == v_heads, + "gate must have shape [1, T, Hv]"); + TORCH_CHECK(beta.sizes() == gate.sizes(), + "beta must have the same shape as gate"); + TORCH_CHECK(v_heads % q_heads == 0, "Hv must be divisible by Hq"); + TORCH_CHECK(k_dim == 128 && v_dim == 128, + "SM70 FlashQLA backend currently supports K=V=128"); + const at::cuda::OptionalCUDAGuard device_guard(device_of(q)); + + const void* initial_ptr = nullptr; + at::ScalarType state_dtype = torch::kFloat32; + if (initial_state.has_value() && initial_state.value().defined()) { + const auto& h0 = initial_state.value(); + validate_tensor(h0, "initial_state", 4); + validate_same_device(h0, q, "initial_state"); + TORCH_CHECK(h0.scalar_type() == torch::kFloat16 || + h0.scalar_type() == torch::kBFloat16 || + h0.scalar_type() == torch::kFloat32, + "initial_state must be fp16, bf16, or fp32"); + TORCH_CHECK(h0.size(0) == num_sequences && h0.size(1) == v_heads && + h0.size(2) == v_dim && h0.size(3) == k_dim, + "initial_state must have shape [N, Hv, V, K]"); + initial_ptr = h0.data_ptr(); + state_dtype = h0.scalar_type(); + } + + torch::Tensor output; + if (output_arg.has_value() && output_arg.value().defined()) { + output = output_arg.value(); + validate_tensor(output, "output", 4); + validate_same_device(output, q, "output"); + TORCH_CHECK(output.scalar_type() == v.scalar_type(), + "output must match v dtype"); + TORCH_CHECK(output.size(0) == batch && output.size(1) == tokens && + output.size(2) == v_heads && output.size(3) == v_dim, + "output must have shape [1, T, Hv, V]"); + } else { + output = torch::empty_like(v); + } + auto final_state = output_final_state + ? torch::empty({num_sequences, v_heads, v_dim, k_dim}, + q.options().dtype(torch::kFloat32)) + : torch::Tensor(); + float* final_state_ptr = + output_final_state ? final_state.data_ptr() : nullptr; + const int column_groups_per_block = + get_column_groups_per_block(tokens, q_heads, v_heads); + const auto stream = at::cuda::getCurrentCUDAStream(q.device().index()).stream(); + + auto dispatch_state = [&](auto q_ptr, auto k_ptr, auto v_ptr, auto gate_ptr, + auto beta_ptr, auto out_ptr) { + using scalar_t = std::remove_pointer_t; + using gate_t = std::remove_pointer_t; + using beta_t = std::remove_pointer_t; + if (initial_ptr == nullptr || state_dtype == torch::kFloat32) { + launch_gdn_forward_vlk_varlen_kv( + q_ptr, k_ptr, v_ptr, gate_ptr, beta_ptr, + reinterpret_cast(initial_ptr), + cu_seqlens.data_ptr(), out_ptr, final_state_ptr, + num_sequences, tokens, q_heads, v_heads, k_dim, v_dim, + static_cast(scale), gate_is_exp, column_groups_per_block, + stream); + } else if (state_dtype == torch::kFloat16) { + launch_gdn_forward_vlk_varlen_kv( + q_ptr, k_ptr, v_ptr, gate_ptr, beta_ptr, + reinterpret_cast(initial_ptr), + cu_seqlens.data_ptr(), out_ptr, final_state_ptr, + num_sequences, tokens, q_heads, v_heads, k_dim, v_dim, + static_cast(scale), gate_is_exp, column_groups_per_block, + stream); + } else { + launch_gdn_forward_vlk_varlen_kv( + q_ptr, k_ptr, v_ptr, gate_ptr, beta_ptr, + reinterpret_cast(initial_ptr), + cu_seqlens.data_ptr(), out_ptr, final_state_ptr, + num_sequences, tokens, q_heads, v_heads, k_dim, v_dim, + static_cast(scale), gate_is_exp, column_groups_per_block, + stream); + } + }; + + auto dispatch_beta = [&](auto q_ptr, auto k_ptr, auto v_ptr, auto gate_ptr, + auto out_ptr) { + if (beta.scalar_type() == torch::kFloat16) { + dispatch_state(q_ptr, k_ptr, v_ptr, gate_ptr, beta.data_ptr(), + out_ptr); + } else if (beta.scalar_type() == torch::kBFloat16) { + dispatch_state(q_ptr, k_ptr, v_ptr, gate_ptr, + beta.data_ptr(), out_ptr); + } else { + dispatch_state(q_ptr, k_ptr, v_ptr, gate_ptr, beta.data_ptr(), + out_ptr); + } + }; + + auto dispatch_gate = [&](auto q_ptr, auto k_ptr, auto v_ptr, auto out_ptr) { + if (gate.scalar_type() == torch::kFloat16) { + dispatch_beta(q_ptr, k_ptr, v_ptr, gate.data_ptr(), out_ptr); + } else if (gate.scalar_type() == torch::kBFloat16) { + dispatch_beta(q_ptr, k_ptr, v_ptr, gate.data_ptr(), out_ptr); + } else { + dispatch_beta(q_ptr, k_ptr, v_ptr, gate.data_ptr(), out_ptr); + } + }; + + if (q.scalar_type() == torch::kFloat16) { + dispatch_gate(q.data_ptr(), k.data_ptr(), + v.data_ptr(), output.data_ptr()); + } else if (q.scalar_type() == torch::kBFloat16) { + dispatch_gate(q.data_ptr(), k.data_ptr(), + v.data_ptr(), + output.data_ptr()); + } else { + dispatch_gate(q.data_ptr(), k.data_ptr(), v.data_ptr(), + output.data_ptr()); + } + + check_cuda(cudaGetLastError(), "gdn_forward_vlk_varlen launch"); + return {output, final_state}; +} + +void gdn_decode_mixed_qkv_global_state(torch::Tensor 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, + double scale, + bool use_qk_l2norm) { + validate_mixed_qkv_tensor(mixed_qkv); + validate_tensor(a, "a", 2); + validate_tensor(b, "b", 2); + validate_tensor(A_log, "A_log", 1); + validate_tensor(dt_bias, "dt_bias", 1); + validate_cuda_rank(state, "state", 4); + validate_tensor(state_indices, "state_indices", 1); + validate_tensor(output, "output", 3); + validate_activation_dtype(mixed_qkv, "mixed_qkv"); + validate_activation_dtype(a, "a"); + validate_activation_dtype(b, "b"); + validate_activation_dtype(state, "state"); + validate_activation_dtype(output, "output"); + validate_same_device(a, mixed_qkv, "a"); + validate_same_device(b, mixed_qkv, "b"); + validate_same_device(A_log, mixed_qkv, "A_log"); + validate_same_device(dt_bias, mixed_qkv, "dt_bias"); + validate_same_device(state, mixed_qkv, "state"); + validate_same_device(state_indices, mixed_qkv, "state_indices"); + validate_same_device(output, mixed_qkv, "output"); + + TORCH_CHECK(mixed_qkv.scalar_type() == a.scalar_type(), + "a must match mixed_qkv dtype"); + TORCH_CHECK(mixed_qkv.scalar_type() == b.scalar_type(), + "b must match mixed_qkv dtype"); + TORCH_CHECK(mixed_qkv.scalar_type() == output.scalar_type(), + "output must match mixed_qkv dtype"); + TORCH_CHECK(A_log.scalar_type() == torch::kFloat32, + "A_log must be float32"); + validate_activation_dtype(dt_bias, "dt_bias"); + TORCH_CHECK(state_indices.scalar_type() == torch::kInt32, + "state_indices must be int32"); + const int tokens = static_cast(mixed_qkv.size(0)); + const int slots = static_cast(state.size(0)); + const int v_heads = static_cast(state.size(1)); + const int v_dim = static_cast(state.size(2)); + const int k_dim = static_cast(state.size(3)); + const int64_t state_slot_stride = state.stride(0); + TORCH_CHECK(tokens > 0, "decode tokens must be positive"); + TORCH_CHECK(v_dim == 128 && k_dim == 128, + "SM70 FlashQLA decode currently supports K=V=128"); + TORCH_CHECK(state.stride(1) == v_dim * k_dim && state.stride(2) == k_dim && + state.stride(3) == 1, + "state inner layout must be [slots,Hv,V,K] with contiguous " + "[Hv,V,K] pages"); + TORCH_CHECK(a.size(0) == tokens && b.size(0) == tokens, + "a/b must match mixed_qkv token count"); + TORCH_CHECK(a.size(1) == v_heads && b.size(1) == v_heads, + "a/b must have HV columns matching state"); + TORCH_CHECK(A_log.size(0) == v_heads && dt_bias.size(0) == v_heads, + "A_log/dt_bias must have HV elements"); + TORCH_CHECK(state_indices.size(0) == tokens, + "state_indices must have one entry per decode token"); + TORCH_CHECK(output.size(0) == tokens && output.size(1) == v_heads && + output.size(2) == v_dim, + "output must have shape [tokens, Hv, V]"); + + const int qkv_dim = static_cast(mixed_qkv.size(1)); + const int qk_dim = qkv_dim - v_heads * v_dim; + TORCH_CHECK(qk_dim > 0 && qk_dim % 2 == 0, + "invalid packed mixed_qkv last dimension"); + const int q_dim = qk_dim / 2; + TORCH_CHECK(q_dim % k_dim == 0, + "packed Q dimension must be divisible by K"); + const int q_heads = q_dim / k_dim; + TORCH_CHECK(q_heads > 0 && v_heads % q_heads == 0, + "invalid H/HV inferred from mixed_qkv and state"); + const int qkv_stride = static_cast(mixed_qkv.stride(0)); + const int column_groups_per_block = + get_column_groups_per_block(tokens, q_heads, v_heads); + const at::cuda::OptionalCUDAGuard device_guard(device_of(mixed_qkv)); + const auto stream = + at::cuda::getCurrentCUDAStream(mixed_qkv.device().index()).stream(); + + auto dispatch_state = [&](auto mixed_ptr, auto a_ptr, auto b_ptr, + auto dt_bias_ptr, auto out_ptr) { + using scalar_t = std::remove_pointer_t; + using bias_t = std::remove_pointer_t; + if (state.scalar_type() == torch::kFloat32) { + launch_gdn_decode_mixed_qkv_global_state_kv( + mixed_ptr, a_ptr, b_ptr, A_log.data_ptr(), + dt_bias_ptr, state.data_ptr(), + state_indices.data_ptr(), out_ptr, tokens, slots, + state_slot_stride, q_heads, v_heads, k_dim, v_dim, qkv_stride, + static_cast(scale), use_qk_l2norm, column_groups_per_block, + stream); + } else if (state.scalar_type() == torch::kFloat16) { + launch_gdn_decode_mixed_qkv_global_state_kv( + mixed_ptr, a_ptr, b_ptr, A_log.data_ptr(), + dt_bias_ptr, state.data_ptr(), + state_indices.data_ptr(), out_ptr, tokens, slots, + state_slot_stride, q_heads, v_heads, k_dim, v_dim, qkv_stride, + static_cast(scale), use_qk_l2norm, column_groups_per_block, + stream); + } else { + launch_gdn_decode_mixed_qkv_global_state_kv( + mixed_ptr, a_ptr, b_ptr, A_log.data_ptr(), + dt_bias_ptr, state.data_ptr(), + state_indices.data_ptr(), out_ptr, tokens, slots, + state_slot_stride, q_heads, v_heads, k_dim, v_dim, qkv_stride, + static_cast(scale), use_qk_l2norm, column_groups_per_block, + stream); + } + }; + + auto dispatch_dt_bias = [&](auto mixed_ptr, auto a_ptr, auto b_ptr, + auto out_ptr) { + if (dt_bias.scalar_type() == torch::kFloat16) { + dispatch_state(mixed_ptr, a_ptr, b_ptr, dt_bias.data_ptr(), + out_ptr); + } else if (dt_bias.scalar_type() == torch::kBFloat16) { + dispatch_state(mixed_ptr, a_ptr, b_ptr, dt_bias.data_ptr(), + out_ptr); + } else { + dispatch_state(mixed_ptr, a_ptr, b_ptr, dt_bias.data_ptr(), + out_ptr); + } + }; + + if (mixed_qkv.scalar_type() == torch::kFloat16) { + dispatch_dt_bias(mixed_qkv.data_ptr(), a.data_ptr(), + b.data_ptr(), output.data_ptr()); + } else if (mixed_qkv.scalar_type() == torch::kBFloat16) { + dispatch_dt_bias(mixed_qkv.data_ptr(), + a.data_ptr(), + b.data_ptr(), + output.data_ptr()); + } else { + dispatch_dt_bias(mixed_qkv.data_ptr(), a.data_ptr(), + b.data_ptr(), output.data_ptr()); + } + + check_cuda(cudaGetLastError(), "gdn_decode_mixed_qkv_global_state launch"); +} + +void gdn_decode_mixed_qkv_ddtree_state(torch::Tensor 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, + double scale, + bool use_qk_l2norm) { + validate_mixed_qkv_tensor(mixed_qkv); + validate_tensor(a, "a", 2); + validate_tensor(b, "b", 2); + validate_tensor(A_log, "A_log", 1); + validate_tensor(dt_bias, "dt_bias", 1); + validate_cuda_rank(state, "state", 4); + validate_tensor(state_indices, "state_indices", 2); + validate_tensor(parent_ids, "parent_ids", 2); + validate_tensor(num_accepted_tokens, "num_accepted_tokens", 1); + validate_tensor(cu_seqlens, "cu_seqlens", 1); + validate_tensor(output, "output", 3); + validate_activation_dtype(mixed_qkv, "mixed_qkv"); + validate_activation_dtype(a, "a"); + validate_activation_dtype(b, "b"); + validate_activation_dtype(state, "state"); + validate_activation_dtype(output, "output"); + validate_same_device(a, mixed_qkv, "a"); + validate_same_device(b, mixed_qkv, "b"); + validate_same_device(A_log, mixed_qkv, "A_log"); + validate_same_device(dt_bias, mixed_qkv, "dt_bias"); + validate_same_device(state, mixed_qkv, "state"); + validate_same_device(state_indices, mixed_qkv, "state_indices"); + validate_same_device(parent_ids, mixed_qkv, "parent_ids"); + validate_same_device(num_accepted_tokens, mixed_qkv, "num_accepted_tokens"); + validate_same_device(cu_seqlens, mixed_qkv, "cu_seqlens"); + validate_same_device(output, mixed_qkv, "output"); + + TORCH_CHECK(mixed_qkv.scalar_type() == a.scalar_type(), + "a must match mixed_qkv dtype"); + TORCH_CHECK(mixed_qkv.scalar_type() == b.scalar_type(), + "b must match mixed_qkv dtype"); + TORCH_CHECK(mixed_qkv.scalar_type() == output.scalar_type(), + "output must match mixed_qkv dtype"); + TORCH_CHECK(A_log.scalar_type() == torch::kFloat32, + "A_log must be float32"); + validate_activation_dtype(dt_bias, "dt_bias"); + TORCH_CHECK(state_indices.scalar_type() == torch::kInt32, + "state_indices must be int32"); + TORCH_CHECK(parent_ids.scalar_type() == torch::kInt32, + "parent_ids must be int32"); + TORCH_CHECK(num_accepted_tokens.scalar_type() == torch::kInt32, + "num_accepted_tokens must be int32"); + TORCH_CHECK(cu_seqlens.scalar_type() == torch::kInt32, + "cu_seqlens must be int32"); + + const int tokens = static_cast(mixed_qkv.size(0)); + const int num_sequences = static_cast(state_indices.size(0)); + const int max_state_tokens = static_cast(state_indices.size(1)); + const int slots = static_cast(state.size(0)); + const int v_heads = static_cast(state.size(1)); + const int v_dim = static_cast(state.size(2)); + const int k_dim = static_cast(state.size(3)); + const int64_t state_slot_stride = state.stride(0); + + TORCH_CHECK(tokens > 0, "decode tokens must be positive"); + TORCH_CHECK(num_sequences > 0, "state_indices must have at least one row"); + TORCH_CHECK(max_state_tokens > 0, + "state_indices must have at least one state token column"); + TORCH_CHECK(v_dim == 128 && k_dim == 128, + "SM70 FlashQLA DDTree decode currently supports K=V=128"); + TORCH_CHECK(state.stride(1) == v_dim * k_dim && state.stride(2) == k_dim && + state.stride(3) == 1, + "state inner layout must be [slots,Hv,V,K] with contiguous " + "[Hv,V,K] pages"); + TORCH_CHECK(parent_ids.sizes() == state_indices.sizes(), + "parent_ids must match state_indices shape"); + TORCH_CHECK(num_accepted_tokens.size(0) == num_sequences, + "num_accepted_tokens must have one entry per sequence"); + TORCH_CHECK(cu_seqlens.size(0) == num_sequences + 1, + "cu_seqlens must have N + 1 entries"); + TORCH_CHECK(a.size(0) == tokens && b.size(0) == tokens, + "a/b must match mixed_qkv token count"); + TORCH_CHECK(a.size(1) == v_heads && b.size(1) == v_heads, + "a/b must have HV columns matching state"); + TORCH_CHECK(A_log.size(0) == v_heads && dt_bias.size(0) == v_heads, + "A_log/dt_bias must have HV elements"); + TORCH_CHECK(output.size(0) == tokens && output.size(1) == v_heads && + output.size(2) == v_dim, + "output must have shape [tokens, Hv, V]"); + + const int qkv_dim = static_cast(mixed_qkv.size(1)); + const int qk_dim = qkv_dim - v_heads * v_dim; + TORCH_CHECK(qk_dim > 0 && qk_dim % 2 == 0, + "mixed_qkv width must be q + k + v"); + TORCH_CHECK(qk_dim % (2 * k_dim) == 0, + "q/k packed width must be divisible by K"); + const int q_heads = qk_dim / (2 * k_dim); + TORCH_CHECK(q_heads > 0, "q_heads must be positive"); + TORCH_CHECK(v_heads % q_heads == 0, "Hv must be divisible by Hq"); + const int qkv_stride = static_cast(mixed_qkv.stride(0)); + const at::cuda::OptionalCUDAGuard device_guard(device_of(mixed_qkv)); + const int column_groups_per_block = + get_ddtree_column_groups_per_block(tokens, q_heads, v_heads); + const auto stream = + at::cuda::getCurrentCUDAStream(mixed_qkv.device().index()).stream(); + + auto dispatch_state = [&](auto mixed_ptr, auto a_ptr, auto b_ptr, + auto dt_bias_ptr, auto out_ptr) { + using scalar_t = std::remove_pointer_t; + using bias_t = std::remove_pointer_t; + if (state.scalar_type() == torch::kFloat32) { + launch_gdn_decode_mixed_qkv_ddtree_state_kv( + mixed_ptr, a_ptr, b_ptr, A_log.data_ptr(), dt_bias_ptr, + state.data_ptr(), state_indices.data_ptr(), + parent_ids.data_ptr(), + num_accepted_tokens.data_ptr(), + cu_seqlens.data_ptr(), out_ptr, num_sequences, + max_state_tokens, tokens, slots, state_slot_stride, q_heads, v_heads, + k_dim, v_dim, qkv_stride, static_cast(scale), use_qk_l2norm, + column_groups_per_block, stream); + } else if (state.scalar_type() == torch::kFloat16) { + launch_gdn_decode_mixed_qkv_ddtree_state_kv( + mixed_ptr, a_ptr, b_ptr, A_log.data_ptr(), dt_bias_ptr, + state.data_ptr(), state_indices.data_ptr(), + parent_ids.data_ptr(), + num_accepted_tokens.data_ptr(), + cu_seqlens.data_ptr(), out_ptr, num_sequences, + max_state_tokens, tokens, slots, state_slot_stride, q_heads, v_heads, + k_dim, v_dim, qkv_stride, static_cast(scale), use_qk_l2norm, + column_groups_per_block, stream); + } else { + launch_gdn_decode_mixed_qkv_ddtree_state_kv( + mixed_ptr, a_ptr, b_ptr, A_log.data_ptr(), dt_bias_ptr, + state.data_ptr(), state_indices.data_ptr(), + parent_ids.data_ptr(), + num_accepted_tokens.data_ptr(), + cu_seqlens.data_ptr(), out_ptr, num_sequences, + max_state_tokens, tokens, slots, state_slot_stride, q_heads, v_heads, + k_dim, v_dim, qkv_stride, static_cast(scale), use_qk_l2norm, + column_groups_per_block, stream); + } + }; + + auto dispatch_dt_bias = [&](auto mixed_ptr, auto a_ptr, auto b_ptr, + auto out_ptr) { + if (dt_bias.scalar_type() == torch::kFloat16) { + dispatch_state(mixed_ptr, a_ptr, b_ptr, dt_bias.data_ptr(), + out_ptr); + } else if (dt_bias.scalar_type() == torch::kBFloat16) { + dispatch_state(mixed_ptr, a_ptr, b_ptr, + dt_bias.data_ptr(), out_ptr); + } else { + dispatch_state(mixed_ptr, a_ptr, b_ptr, dt_bias.data_ptr(), + out_ptr); + } + }; + + if (mixed_qkv.scalar_type() == torch::kFloat16) { + dispatch_dt_bias(mixed_qkv.data_ptr(), a.data_ptr(), + b.data_ptr(), output.data_ptr()); + } else if (mixed_qkv.scalar_type() == torch::kBFloat16) { + dispatch_dt_bias(mixed_qkv.data_ptr(), + a.data_ptr(), b.data_ptr(), + output.data_ptr()); + } else { + dispatch_dt_bias(mixed_qkv.data_ptr(), a.data_ptr(), + b.data_ptr(), output.data_ptr()); + } + + check_cuda(cudaGetLastError(), "gdn_decode_mixed_qkv_ddtree_state launch"); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("gdn_forward", &gdn_forward, "SM70/SM75 FlashQLA GDN forward"); + m.def("gdn_forward_vlk_varlen", + &gdn_forward_vlk_varlen, + "SM70/SM75 FlashQLA GDN forward for vLLM [N,Hv,V,K] state"); + m.def("gdn_decode_mixed_qkv_global_state", + &gdn_decode_mixed_qkv_global_state, + "SM70/SM75 FlashQLA fused mixed-QKV decode for vLLM global state"); + m.def("gdn_decode_mixed_qkv_ddtree_state", + &gdn_decode_mixed_qkv_ddtree_state, + "SM70/SM75 FlashQLA mixed-QKV DDTree decode for vLLM global state"); + m.def("resolve_column_groups_per_block", + &resolve_column_groups_per_block, + "Resolve SM70/SM75 FlashQLA GDN column groups per block"); +} diff --git a/qwen3_6_scripts/flash_qla_sm70/fused_fwd.py b/qwen3_6_scripts/flash_qla_sm70/fused_fwd.py new file mode 100644 index 00000000..6e71c2a3 --- /dev/null +++ b/qwen3_6_scripts/flash_qla_sm70/fused_fwd.py @@ -0,0 +1,490 @@ +# 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") + 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)) diff --git a/qwen3_6_scripts/flash_qla_sm70/naive_gdn.py b/qwen3_6_scripts/flash_qla_sm70/naive_gdn.py new file mode 100644 index 00000000..cd0cf0d1 --- /dev/null +++ b/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 diff --git a/qwen3_6_scripts/patch_ops.sh b/qwen3_6_scripts/patch_ops.sh index f71961a6..31535aa2 100755 --- a/qwen3_6_scripts/patch_ops.sh +++ b/qwen3_6_scripts/patch_ops.sh @@ -175,6 +175,22 @@ 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 + # 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_sm70 "$VLLM2/model_executor/models/flash_qla_sm70" 2>/dev/null || true + fi +fi diff --git a/qwen3_6_scripts/qwen3_5.py b/qwen3_6_scripts/qwen3_5.py index 93fab699..8754b732 100644 --- a/qwen3_6_scripts/qwen3_5.py +++ b/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 # ---------------------------------------------------------------------------