141 lines
4.4 KiB
Python
141 lines
4.4 KiB
Python
"""
|
|
ex_engine/csrc/factor_gdn_flashqla.py — GDN Factor 5 via FlashQLA
|
|
|
|
Instead of a custom CUDA kernel, this loads the FlashQLA .so (compiled by
|
|
torch.utils.cpp_extension from gdn_forward.cu) and calls gdn_forward().
|
|
|
|
Real test on BI-V100 (from user doc):
|
|
output: torch.Size([1, 64, 4, 128]), state: torch.Size([1, 4, 128, 128])
|
|
NaN: False, abs mean: inf ← need to investigate inf issue
|
|
|
|
The FlashQLA kernel:
|
|
- Compiled via corex clang/16 with --cuda-gpu-arch=ivcore10
|
|
- Provides: gdn_forward(q, k, v, g, beta, initial_state, scale, output_final_state, head_first)
|
|
- Returns: (output, final_state)
|
|
- Full fp32 accumulation (no NaN)
|
|
"""
|
|
|
|
import os
|
|
import logging
|
|
import torch
|
|
from typing import Optional, Tuple
|
|
|
|
logger = logging.getLogger("ex_engine.gdn")
|
|
|
|
_flash_qla_ext = None
|
|
_flash_qla_available = False
|
|
|
|
|
|
def _load_flash_qla(build_dir: str = "/workspace/flash_qla_sm70") -> bool:
|
|
"""Load the pre-compiled FlashQLA extension."""
|
|
global _flash_qla_ext, _flash_qla_available
|
|
|
|
if _flash_qla_available:
|
|
return True
|
|
|
|
so_path = os.path.join(build_dir, "flash_qla_sm70_gdn.so")
|
|
|
|
# Try pre-compiled .so first
|
|
if os.path.exists(so_path):
|
|
try:
|
|
torch.ops.load_library(so_path)
|
|
_flash_qla_available = True
|
|
logger.info("FlashQLA GDN loaded from %s", so_path)
|
|
return True
|
|
except Exception as e:
|
|
logger.warning("FlashQLA .so load failed: %s, trying JIT compile", e)
|
|
|
|
# Try JIT compile
|
|
cu_path = os.path.join(build_dir, "csrc", "gdn_forward.cu")
|
|
if not os.path.exists(cu_path):
|
|
# Try alternate locations
|
|
for alt in [
|
|
"/workspace/qwen3_6_scripts/flash_qla_sm70/csrc/gdn_forward.cu",
|
|
"/workspace/flash_qla_sm70/csrc/gdn_forward.cu",
|
|
]:
|
|
if os.path.exists(alt):
|
|
cu_path = alt
|
|
break
|
|
|
|
if os.path.exists(cu_path):
|
|
try:
|
|
os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "7.0")
|
|
from torch.utils.cpp_extension import load
|
|
_flash_qla_ext = load(
|
|
name="flash_qla_sm70_gdn",
|
|
sources=[cu_path],
|
|
extra_cuda_cflags=["-O3"],
|
|
extra_cflags=["-O3"],
|
|
verbose=False,
|
|
)
|
|
_flash_qla_available = True
|
|
logger.info("FlashQLA GDN JIT compiled from %s", cu_path)
|
|
return True
|
|
except Exception as e:
|
|
logger.error("FlashQLA JIT compile failed: %s", e)
|
|
return False
|
|
|
|
logger.warning("FlashQLA GDN not found at %s", cu_path)
|
|
return False
|
|
|
|
|
|
def gdn_forward_flashqla(
|
|
query: torch.Tensor, # (B, L, H, D) half
|
|
key: torch.Tensor, # (B, L, H, D) half
|
|
value: torch.Tensor, # (B, L, Hv, V) half
|
|
gate: torch.Tensor, # (B, L, Hv) half
|
|
beta: torch.Tensor, # (B, L, Hv) half — already sigmoid'd
|
|
initial_state: Optional[torch.Tensor], # (B, Hv, K, V) or None
|
|
scale: float = None,
|
|
output_final_state: bool = True,
|
|
head_first: bool = False,
|
|
) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
"""
|
|
Call FlashQLA's gdn_forward on BI-V100.
|
|
|
|
This is the PROVEN path: compiles and runs without NaN on real hardware.
|
|
"""
|
|
if not _flash_qla_available:
|
|
if not _load_flash_qla():
|
|
raise RuntimeError("FlashQLA GDN not available")
|
|
|
|
if scale is None:
|
|
K = query.shape[-1]
|
|
scale = float(K ** -0.5)
|
|
|
|
output, state = _flash_qla_ext.gdn_forward(
|
|
query, key, value, gate, beta,
|
|
initial_state, scale, output_final_state, head_first
|
|
)
|
|
|
|
return output, state
|
|
|
|
|
|
def gdn_decode_flashqla(
|
|
query: torch.Tensor,
|
|
key: torch.Tensor,
|
|
value: torch.Tensor,
|
|
gate: torch.Tensor,
|
|
beta: torch.Tensor,
|
|
state: torch.Tensor,
|
|
scale: float = None,
|
|
) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
"""
|
|
FlashQLA decode step (single token, update state).
|
|
Uses gdn_decode_mixed_qkv_global_state.
|
|
"""
|
|
if not _flash_qla_available:
|
|
if not _load_flash_qla():
|
|
raise RuntimeError("FlashQLA GDN not available")
|
|
|
|
if scale is None:
|
|
K = query.shape[-1]
|
|
scale = float(K ** -0.5)
|
|
|
|
# FlashQLA decode expects different format — adapt as needed
|
|
output = _flash_qla_ext.gdn_decode_mixed_qkv_global_state(
|
|
query, key, value, gate, beta, state, scale
|
|
)
|
|
|
|
return output, state
|