feat(SM70): add 1Cat-vLLM FlashQLA fused GDN CUDA kernel for BI-V100
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
This commit is contained in:
Claude
14
qwen3_6_scripts/flash_qla_sm70/__init__.py
Normal file
14
qwen3_6_scripts/flash_qla_sm70/__init__.py
Normal file
@@ -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",
|
||||
]
|
||||
1919
qwen3_6_scripts/flash_qla_sm70/csrc/gdn_forward.cu
Normal file
1919
qwen3_6_scripts/flash_qla_sm70/csrc/gdn_forward.cu
Normal file
File diff suppressed because it is too large
Load Diff
490
qwen3_6_scripts/flash_qla_sm70/fused_fwd.py
Normal file
490
qwen3_6_scripts/flash_qla_sm70/fused_fwd.py
Normal file
@@ -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))
|
||||
161
qwen3_6_scripts/flash_qla_sm70/naive_gdn.py
Normal file
161
qwen3_6_scripts/flash_qla_sm70/naive_gdn.py
Normal file
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
Reference in New Issue
Block a user