[AMD] Dsv4/pr2 compressor opt (#26208)
Co-authored-by: wunhuang <wunhuang@amd.com> Co-authored-by: Thomas Wang <1am9trash@gmail.com> Co-authored-by: Xinyi Song <86638975+RolaoDenthu@users.noreply.github.com> Co-authored-by: HaiShaw <hixiao@gmail.com> Co-authored-by: amd-danli103 <danli103@amd.com> Co-authored-by: Lin, Soga <soga.lin@amd.com> Co-authored-by: Raiden-Makoto <Raiden-Makoto@users.noreply.github.com> Co-authored-by: Hubert Lu <55214931+hubertlu-tw@users.noreply.github.com> Co-authored-by: yichiche@amd.com <jacky.cheng> Co-authored-by: yctseng0211 <yctseng@amd.com> Co-authored-by: Bingxu Chen <bingxche@amd.com>
This commit is contained in:
kk
co-authored by
wunhuang
Thomas Wang
Xinyi Song
HaiShaw
amd-danli103
Lin, Soga
Raiden-Makoto
Hubert Lu
yichiche@amd.com
yctseng0211
Bingxu Chen
wunhuang
Thomas Wang
Xinyi Song
HaiShaw
amd-danli103
Lin, Soga
Raiden-Makoto
Hubert Lu
yichiche@amd.com
yctseng0211
Bingxu Chen
parent
7c0fbc8c2e
commit
3f5e2c7688
@@ -627,7 +627,9 @@ class Envs:
|
||||
SGLANG_OPT_USE_TRITON_SWA_PREPARE = EnvBool(True)
|
||||
SGLANG_OPT_USE_AITER_MHC_PRE = EnvBool(True)
|
||||
SGLANG_OPT_USE_AITER_MHC_POST = EnvBool(True)
|
||||
SGLANG_OPT_USE_AITER_SILU_MUL = EnvBool(False)
|
||||
SGLANG_OPT_USE_FUSED_COMPRESS = EnvBool(False)
|
||||
SGLANG_OPT_USE_FUSED_COMPRESS_TRITON = EnvBool(False)
|
||||
SGLANG_OPT_USE_FUSED_QK_NORM_ROPE = EnvBool(True)
|
||||
SGLANG_OPT_USE_FUSED_CLAMP_ACT_MUL = EnvBool(True)
|
||||
SGLANG_FIX_MTP_HC_HIDDEN = EnvBool(False)
|
||||
@@ -644,6 +646,7 @@ class Envs:
|
||||
SGLANG_OPT_USE_TILELANG_MHC_PRE = EnvBool(True)
|
||||
SGLANG_OPT_USE_TILELANG_MHC_POST = EnvBool(True)
|
||||
SGLANG_OPT_USE_TILELANG_INDEXER = EnvBool(False)
|
||||
SGLANG_OPT_USE_AITER_INDEXER = EnvBool(False)
|
||||
SGLANG_OPT_USE_JIT_INDEXER_METADATA = EnvBool(True)
|
||||
SGLANG_OPT_USE_ONLINE_COMPRESS = EnvBool(False)
|
||||
SGLANG_OPT_USE_COMPRESSOR_V2 = EnvBool(True)
|
||||
@@ -688,6 +691,7 @@ class Envs:
|
||||
|
||||
# Cache / overlap
|
||||
SGLANG_OPT_USE_FUSED_STORE_CACHE = EnvBool(True)
|
||||
SGLANG_OPT_USE_JIT_NORM = EnvBool(True)
|
||||
SGLANG_OPT_USE_MULTI_STREAM_OVERLAP = EnvBool(True)
|
||||
|
||||
# CUDA graph
|
||||
|
||||
@@ -33,6 +33,7 @@ from sglang.srt.layers.utils import MultiPlatformOp
|
||||
from sglang.srt.server_args import get_global_server_args
|
||||
from sglang.srt.utils import (
|
||||
cpu_has_amx_support,
|
||||
get_bool_env_var,
|
||||
is_cpu,
|
||||
is_cuda,
|
||||
is_hip,
|
||||
@@ -50,6 +51,7 @@ _is_cpu_amx_available = cpu_has_amx_support()
|
||||
_is_cpu = is_cpu()
|
||||
_is_hip = is_hip()
|
||||
_is_xpu = is_xpu()
|
||||
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
|
||||
|
||||
if _is_cuda:
|
||||
from sglang.jit_kernel.activation import (
|
||||
@@ -71,6 +73,9 @@ elif _is_musa:
|
||||
return torch.empty(output_shape, dtype=x.dtype, device=x.device)
|
||||
|
||||
|
||||
if _use_aiter:
|
||||
from aiter import silu_and_mul as _aiter_silu_and_mul
|
||||
|
||||
if is_npu():
|
||||
import torch_npu
|
||||
|
||||
@@ -82,6 +87,8 @@ class SiluAndMul(MultiPlatformOp):
|
||||
super().__init__(*args, **kwargs)
|
||||
if get_global_server_args().rl_on_policy_target is not None:
|
||||
self._forward_method = self.forward_native
|
||||
elif _use_aiter and envs.SGLANG_OPT_USE_AITER_SILU_MUL.get():
|
||||
self._forward_method = self.forward_aiter
|
||||
|
||||
def forward_native(self, x: torch.Tensor) -> torch.Tensor:
|
||||
d = x.shape[-1] // 2
|
||||
@@ -94,6 +101,13 @@ class SiluAndMul(MultiPlatformOp):
|
||||
silu_and_mul(x, out)
|
||||
return out
|
||||
|
||||
def forward_aiter(self, x: torch.Tensor, limit: float = 0.0) -> torch.Tensor:
|
||||
d = x.shape[-1] // 2
|
||||
output_shape = x.shape[:-1] + (d,)
|
||||
out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
|
||||
_aiter_silu_and_mul(out, x, limit)
|
||||
return out
|
||||
|
||||
def forward_cpu(self, x: torch.Tensor) -> torch.Tensor:
|
||||
if _is_cpu_amx_available:
|
||||
out = torch.ops.sgl_kernel.silu_and_mul_cpu(x)
|
||||
|
||||
@@ -20,11 +20,19 @@ import torch.nn.functional as F
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
|
||||
from sglang.srt.layers.attention.dsv4.compressor import (
|
||||
CompressorBackendMixin,
|
||||
FusedCompressMetadata,
|
||||
create_paged_compressor_data,
|
||||
)
|
||||
|
||||
if envs.SGLANG_OPT_USE_COMPRESSOR_V2.get():
|
||||
from sglang.srt.layers.attention.dsv4.compressor_v2 import (
|
||||
CompressorBackendMixin,
|
||||
FusedCompressMetadata,
|
||||
create_paged_compressor_data,
|
||||
)
|
||||
else:
|
||||
from sglang.srt.layers.attention.dsv4.compressor import (
|
||||
CompressorBackendMixin,
|
||||
FusedCompressMetadata,
|
||||
create_paged_compressor_data,
|
||||
)
|
||||
from sglang.srt.layers.attention.dsv4.indexer import C4IndexerBackendMixin
|
||||
from sglang.srt.layers.attention.dsv4.metadata import (
|
||||
PagedIndexerMetadata,
|
||||
|
||||
@@ -12,9 +12,14 @@ import triton.language as tl
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.layers.attention.dsa.dsa_indexer import rotate_activation
|
||||
from sglang.srt.layers.attention.dsv4.compressor import Compressor as _CompressorBase
|
||||
from sglang.srt.layers.attention.dsv4.fused_compress_triton import (
|
||||
fused_ape_pool_norm_rope,
|
||||
)
|
||||
from sglang.srt.layers.attention.nsa.nsa_indexer import rotate_activation
|
||||
from sglang.srt.layers.deepseek_v4_rope import (
|
||||
apply_rotary_emb_triton,
|
||||
fused_norm_rope_inplace_triton,
|
||||
fused_softmax_pool_triton,
|
||||
)
|
||||
|
||||
try:
|
||||
@@ -372,10 +377,6 @@ class CompressorHip(_CompressorBase):
|
||||
freqs_real_table = self._get_freqs_cis_real()
|
||||
freqs_batch = freqs_real_table[comp_positions]
|
||||
|
||||
from sglang.srt.layers.attention.dsv4.fused_compress_kernel import (
|
||||
fused_ape_pool_norm_rope,
|
||||
)
|
||||
|
||||
kv_compressed = fused_ape_pool_norm_rope(
|
||||
kv_score_gathered=gathered,
|
||||
ape=self.ape,
|
||||
|
||||
@@ -57,6 +57,9 @@ class CompressorBackendMixin:
|
||||
assert isinstance(metadata, FusedCompressMetadata)
|
||||
return metadata
|
||||
|
||||
def _maybe_upgrade_forward_metadata(self) -> None:
|
||||
pass
|
||||
|
||||
def forward_compress(
|
||||
self,
|
||||
*,
|
||||
@@ -91,6 +94,37 @@ class CompressorBackendMixin:
|
||||
metadata = (forward_batch.req_pool_indices.to(torch.int32), None, plan)
|
||||
indices, extra_data, plan = metadata
|
||||
|
||||
if _is_hip:
|
||||
if not is_paged:
|
||||
raise NotImplementedError("HIP fused compressor expects paged metadata")
|
||||
|
||||
from sglang.srt.layers.attention.dsv4.fused_compress_triton import (
|
||||
hip_compress_forward,
|
||||
hip_compress_fused_norm_rope_inplace,
|
||||
)
|
||||
|
||||
kv_compressed = hip_compress_forward(
|
||||
kv_score_buffer=kv_score_buffer,
|
||||
kv_score_input=kv_score_input,
|
||||
ape=ape,
|
||||
indices=indices,
|
||||
plan=plan,
|
||||
compress_ratio=compress_ratio,
|
||||
head_dim=head_dim,
|
||||
extra_data=extra_data,
|
||||
)
|
||||
norm_eps = (
|
||||
norm.variance_epsilon if hasattr(norm, "variance_epsilon") else norm.eps
|
||||
)
|
||||
hip_compress_fused_norm_rope_inplace(
|
||||
kv_compressed,
|
||||
norm.weight,
|
||||
norm_eps,
|
||||
freqs_cis_cache,
|
||||
plan,
|
||||
)
|
||||
return rotate_activation(kv_compressed) if rotate else kv_compressed
|
||||
|
||||
kv_compressed = compress_forward(
|
||||
kv_score_buffer=kv_score_buffer,
|
||||
kv_score_input=kv_score_input,
|
||||
@@ -279,6 +313,8 @@ def create_paged_compressor_data(
|
||||
if is_overlap:
|
||||
write_overlap_loc = get_raw_loc(write_positions - compress_ratio)
|
||||
extra_data = write_overlap_loc.view(-1, 1)
|
||||
elif _is_hip:
|
||||
extra_data = get_raw_loc(write_positions - compress_ratio)
|
||||
else:
|
||||
extra_data = None
|
||||
plan = CompressorDecodePlan(compress_ratio, seq_lens.to(torch.int32))
|
||||
@@ -392,7 +428,7 @@ class Compressor(nn.Module):
|
||||
)
|
||||
|
||||
|
||||
if _is_hip:
|
||||
if _is_hip and not envs.SGLANG_OPT_USE_COMPRESSOR_V2.get():
|
||||
from sglang.srt.layers.attention.dsv4.compress_hip import ( # noqa: F811
|
||||
CompressorHip as Compressor,
|
||||
)
|
||||
|
||||
@@ -10,6 +10,7 @@ from sglang.jit_kernel.dsv4 import (
|
||||
compress_forward,
|
||||
compress_norm_rope_store,
|
||||
)
|
||||
from sglang.jit_kernel.utils import is_hip_runtime
|
||||
from sglang.srt.environ import envs
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -24,12 +25,380 @@ CompressMetadata: TypeAlias = Union[CompressorDecodePlan, CompressorPrefillPlan]
|
||||
# NOTE: alias for backward compatibility
|
||||
FusedCompressMetadata: TypeAlias = CompressMetadata
|
||||
|
||||
_is_hip = is_hip_runtime()
|
||||
|
||||
if _is_hip:
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
@triton.jit
|
||||
def _c128_compress_decode_kernel(
|
||||
buf_ptr,
|
||||
input_ptr,
|
||||
ape_ptr,
|
||||
out_ptr,
|
||||
plan_ptr,
|
||||
buf_stride_slot,
|
||||
input_stride_b,
|
||||
ape_stride_r,
|
||||
out_stride_b,
|
||||
bs,
|
||||
HEAD_DIM: tl.constexpr,
|
||||
BLOCK_D: tl.constexpr,
|
||||
COMPRESS_RATIO: tl.constexpr,
|
||||
):
|
||||
"""Fused C128 decode: write to state buffer + online softmax-pool.
|
||||
|
||||
plan_ptr points to int32 view: [bs, 4] where each row is
|
||||
{seq_len, write_loc, read_page_0, read_page_1}.
|
||||
"""
|
||||
bid = tl.program_id(0)
|
||||
if bid >= bs:
|
||||
return
|
||||
|
||||
# Parse plan
|
||||
plan_base = plan_ptr + bid * 4
|
||||
seq_len = tl.load(plan_base).to(tl.int32)
|
||||
write_loc = tl.load(plan_base + 1).to(tl.int32)
|
||||
read_page_0 = tl.load(plan_base + 2).to(tl.int32)
|
||||
|
||||
d = tl.arange(0, BLOCK_D)
|
||||
last_dim: tl.constexpr = HEAD_DIM * 2
|
||||
|
||||
# Step 1: Write kv_score_input to state buffer at write_loc
|
||||
d_mask_full = d < last_dim
|
||||
input_val = tl.load(
|
||||
input_ptr + bid * input_stride_b + d, mask=d_mask_full, other=0.0
|
||||
)
|
||||
tl.store(buf_ptr + write_loc * buf_stride_slot + d, input_val, mask=d_mask_full)
|
||||
|
||||
# Step 2: Check boundary condition
|
||||
d_mask_hd = d < HEAD_DIM
|
||||
if seq_len % COMPRESS_RATIO != 0:
|
||||
tl.store(
|
||||
out_ptr + bid * out_stride_b + d,
|
||||
tl.zeros([BLOCK_D], tl.float32),
|
||||
mask=d_mask_hd,
|
||||
)
|
||||
return
|
||||
|
||||
# Step 3: Online softmax-pool over 128 slots in the page
|
||||
page_base = read_page_0 * COMPRESS_RATIO * buf_stride_slot
|
||||
m_prev = tl.full([BLOCK_D], float("-inf"), tl.float32)
|
||||
kv_acc = tl.zeros([BLOCK_D], tl.float32)
|
||||
w_acc = tl.zeros([BLOCK_D], tl.float32)
|
||||
|
||||
for k in tl.static_range(COMPRESS_RATIO):
|
||||
slot_addr = page_base + k * buf_stride_slot
|
||||
kv_val = tl.load(buf_ptr + slot_addr + d, mask=d_mask_hd, other=0.0).to(
|
||||
tl.float32
|
||||
)
|
||||
sc_val = tl.load(
|
||||
buf_ptr + slot_addr + HEAD_DIM + d, mask=d_mask_hd, other=0.0
|
||||
).to(tl.float32)
|
||||
ape_val = tl.load(
|
||||
ape_ptr + k * ape_stride_r + d, mask=d_mask_hd, other=0.0
|
||||
).to(tl.float32)
|
||||
score_k = sc_val + ape_val
|
||||
|
||||
m_new = tl.maximum(m_prev, score_k)
|
||||
exp_old = tl.where(m_prev == float("-inf"), 0.0, tl.exp(m_prev - m_new))
|
||||
exp_cur = tl.where(score_k == float("-inf"), 0.0, tl.exp(score_k - m_new))
|
||||
kv_acc = kv_acc * exp_old + exp_cur * kv_val
|
||||
w_acc = w_acc * exp_old + exp_cur
|
||||
m_prev = m_new
|
||||
|
||||
compressed = kv_acc / w_acc
|
||||
tl.store(out_ptr + bid * out_stride_b + d, compressed, mask=d_mask_hd)
|
||||
|
||||
@triton.jit
|
||||
def _c128_compress_prefill_write_kernel(
|
||||
buf_ptr,
|
||||
input_ptr,
|
||||
plan_w_ptr,
|
||||
buf_stride_slot,
|
||||
input_stride_b,
|
||||
num_w,
|
||||
BLOCK_D: tl.constexpr,
|
||||
LAST_DIM: tl.constexpr,
|
||||
):
|
||||
"""Prefill write phase: scatter kv_score_input tokens into state buffer."""
|
||||
wid = tl.program_id(0)
|
||||
if wid >= num_w:
|
||||
return
|
||||
|
||||
# WritePlan: {ragged_id(u32), write_loc(i32)} = 8 bytes = 2 int32s
|
||||
plan_base = plan_w_ptr + wid * 2
|
||||
ragged_id = (tl.load(plan_base).to(tl.int32)) & 0xFFFF
|
||||
write_loc = tl.load(plan_base + 1).to(tl.int32)
|
||||
|
||||
d = tl.arange(0, BLOCK_D)
|
||||
d_mask = d < LAST_DIM
|
||||
|
||||
if write_loc >= 0:
|
||||
input_val = tl.load(
|
||||
input_ptr + ragged_id * input_stride_b + d, mask=d_mask, other=0.0
|
||||
)
|
||||
tl.store(buf_ptr + write_loc * buf_stride_slot + d, input_val, mask=d_mask)
|
||||
|
||||
@triton.jit
|
||||
def _c128_compress_prefill_compress_kernel(
|
||||
buf_ptr,
|
||||
ape_ptr,
|
||||
out_ptr,
|
||||
plan_c_ptr,
|
||||
buf_stride_slot,
|
||||
ape_stride_r,
|
||||
out_stride_b,
|
||||
num_c,
|
||||
HEAD_DIM: tl.constexpr,
|
||||
BLOCK_D: tl.constexpr,
|
||||
COMPRESS_RATIO: tl.constexpr,
|
||||
):
|
||||
"""Prefill compress phase: online softmax-pool for each compress plan entry."""
|
||||
cid = tl.program_id(0)
|
||||
if cid >= num_c:
|
||||
return
|
||||
|
||||
# CompressPlan: {seq_len(u32), ragged_id(u16)|buffer_len(u16), read_page_0(i32), read_page_1(i32)}
|
||||
plan_base = plan_c_ptr + cid * 4
|
||||
read_page_0 = tl.load(plan_base + 2).to(tl.int32)
|
||||
|
||||
d = tl.arange(0, BLOCK_D)
|
||||
d_mask_hd = d < HEAD_DIM
|
||||
|
||||
if read_page_0 < 0:
|
||||
tl.store(
|
||||
out_ptr + cid * out_stride_b + d,
|
||||
tl.zeros([BLOCK_D], tl.float32),
|
||||
mask=d_mask_hd,
|
||||
)
|
||||
return
|
||||
|
||||
page_base = read_page_0 * COMPRESS_RATIO * buf_stride_slot
|
||||
m_prev = tl.full([BLOCK_D], float("-inf"), tl.float32)
|
||||
kv_acc = tl.zeros([BLOCK_D], tl.float32)
|
||||
w_acc = tl.zeros([BLOCK_D], tl.float32)
|
||||
|
||||
for k in tl.static_range(COMPRESS_RATIO):
|
||||
slot_addr = page_base + k * buf_stride_slot
|
||||
kv_val = tl.load(buf_ptr + slot_addr + d, mask=d_mask_hd, other=0.0).to(
|
||||
tl.float32
|
||||
)
|
||||
sc_val = tl.load(
|
||||
buf_ptr + slot_addr + HEAD_DIM + d, mask=d_mask_hd, other=0.0
|
||||
).to(tl.float32)
|
||||
ape_val = tl.load(
|
||||
ape_ptr + k * ape_stride_r + d, mask=d_mask_hd, other=0.0
|
||||
).to(tl.float32)
|
||||
score_k = sc_val + ape_val
|
||||
|
||||
m_new = tl.maximum(m_prev, score_k)
|
||||
exp_old = tl.where(m_prev == float("-inf"), 0.0, tl.exp(m_prev - m_new))
|
||||
exp_cur = tl.where(score_k == float("-inf"), 0.0, tl.exp(score_k - m_new))
|
||||
kv_acc = kv_acc * exp_old + exp_cur * kv_val
|
||||
w_acc = w_acc * exp_old + exp_cur
|
||||
m_prev = m_new
|
||||
|
||||
compressed = kv_acc / w_acc
|
||||
tl.store(out_ptr + cid * out_stride_b + d, compressed, mask=d_mask_hd)
|
||||
|
||||
|
||||
def _compress_forward_c128_triton(
|
||||
kv_score_buffer: torch.Tensor,
|
||||
kv_score_input: torch.Tensor,
|
||||
ape: torch.Tensor,
|
||||
plan: Union[CompressorDecodePlan, CompressorPrefillPlan],
|
||||
head_dim: int,
|
||||
) -> torch.Tensor:
|
||||
"""Triton C128 compress_forward for HIP (wave64).
|
||||
|
||||
Fuses write + online-softmax-pool into Triton kernels.
|
||||
CUDA graph compatible.
|
||||
"""
|
||||
num_total_slots = kv_score_buffer.shape[0] * kv_score_buffer.shape[1]
|
||||
num_pages = kv_score_buffer.shape[0]
|
||||
last_dim = kv_score_buffer.shape[-1]
|
||||
compress_ratio = 128
|
||||
|
||||
buf_flat = kv_score_buffer.view(-1, last_dim)
|
||||
buf_stride_slot = last_dim # elements per slot
|
||||
|
||||
BLOCK_D = triton.next_power_of_2(last_dim)
|
||||
|
||||
if plan.is_decode:
|
||||
# Decode path: single kernel does write + compress
|
||||
plan_raw = plan[1].view(torch.int32) # [bs, 4]
|
||||
bs = plan_raw.shape[0]
|
||||
out = torch.empty(
|
||||
bs, head_dim, dtype=torch.float32, device=kv_score_input.device
|
||||
)
|
||||
|
||||
if bs > 0 and num_total_slots > 0:
|
||||
grid = (bs,)
|
||||
_c128_compress_decode_kernel[grid](
|
||||
buf_flat,
|
||||
kv_score_input,
|
||||
ape,
|
||||
out,
|
||||
plan_raw,
|
||||
buf_stride_slot,
|
||||
kv_score_input.stride(0),
|
||||
ape.stride(0),
|
||||
out.stride(0),
|
||||
bs,
|
||||
HEAD_DIM=head_dim,
|
||||
BLOCK_D=triton.next_power_of_2(head_dim),
|
||||
COMPRESS_RATIO=compress_ratio,
|
||||
num_warps=8,
|
||||
)
|
||||
return out
|
||||
else:
|
||||
# Prefill path: separate write kernel + compress kernel
|
||||
plan_c_raw = plan[1].view(torch.int32) # [num_c, 4]
|
||||
plan_w = plan[2] # [num_w, 8] uint8
|
||||
plan_w_raw = plan_w.view(torch.int32) # [num_w, 2]
|
||||
num_c = plan_c_raw.shape[0]
|
||||
num_w = plan_w_raw.shape[0]
|
||||
|
||||
out = torch.empty(
|
||||
num_c, head_dim, dtype=torch.float32, device=kv_score_input.device
|
||||
)
|
||||
|
||||
# Phase 1: Write
|
||||
if num_w > 0 and num_total_slots > 0:
|
||||
grid_w = (num_w,)
|
||||
_c128_compress_prefill_write_kernel[grid_w](
|
||||
buf_flat,
|
||||
kv_score_input,
|
||||
plan_w_raw,
|
||||
buf_stride_slot,
|
||||
kv_score_input.stride(0),
|
||||
num_w,
|
||||
BLOCK_D=BLOCK_D,
|
||||
LAST_DIM=last_dim,
|
||||
num_warps=4,
|
||||
)
|
||||
|
||||
# Phase 2: Compress
|
||||
if num_c > 0 and num_pages > 0:
|
||||
grid_c = (num_c,)
|
||||
_c128_compress_prefill_compress_kernel[grid_c](
|
||||
buf_flat,
|
||||
ape,
|
||||
out,
|
||||
plan_c_raw,
|
||||
buf_stride_slot,
|
||||
ape.stride(0),
|
||||
out.stride(0),
|
||||
num_c,
|
||||
HEAD_DIM=head_dim,
|
||||
BLOCK_D=triton.next_power_of_2(head_dim),
|
||||
COMPRESS_RATIO=compress_ratio,
|
||||
num_warps=8,
|
||||
)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
def _use_online_compress(compress_ratio: int) -> bool:
|
||||
"""Online state-pool path is c128-only."""
|
||||
return compress_ratio == 128 and envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get()
|
||||
|
||||
|
||||
def _extract_positions_from_plan(
|
||||
plan: Union[CompressorDecodePlan, CompressorPrefillPlan],
|
||||
compress_ratio: int,
|
||||
) -> torch.Tensor:
|
||||
"""Extract RoPE positions from plan tensors (decode or prefill).
|
||||
|
||||
DecodePlan layout: [bs, 16] uint8, first 4 bytes = uint32 seq_len.
|
||||
CompressPlan layout: [num_c, 16] uint8, first 4 bytes = uint32 seq_len.
|
||||
Position for RoPE = seq_len - compress_ratio.
|
||||
"""
|
||||
plan_tensor = plan[1] # plan_d or plan_c
|
||||
seq_lens = plan_tensor[:, :4].contiguous().view(torch.int32).squeeze(-1)
|
||||
positions = seq_lens.to(torch.int32) - compress_ratio
|
||||
return positions
|
||||
|
||||
|
||||
def _compress_forward_c128_fallback(
|
||||
kv_score_buffer: torch.Tensor,
|
||||
kv_score_input: torch.Tensor,
|
||||
ape: torch.Tensor,
|
||||
plan: Union[CompressorDecodePlan, CompressorPrefillPlan],
|
||||
head_dim: int,
|
||||
) -> torch.Tensor:
|
||||
"""PyTorch fallback for C128 compress_forward on HIP (wave64).
|
||||
|
||||
Fully vectorized, compatible with CUDA graph capture.
|
||||
kv_score_buffer: [num_pages, 128, head_dim * 2]
|
||||
ape: [128, head_dim]
|
||||
|
||||
IMPORTANT: This also performs the write to state buffer (like the JIT kernel).
|
||||
The JIT kernel does: (1) write kv_score_input to buffer, (2) compress from buffer.
|
||||
"""
|
||||
num_total_slots = kv_score_buffer.shape[0] * kv_score_buffer.shape[1]
|
||||
num_pages = kv_score_buffer.shape[0]
|
||||
last_dim = kv_score_buffer.shape[-1]
|
||||
|
||||
# Step 1: WRITE kv_score_input to state buffer
|
||||
if num_total_slots > 0:
|
||||
buf_flat = kv_score_buffer.view(-1, last_dim)
|
||||
if plan.is_decode:
|
||||
# Decode: plan_d has write_loc per batch item
|
||||
plan_raw = plan[1].view(torch.int32) # [bs, 4]
|
||||
write_locs = plan_raw[:, 1].long()
|
||||
# Only write valid locations (>= 0 and < buffer size)
|
||||
valid_write = (write_locs >= 0) & (write_locs < num_total_slots)
|
||||
if valid_write.any():
|
||||
buf_flat[write_locs[valid_write]] = kv_score_input[valid_write]
|
||||
else:
|
||||
# Prefill: plan_w has {ragged_id, write_loc} per write entry
|
||||
plan_w = plan[2] # [num_w, 8] uint8 = WritePlan
|
||||
if plan_w.shape[0] > 0:
|
||||
plan_w_raw = plan_w.view(torch.int32) # [num_w, 2]
|
||||
ragged_ids = plan_w_raw[:, 0].long() & 0xFFFF
|
||||
write_locs = plan_w_raw[:, 1].long()
|
||||
valid_write = (write_locs >= 0) & (write_locs < num_total_slots)
|
||||
ragged_ids_safe = ragged_ids.clamp(
|
||||
min=0, max=kv_score_input.shape[0] - 1
|
||||
)
|
||||
if valid_write.any():
|
||||
buf_flat[write_locs[valid_write]] = kv_score_input[
|
||||
ragged_ids_safe[valid_write]
|
||||
]
|
||||
|
||||
# Step 2: COMPRESS (read from buffer page and do softmax-pool)
|
||||
plan_c = plan[1] # plan_d for decode, plan_c for prefill
|
||||
num_tokens = plan_c.shape[0]
|
||||
if num_pages == 0 or num_tokens == 0:
|
||||
return kv_score_input.new_zeros(num_tokens, head_dim)
|
||||
|
||||
plan_c_raw = plan_c.view(torch.int32) # [N, 4]
|
||||
read_page_0 = plan_c_raw[:, 2].long()
|
||||
# Use torch.where instead of clamp to handle -1 (invalid) gracefully
|
||||
valid_read = (read_page_0 >= 0) & (read_page_0 < num_pages)
|
||||
read_page_0_safe = torch.where(
|
||||
valid_read, read_page_0, torch.zeros_like(read_page_0)
|
||||
)
|
||||
|
||||
gathered = kv_score_buffer[read_page_0_safe] # [N, 128, head_dim*2]
|
||||
kv = gathered[:, :, :head_dim].float()
|
||||
score = gathered[:, :, head_dim:].float() + ape.float().unsqueeze(0)
|
||||
weights = score.softmax(dim=1)
|
||||
out = (weights * kv).sum(dim=1)
|
||||
|
||||
# For decode: zero out non-boundary tokens (seq_len % 128 != 0)
|
||||
# so they don't corrupt kvcache location 0 when stored.
|
||||
if plan.is_decode:
|
||||
seq_lens = plan_c_raw[:, 0].to(torch.int32)
|
||||
is_boundary = (seq_lens % 128 == 0).unsqueeze(-1) # [N, 1]
|
||||
out = torch.where(is_boundary, out, torch.zeros_like(out))
|
||||
|
||||
return out.to(kv_score_input.dtype)
|
||||
|
||||
|
||||
class CompressorBackendMixin:
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
@@ -74,6 +443,8 @@ class CompressorBackendMixin:
|
||||
last_dim = 2 * head_dim * coff
|
||||
assert kv_score_buffer.shape[-1] == last_dim
|
||||
kv_score_buffer = kv_score_buffer.view(-1, compress_ratio, last_dim)
|
||||
|
||||
# Step 1: compress_forward
|
||||
kv_compressed = compress_forward(
|
||||
kv_score_buffer=kv_score_buffer,
|
||||
kv_score_input=kv_score_input,
|
||||
@@ -83,7 +454,8 @@ class CompressorBackendMixin:
|
||||
head_dim=head_dim,
|
||||
is_online=is_online,
|
||||
)
|
||||
# NOTE: we use some hack here...
|
||||
|
||||
# Step 2: norm + rope + store
|
||||
compress_norm_rope_store(
|
||||
kv_compressed,
|
||||
plan,
|
||||
@@ -109,34 +481,153 @@ class CompressorBackendMixin:
|
||||
token_to_kv_pool = self.token_to_kv_pool
|
||||
token_to_kv_pool = cast("DeepSeekV4TokenToKVPool", token_to_kv_pool)
|
||||
kv_score_input = compressor.compute_kv_score(x, forward_batch)
|
||||
|
||||
state_pool = compressor.get_state_pool(self)
|
||||
out_loc = self._get_out_loc(compressor.ratio)
|
||||
if compressor.is_in_indexer:
|
||||
kv_cache = token_to_kv_pool.get_index_k_with_scale_buffer(layer_id)
|
||||
page_size = token_to_kv_pool.get_index_k_page_size()
|
||||
if _is_hip and not envs.SGLANG_OPT_USE_JIT_NORM.get():
|
||||
self._forward_unified_hip(
|
||||
token_to_kv_pool=token_to_kv_pool,
|
||||
kv_score_input=kv_score_input,
|
||||
state_pool=state_pool,
|
||||
compressor=compressor,
|
||||
layer_id=layer_id,
|
||||
)
|
||||
else:
|
||||
_, _, compress_kv_pool = token_to_kv_pool.layer_mapping[layer_id]
|
||||
assert compress_kv_pool is not None
|
||||
kv_cache = token_to_kv_pool.get_extra_key_buffer(layer_id)
|
||||
page_size = token_to_kv_pool.get_extra_key_page_size(layer_id)
|
||||
if hasattr(compress_kv_pool, "translate_loc_to_hisparse_device"):
|
||||
# The v2 compressor writes directly into the raw C4 KV tensor.
|
||||
# HiSparse C4 therefore needs the physical C4 location here.
|
||||
out_loc = compress_kv_pool.translate_loc_to_hisparse_device(out_loc)
|
||||
self._forward_compress_all_in_one(
|
||||
kv_score_buffer=state_pool.kv_score_buffer.kv_score,
|
||||
kv_score_input=kv_score_input,
|
||||
ape=compressor.ape,
|
||||
head_dim=compressor.head_dim,
|
||||
norm=compressor.norm,
|
||||
freqs_cis_cache=compressor.freqs_cis,
|
||||
kv_cache=kv_cache.view(dtype=torch.uint8),
|
||||
is_indexer=compressor.is_in_indexer,
|
||||
rotate=compressor.rotate,
|
||||
compress_ratio=compressor.ratio,
|
||||
page_size=page_size,
|
||||
out_loc=out_loc,
|
||||
out_loc = self._get_out_loc(compressor.ratio)
|
||||
if compressor.is_in_indexer:
|
||||
kv_cache = token_to_kv_pool.get_index_k_with_scale_buffer(layer_id)
|
||||
page_size = token_to_kv_pool.get_index_k_page_size()
|
||||
else:
|
||||
_, _, compress_kv_pool = token_to_kv_pool.layer_mapping[layer_id]
|
||||
assert compress_kv_pool is not None
|
||||
kv_cache = token_to_kv_pool.get_extra_key_buffer(layer_id)
|
||||
page_size = token_to_kv_pool.get_extra_key_page_size(layer_id)
|
||||
if hasattr(compress_kv_pool, "translate_loc_to_hisparse_device"):
|
||||
# The v2 compressor writes directly into the raw C4 KV tensor.
|
||||
# HiSparse C4 therefore needs the physical C4 location here.
|
||||
out_loc = compress_kv_pool.translate_loc_to_hisparse_device(out_loc)
|
||||
self._forward_compress_all_in_one(
|
||||
kv_score_buffer=state_pool.kv_score_buffer.kv_score,
|
||||
kv_score_input=kv_score_input,
|
||||
ape=compressor.ape,
|
||||
head_dim=compressor.head_dim,
|
||||
norm=compressor.norm,
|
||||
freqs_cis_cache=compressor.freqs_cis,
|
||||
kv_cache=kv_cache.view(dtype=torch.uint8),
|
||||
is_indexer=compressor.is_in_indexer,
|
||||
rotate=compressor.rotate,
|
||||
compress_ratio=compressor.ratio,
|
||||
page_size=page_size,
|
||||
out_loc=out_loc,
|
||||
)
|
||||
|
||||
def _forward_unified_hip(
|
||||
self,
|
||||
token_to_kv_pool: DeepSeekV4TokenToKVPool,
|
||||
kv_score_input: torch.Tensor,
|
||||
state_pool,
|
||||
compressor: Compressor,
|
||||
layer_id: int,
|
||||
) -> None:
|
||||
"""HIP-specific forward path using PyTorch/Triton fallbacks."""
|
||||
from sglang.srt.layers.attention.dsv4.quant_k_cache import (
|
||||
quant_to_nope_fp8_rope_bf16_pack_triton,
|
||||
)
|
||||
from sglang.srt.layers.attention.nsa.nsa_indexer import rotate_activation
|
||||
from sglang.srt.layers.attention.nsa.triton_kernel import act_quant
|
||||
from sglang.srt.layers.deepseek_v4_rope import fused_norm_rope_inplace_triton
|
||||
|
||||
compress_ratio = compressor.ratio
|
||||
head_dim = compressor.head_dim
|
||||
is_indexer = compressor.is_in_indexer
|
||||
|
||||
plan = self._get_paged_compress_metadata(compress_ratio)
|
||||
out_loc = self._get_out_loc(compress_ratio)
|
||||
|
||||
# Step 1: compress_forward (always use JIT for both C4 and C128)
|
||||
coff = 2 if is_overlap_compress(compress_ratio) else 1
|
||||
last_dim = 2 * head_dim * coff
|
||||
kv_score_buffer = state_pool.kv_score_buffer.kv_score
|
||||
kv_score_buffer = kv_score_buffer.view(-1, compress_ratio, last_dim)
|
||||
|
||||
kv_compressed = compress_forward(
|
||||
kv_score_buffer=kv_score_buffer,
|
||||
kv_score_input=kv_score_input,
|
||||
ape=compressor.ape.view(-1, head_dim),
|
||||
plan=plan,
|
||||
compress_ratio=compress_ratio,
|
||||
head_dim=head_dim,
|
||||
is_online=False,
|
||||
)
|
||||
|
||||
if kv_compressed.shape[0] == 0:
|
||||
return
|
||||
|
||||
# For decode: zero out non-boundary tokens to prevent corrupting kvcache loc 0.
|
||||
if plan.is_decode:
|
||||
plan_raw = plan[1].view(torch.int32)
|
||||
seq_lens_plan = plan_raw[:, 0].to(torch.int32)
|
||||
is_boundary = (seq_lens_plan % compress_ratio == 0).unsqueeze(-1)
|
||||
kv_compressed = torch.where(
|
||||
is_boundary, kv_compressed, torch.zeros_like(kv_compressed)
|
||||
)
|
||||
|
||||
# Step 2: norm + rope (Triton fallback for precision parity with V1)
|
||||
positions = _extract_positions_from_plan(plan, compress_ratio)
|
||||
positions_safe = positions.clamp(min=0)
|
||||
|
||||
fused_norm_rope_inplace_triton(
|
||||
kv_compressed,
|
||||
compressor.norm.weight,
|
||||
compressor.norm.variance_epsilon,
|
||||
compressor.freqs_cis,
|
||||
positions=positions_safe,
|
||||
)
|
||||
|
||||
# Step 3: optional Hadamard rotation for indexer
|
||||
if compressor.rotate:
|
||||
kv_compressed = rotate_activation(kv_compressed)
|
||||
|
||||
# Step 4: store to kvcache
|
||||
# For decode: store ALL tokens. Non-boundary tokens have out_loc=0 (safe).
|
||||
# For prefill: plan_c already only contains valid entries.
|
||||
if plan.is_decode:
|
||||
kv_to_store = kv_compressed
|
||||
out_loc_to_store = out_loc
|
||||
else:
|
||||
kv_to_store = kv_compressed
|
||||
plan_raw = plan[1].view(torch.int32)
|
||||
ragged_ids = plan_raw[:, 1].to(torch.int32) & 0xFFFF
|
||||
out_loc_to_store = out_loc[ragged_ids.long()]
|
||||
|
||||
if kv_to_store.shape[0] == 0:
|
||||
return
|
||||
|
||||
if envs.SGLANG_OPT_USE_FUSED_STORE_CACHE.get():
|
||||
# fused kernel: BF16 in -> FP8 quant + paged scatter in one launch
|
||||
if is_indexer:
|
||||
token_to_kv_pool.set_index_k_fused(
|
||||
layer_id=layer_id,
|
||||
loc=out_loc_to_store,
|
||||
cache_k=kv_to_store,
|
||||
)
|
||||
else:
|
||||
token_to_kv_pool.set_extra_key_buffer_fused(
|
||||
layer_id=layer_id,
|
||||
loc=out_loc_to_store,
|
||||
cache_k=kv_to_store,
|
||||
)
|
||||
else:
|
||||
if is_indexer:
|
||||
kv_fp8, kv_scale = act_quant(kv_to_store)
|
||||
token_to_kv_pool.set_index_k_scale_buffer(
|
||||
layer_id=layer_id,
|
||||
loc=out_loc_to_store,
|
||||
index_k=kv_fp8,
|
||||
index_k_scale=kv_scale,
|
||||
)
|
||||
else:
|
||||
pack = quant_to_nope_fp8_rope_bf16_pack_triton(kv_to_store.bfloat16())
|
||||
token_to_kv_pool.set_extra_key_buffer(layer_id, out_loc_to_store, pack)
|
||||
|
||||
# NOTE: alias for backward compatibility
|
||||
forward_indexer_compressor = forward_unified
|
||||
|
||||
@@ -0,0 +1,954 @@
|
||||
"""HIP fused compressor kernels using the NV/main metadata contract.
|
||||
|
||||
The public wrappers mirror ``compress_forward``:
|
||||
|
||||
decode: indices, seq_lens, extra_data
|
||||
prefill: indices, compress_plan, write_plan, extra_data
|
||||
|
||||
Prefill plans are the upstream 16-byte ``PrefillPlan`` structs stored as
|
||||
``uint8[:, 16]``. The wrappers reinterpret them as ``int32[:, 4]`` before
|
||||
launching Triton kernels.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Optional, Union
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
from sglang.jit_kernel.dsv4.compress_old import (
|
||||
CompressorDecodePlan,
|
||||
CompressorPrefillPlan,
|
||||
)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _fused_ape_pool_norm_rope_kernel(
|
||||
kv_score_ptr,
|
||||
kv_score_stride_b,
|
||||
kv_score_stride_k,
|
||||
ape_ptr,
|
||||
ape_stride_r,
|
||||
rms_weight_ptr,
|
||||
rms_eps,
|
||||
freqs_ptr,
|
||||
freqs_stride_b,
|
||||
out_ptr,
|
||||
out_stride_b,
|
||||
head_dim,
|
||||
rope_head_dim,
|
||||
half_dim,
|
||||
RATIO: tl.constexpr,
|
||||
K_POOL: tl.constexpr,
|
||||
BLOCK_D: tl.constexpr,
|
||||
HALF_ROPE: tl.constexpr,
|
||||
OVERLAP: tl.constexpr,
|
||||
):
|
||||
bid = tl.program_id(0)
|
||||
d = tl.arange(0, BLOCK_D)
|
||||
d_mask = d < head_dim
|
||||
|
||||
m_prev = tl.full([BLOCK_D], float("-inf"), tl.float32)
|
||||
kv_acc = tl.zeros([BLOCK_D], tl.float32)
|
||||
w_acc = tl.zeros([BLOCK_D], tl.float32)
|
||||
|
||||
batch_base = bid * kv_score_stride_b
|
||||
|
||||
for k in tl.range(0, K_POOL):
|
||||
if OVERLAP:
|
||||
is_b = k >= RATIO
|
||||
col_off = tl.where(is_b, head_dim, 0)
|
||||
else:
|
||||
col_off = 0
|
||||
|
||||
row_off = batch_base + k * kv_score_stride_k
|
||||
kv_val = tl.load(
|
||||
kv_score_ptr + row_off + col_off + d, mask=d_mask, other=0.0
|
||||
).to(tl.float32)
|
||||
sc_val = tl.load(
|
||||
kv_score_ptr + row_off + half_dim + col_off + d, mask=d_mask, other=0.0
|
||||
).to(tl.float32)
|
||||
|
||||
ape_val = tl.load(
|
||||
ape_ptr + (k % RATIO) * ape_stride_r + col_off + d, mask=d_mask, other=0.0
|
||||
).to(tl.float32)
|
||||
score_k = sc_val + ape_val
|
||||
|
||||
m_new = tl.maximum(m_prev, score_k)
|
||||
exp_old = tl.where(m_prev == float("-inf"), 0.0, tl.exp(m_prev - m_new))
|
||||
exp_cur = tl.where(score_k == float("-inf"), 0.0, tl.exp(score_k - m_new))
|
||||
kv_acc = kv_acc * exp_old + exp_cur * kv_val
|
||||
w_acc = w_acc * exp_old + exp_cur
|
||||
m_prev = m_new
|
||||
|
||||
compressed = kv_acc / w_acc
|
||||
rms_w = tl.load(rms_weight_ptr + d, mask=d_mask, other=0.0)
|
||||
c_sq = tl.where(d_mask, compressed * compressed, 0.0)
|
||||
var = tl.sum(c_sq, axis=0) / head_dim
|
||||
normed = compressed * tl.rsqrt(var + rms_eps) * rms_w
|
||||
|
||||
out_base = out_ptr + bid * out_stride_b
|
||||
tl.store(out_base + d, normed.to(out_ptr.dtype.element_ty), mask=d_mask)
|
||||
|
||||
rope_start = head_dim - rope_head_dim
|
||||
p = tl.arange(0, HALF_ROPE)
|
||||
pmask = p < (rope_head_dim // 2)
|
||||
xr = tl.load(out_base + rope_start + 2 * p, mask=pmask, other=0.0).to(tl.float32)
|
||||
xi = tl.load(out_base + rope_start + 2 * p + 1, mask=pmask, other=0.0).to(
|
||||
tl.float32
|
||||
)
|
||||
|
||||
freq_base = bid * freqs_stride_b
|
||||
fr = tl.load(freqs_ptr + freq_base + 2 * p, mask=pmask, other=1.0).to(tl.float32)
|
||||
fi = tl.load(freqs_ptr + freq_base + 2 * p + 1, mask=pmask, other=0.0).to(
|
||||
tl.float32
|
||||
)
|
||||
|
||||
tl.store(
|
||||
out_base + rope_start + 2 * p,
|
||||
(xr * fr - xi * fi).to(out_ptr.dtype.element_ty),
|
||||
mask=pmask,
|
||||
)
|
||||
tl.store(
|
||||
out_base + rope_start + 2 * p + 1,
|
||||
(xr * fi + xi * fr).to(out_ptr.dtype.element_ty),
|
||||
mask=pmask,
|
||||
)
|
||||
|
||||
|
||||
def fused_ape_pool_norm_rope(
|
||||
kv_score_gathered: torch.Tensor,
|
||||
ape: torch.Tensor,
|
||||
rms_weight: torch.Tensor,
|
||||
rms_eps: float,
|
||||
freqs_cis_real: torch.Tensor,
|
||||
head_dim: int,
|
||||
rope_head_dim: int,
|
||||
ratio: int,
|
||||
overlap: bool,
|
||||
) -> torch.Tensor:
|
||||
"""Fused APE-add + overlap-transform + softmax-pool + RMSNorm + RoPE."""
|
||||
coff = 2 if overlap else 1
|
||||
bs = kv_score_gathered.shape[0]
|
||||
k_in = kv_score_gathered.shape[1]
|
||||
last_dim = kv_score_gathered.shape[2]
|
||||
half_dim = last_dim // 2
|
||||
assert k_in == ratio * coff, f"k_in={k_in} != ratio*coff={ratio}*{coff}"
|
||||
|
||||
out = torch.empty(
|
||||
bs, head_dim, dtype=torch.float32, device=kv_score_gathered.device
|
||||
)
|
||||
if bs == 0:
|
||||
return out
|
||||
|
||||
block_d = triton.next_power_of_2(head_dim)
|
||||
half_rope = triton.next_power_of_2(rope_head_dim // 2)
|
||||
num_warps = 4 if head_dim <= 256 else 8
|
||||
|
||||
_fused_ape_pool_norm_rope_kernel[(bs,)](
|
||||
kv_score_gathered,
|
||||
kv_score_gathered.stride(0),
|
||||
kv_score_gathered.stride(1),
|
||||
ape,
|
||||
ape.stride(0),
|
||||
rms_weight,
|
||||
rms_eps,
|
||||
freqs_cis_real,
|
||||
freqs_cis_real.stride(0),
|
||||
out,
|
||||
out.stride(0),
|
||||
head_dim,
|
||||
rope_head_dim,
|
||||
half_dim,
|
||||
RATIO=ratio,
|
||||
K_POOL=k_in,
|
||||
BLOCK_D=block_d,
|
||||
HALF_ROPE=half_rope,
|
||||
OVERLAP=int(overlap),
|
||||
num_warps=num_warps,
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _c4_decode_kernel(
|
||||
kv_in_ptr,
|
||||
out_ptr,
|
||||
buffer_ptr,
|
||||
ape_ptr,
|
||||
indices_ptr,
|
||||
seq_lens_ptr,
|
||||
extra_ptr,
|
||||
kv_in_row_stride,
|
||||
out_row_stride,
|
||||
buffer_page_stride,
|
||||
buffer_slot_stride,
|
||||
ape_row_stride,
|
||||
HEAD_DIM: tl.constexpr,
|
||||
BLOCK_D: tl.constexpr,
|
||||
):
|
||||
bid = tl.program_id(0)
|
||||
pid_d = tl.program_id(1)
|
||||
d_offs = pid_d * BLOCK_D + tl.arange(0, BLOCK_D)
|
||||
d_mask = d_offs < HEAD_DIM
|
||||
|
||||
index = tl.load(indices_ptr + bid).to(tl.int64)
|
||||
index_prev = tl.load(extra_ptr + bid).to(tl.int64)
|
||||
seq_len = tl.load(seq_lens_ptr + bid).to(tl.int32)
|
||||
write_slot = (seq_len + 3) % 4
|
||||
|
||||
in_base = bid.to(tl.int64) * kv_in_row_stride
|
||||
page_base = (
|
||||
index * buffer_page_stride + write_slot.to(tl.int64) * buffer_slot_stride
|
||||
)
|
||||
|
||||
valid_index = index >= 0
|
||||
for ch in tl.static_range(4):
|
||||
ch_off = ch * HEAD_DIM
|
||||
val = tl.load(kv_in_ptr + in_base + ch_off + d_offs, mask=d_mask, other=0.0)
|
||||
tl.store(
|
||||
buffer_ptr + page_base + ch_off + d_offs,
|
||||
val,
|
||||
mask=d_mask & valid_index,
|
||||
)
|
||||
|
||||
NEG_BIG: tl.constexpr = -1.0e9
|
||||
running_max = tl.full((BLOCK_D,), NEG_BIG, tl.float32)
|
||||
running_sum = tl.zeros((BLOCK_D,), tl.float32)
|
||||
weighted = tl.zeros((BLOCK_D,), tl.float32)
|
||||
|
||||
for slot in tl.static_range(8):
|
||||
if slot < 4:
|
||||
page = index_prev
|
||||
kv_off = 0
|
||||
score_off = 2 * HEAD_DIM
|
||||
else:
|
||||
page = index
|
||||
kv_off = HEAD_DIM
|
||||
score_off = 3 * HEAD_DIM
|
||||
|
||||
src_pos = seq_len - 8 + slot
|
||||
is_input = slot == 7
|
||||
write_pos = ((seq_len - 1) // 4) * 4
|
||||
page = tl.where(src_pos < write_pos, index_prev, index)
|
||||
slot_in_page = src_pos % 4
|
||||
slot_base = (
|
||||
page * buffer_page_stride + slot_in_page.to(tl.int64) * buffer_slot_stride
|
||||
)
|
||||
valid = src_pos >= 0
|
||||
if slot == 7:
|
||||
kv = tl.load(
|
||||
kv_in_ptr + in_base + kv_off + d_offs,
|
||||
mask=d_mask & valid,
|
||||
other=0.0,
|
||||
)
|
||||
score = tl.load(
|
||||
kv_in_ptr + in_base + score_off + d_offs,
|
||||
mask=d_mask & valid,
|
||||
other=NEG_BIG,
|
||||
)
|
||||
else:
|
||||
kv = tl.load(
|
||||
buffer_ptr + slot_base + kv_off + d_offs,
|
||||
mask=d_mask & valid,
|
||||
other=0.0,
|
||||
)
|
||||
score = tl.load(
|
||||
buffer_ptr + slot_base + score_off + d_offs,
|
||||
mask=d_mask & valid,
|
||||
other=NEG_BIG,
|
||||
)
|
||||
bias = tl.load(ape_ptr + slot * ape_row_stride + d_offs, mask=d_mask, other=0.0)
|
||||
s = score + bias
|
||||
new_max = tl.maximum(running_max, s)
|
||||
factor = tl.exp(running_max - new_max)
|
||||
e = tl.where(valid, tl.exp(s - new_max), 0.0)
|
||||
running_sum = running_sum * factor + e
|
||||
weighted = weighted * factor + kv * e
|
||||
running_max = new_max
|
||||
|
||||
tl.store(
|
||||
out_ptr + bid.to(tl.int64) * out_row_stride + d_offs,
|
||||
weighted / running_sum,
|
||||
mask=d_mask,
|
||||
)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _c4_prefill_compress_kernel(
|
||||
kv_in_ptr,
|
||||
out_ptr,
|
||||
buffer_ptr,
|
||||
ape_ptr,
|
||||
indices_ptr,
|
||||
extra_ptr,
|
||||
plan_ptr,
|
||||
kv_in_row_stride,
|
||||
out_row_stride,
|
||||
buffer_page_stride,
|
||||
buffer_slot_stride,
|
||||
ape_row_stride,
|
||||
plan_row_stride,
|
||||
HEAD_DIM: tl.constexpr,
|
||||
BLOCK_D: tl.constexpr,
|
||||
):
|
||||
pid_p = tl.program_id(0)
|
||||
pid_d = tl.program_id(1)
|
||||
d_offs = pid_d * BLOCK_D + tl.arange(0, BLOCK_D)
|
||||
d_mask = d_offs < HEAD_DIM
|
||||
|
||||
plan_base = plan_ptr + pid_p * plan_row_stride
|
||||
ragged_id = tl.load(plan_base + 0).to(tl.int32)
|
||||
batch_id = tl.load(plan_base + 1).to(tl.int32)
|
||||
position = tl.load(plan_base + 2).to(tl.int32)
|
||||
window_len = tl.load(plan_base + 3).to(tl.int32)
|
||||
if ragged_id < 0:
|
||||
return
|
||||
|
||||
extra_base = extra_ptr + batch_id.to(tl.int64) * 4
|
||||
load_first_page = tl.load(extra_base + 0).to(tl.int64)
|
||||
load_second_page = tl.load(extra_base + 1).to(tl.int64)
|
||||
|
||||
NEG_BIG: tl.constexpr = -1.0e9
|
||||
running_max = tl.full((BLOCK_D,), NEG_BIG, tl.float32)
|
||||
running_sum = tl.zeros((BLOCK_D,), tl.float32)
|
||||
weighted = tl.zeros((BLOCK_D,), tl.float32)
|
||||
|
||||
for slot in tl.static_range(8):
|
||||
in_state = slot < window_len
|
||||
if slot < 4:
|
||||
page = tl.where(window_len <= 4, load_second_page, load_first_page)
|
||||
kv_off = 0
|
||||
score_off = 2 * HEAD_DIM
|
||||
slot_in_page = slot
|
||||
else:
|
||||
page = load_second_page
|
||||
kv_off = HEAD_DIM
|
||||
score_off = 3 * HEAD_DIM
|
||||
slot_in_page = slot - 4
|
||||
|
||||
src_pos = position - 7 + slot
|
||||
state_valid = in_state & (src_pos >= 0)
|
||||
slot_base = page * buffer_page_stride + slot_in_page * buffer_slot_stride
|
||||
in_row = ragged_id - (7 - slot)
|
||||
in_row_safe = tl.where(in_state, 0, in_row)
|
||||
in_base = in_row_safe.to(tl.int64) * kv_in_row_stride
|
||||
|
||||
kv_state = tl.load(
|
||||
buffer_ptr + slot_base + kv_off + d_offs,
|
||||
mask=d_mask & state_valid,
|
||||
other=0.0,
|
||||
)
|
||||
score_state = tl.load(
|
||||
buffer_ptr + slot_base + score_off + d_offs,
|
||||
mask=d_mask & state_valid,
|
||||
other=NEG_BIG,
|
||||
)
|
||||
kv_input = tl.load(
|
||||
kv_in_ptr + in_base + kv_off + d_offs,
|
||||
mask=d_mask & (~in_state),
|
||||
other=0.0,
|
||||
)
|
||||
score_input = tl.load(
|
||||
kv_in_ptr + in_base + score_off + d_offs,
|
||||
mask=d_mask & (~in_state),
|
||||
other=NEG_BIG,
|
||||
)
|
||||
kv = tl.where(in_state, kv_state, kv_input)
|
||||
score = tl.where(in_state, score_state, score_input)
|
||||
bias = tl.load(ape_ptr + slot * ape_row_stride + d_offs, mask=d_mask, other=0.0)
|
||||
|
||||
s = score + bias
|
||||
new_max = tl.maximum(running_max, s)
|
||||
factor = tl.exp(running_max - new_max)
|
||||
e = tl.exp(s - new_max)
|
||||
running_sum = running_sum * factor + e
|
||||
weighted = weighted * factor + kv * e
|
||||
running_max = new_max
|
||||
|
||||
tl.store(
|
||||
out_ptr + ragged_id.to(tl.int64) * out_row_stride + d_offs,
|
||||
weighted / running_sum,
|
||||
mask=d_mask,
|
||||
)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _c4_prefill_write_kernel(
|
||||
kv_in_ptr,
|
||||
buffer_ptr,
|
||||
indices_ptr,
|
||||
extra_ptr,
|
||||
plan_ptr,
|
||||
kv_in_row_stride,
|
||||
buffer_page_stride,
|
||||
buffer_slot_stride,
|
||||
plan_row_stride,
|
||||
HEAD_DIM: tl.constexpr,
|
||||
BLOCK_D: tl.constexpr,
|
||||
):
|
||||
pid_p = tl.program_id(0)
|
||||
pid_d = tl.program_id(1)
|
||||
d_offs = pid_d * BLOCK_D + tl.arange(0, BLOCK_D)
|
||||
d_mask = d_offs < HEAD_DIM
|
||||
|
||||
plan_base = plan_ptr + pid_p * plan_row_stride
|
||||
ragged_id = tl.load(plan_base + 0).to(tl.int32)
|
||||
batch_id = tl.load(plan_base + 1).to(tl.int32)
|
||||
position = tl.load(plan_base + 2).to(tl.int32)
|
||||
if ragged_id < 0:
|
||||
return
|
||||
|
||||
extra_base = extra_ptr + batch_id.to(tl.int64) * 4
|
||||
write_first_page = tl.load(extra_base + 2).to(tl.int64)
|
||||
last_position = tl.load(extra_base + 3).to(tl.int32)
|
||||
write_second_page = tl.load(indices_ptr + batch_id).to(tl.int64)
|
||||
page = tl.where(position < last_position, write_first_page, write_second_page)
|
||||
slot = position % 4
|
||||
|
||||
in_base = ragged_id.to(tl.int64) * kv_in_row_stride
|
||||
dst_base = page * buffer_page_stride + slot.to(tl.int64) * buffer_slot_stride
|
||||
for ch in tl.static_range(4):
|
||||
ch_off = ch * HEAD_DIM
|
||||
val = tl.load(kv_in_ptr + in_base + ch_off + d_offs, mask=d_mask, other=0.0)
|
||||
tl.store(buffer_ptr + dst_base + ch_off + d_offs, val, mask=d_mask)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _c128_decode_kernel(
|
||||
kv_in_ptr,
|
||||
out_ptr,
|
||||
buffer_ptr,
|
||||
ape_ptr,
|
||||
indices_ptr,
|
||||
seq_lens_ptr,
|
||||
extra_ptr,
|
||||
kv_in_row_stride,
|
||||
out_row_stride,
|
||||
buffer_page_stride,
|
||||
buffer_slot_stride,
|
||||
ape_row_stride,
|
||||
HEAD_DIM: tl.constexpr,
|
||||
BLOCK_D: tl.constexpr,
|
||||
BLOCK_S: tl.constexpr,
|
||||
):
|
||||
bid = tl.program_id(0)
|
||||
pid_d = tl.program_id(1)
|
||||
d_offs = pid_d * BLOCK_D + tl.arange(0, BLOCK_D)
|
||||
d_mask = d_offs < HEAD_DIM
|
||||
|
||||
index = tl.load(indices_ptr + bid).to(tl.int64)
|
||||
index_prev = tl.load(extra_ptr + bid).to(tl.int64)
|
||||
seq_len = tl.load(seq_lens_ptr + bid).to(tl.int32)
|
||||
write_slot = (seq_len + 127) % 128
|
||||
in_base = bid.to(tl.int64) * kv_in_row_stride
|
||||
dst_base = index * buffer_page_stride + write_slot.to(tl.int64) * buffer_slot_stride
|
||||
|
||||
for ch in tl.static_range(2):
|
||||
ch_off = ch * HEAD_DIM
|
||||
val = tl.load(kv_in_ptr + in_base + ch_off + d_offs, mask=d_mask, other=0.0)
|
||||
tl.store(buffer_ptr + dst_base + ch_off + d_offs, val, mask=d_mask)
|
||||
|
||||
NEG_BIG: tl.constexpr = -1.0e9
|
||||
running_max = tl.full((BLOCK_D,), NEG_BIG, tl.float32)
|
||||
running_sum = tl.zeros((BLOCK_D,), tl.float32)
|
||||
weighted = tl.zeros((BLOCK_D,), tl.float32)
|
||||
|
||||
for chunk_start in tl.static_range(0, 128, BLOCK_S):
|
||||
slot_offs = chunk_start + tl.arange(0, BLOCK_S)
|
||||
src_pos = seq_len - 128 + slot_offs
|
||||
valid = src_pos >= 0
|
||||
is_input = slot_offs == 127
|
||||
write_pos = ((seq_len - 1) // 128) * 128
|
||||
pages = tl.where(src_pos < write_pos, index_prev, index)
|
||||
slot_in_page = src_pos % 128
|
||||
slot_bases = (
|
||||
pages * buffer_page_stride + slot_in_page.to(tl.int64) * buffer_slot_stride
|
||||
)
|
||||
kv_tile = tl.load(
|
||||
buffer_ptr + slot_bases[:, None] + d_offs[None, :],
|
||||
mask=valid[:, None] & (~is_input)[:, None] & d_mask[None, :],
|
||||
other=0.0,
|
||||
)
|
||||
score_tile = tl.load(
|
||||
buffer_ptr + slot_bases[:, None] + HEAD_DIM + d_offs[None, :],
|
||||
mask=valid[:, None] & (~is_input)[:, None] & d_mask[None, :],
|
||||
other=NEG_BIG,
|
||||
)
|
||||
kv_input_tile = tl.load(
|
||||
kv_in_ptr + in_base + d_offs[None, :],
|
||||
mask=valid[:, None] & is_input[:, None] & d_mask[None, :],
|
||||
other=0.0,
|
||||
)
|
||||
score_input_tile = tl.load(
|
||||
kv_in_ptr + in_base + HEAD_DIM + d_offs[None, :],
|
||||
mask=valid[:, None] & is_input[:, None] & d_mask[None, :],
|
||||
other=NEG_BIG,
|
||||
)
|
||||
kv_tile = tl.where(is_input[:, None], kv_input_tile, kv_tile)
|
||||
score_tile = tl.where(is_input[:, None], score_input_tile, score_tile)
|
||||
bias_tile = tl.load(
|
||||
ape_ptr + slot_offs[:, None] * ape_row_stride + d_offs[None, :],
|
||||
mask=d_mask[None, :],
|
||||
other=0.0,
|
||||
)
|
||||
s = score_tile + bias_tile
|
||||
local_max = tl.max(s, axis=0)
|
||||
new_max = tl.maximum(running_max, local_max)
|
||||
exp_s = tl.exp(s - new_max[None, :])
|
||||
exp_s = tl.where(valid[:, None], exp_s, 0.0)
|
||||
factor = tl.exp(running_max - new_max)
|
||||
running_sum = running_sum * factor + tl.sum(exp_s, axis=0)
|
||||
weighted = weighted * factor + tl.sum(kv_tile * exp_s, axis=0)
|
||||
running_max = new_max
|
||||
|
||||
tl.store(
|
||||
out_ptr + bid.to(tl.int64) * out_row_stride + d_offs,
|
||||
weighted / running_sum,
|
||||
mask=d_mask,
|
||||
)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _c128_prefill_compress_kernel(
|
||||
kv_in_ptr,
|
||||
out_ptr,
|
||||
buffer_ptr,
|
||||
ape_ptr,
|
||||
indices_ptr,
|
||||
plan_ptr,
|
||||
kv_in_row_stride,
|
||||
out_row_stride,
|
||||
buffer_page_stride,
|
||||
buffer_slot_stride,
|
||||
ape_row_stride,
|
||||
plan_row_stride,
|
||||
HEAD_DIM: tl.constexpr,
|
||||
BLOCK_D: tl.constexpr,
|
||||
BLOCK_S: tl.constexpr,
|
||||
):
|
||||
pid_p = tl.program_id(0)
|
||||
pid_d = tl.program_id(1)
|
||||
d_offs = pid_d * BLOCK_D + tl.arange(0, BLOCK_D)
|
||||
d_mask = d_offs < HEAD_DIM
|
||||
|
||||
plan_base = plan_ptr + pid_p * plan_row_stride
|
||||
ragged_id = tl.load(plan_base + 0).to(tl.int32)
|
||||
batch_id = tl.load(plan_base + 1).to(tl.int32)
|
||||
position = tl.load(plan_base + 2).to(tl.int32)
|
||||
window_len = tl.load(plan_base + 3).to(tl.int32)
|
||||
if ragged_id < 0:
|
||||
return
|
||||
|
||||
index = tl.load(indices_ptr + batch_id).to(tl.int64)
|
||||
NEG_BIG: tl.constexpr = -1.0e9
|
||||
running_max = tl.full((BLOCK_D,), NEG_BIG, tl.float32)
|
||||
running_sum = tl.zeros((BLOCK_D,), tl.float32)
|
||||
weighted = tl.zeros((BLOCK_D,), tl.float32)
|
||||
|
||||
for chunk_start in tl.static_range(0, 128, BLOCK_S):
|
||||
slot_offs = chunk_start + tl.arange(0, BLOCK_S)
|
||||
is_state = slot_offs < window_len
|
||||
src_pos = position - 127 + slot_offs
|
||||
state_valid = is_state & (src_pos >= 0)
|
||||
slot_bases = (
|
||||
index * buffer_page_stride + slot_offs.to(tl.int64) * buffer_slot_stride
|
||||
)
|
||||
in_rows = ragged_id - (127 - slot_offs)
|
||||
in_rows_safe = tl.where(is_state, tl.zeros_like(in_rows), in_rows)
|
||||
in_bases = in_rows_safe.to(tl.int64) * kv_in_row_stride
|
||||
|
||||
kv_state = tl.load(
|
||||
buffer_ptr + slot_bases[:, None] + d_offs[None, :],
|
||||
mask=state_valid[:, None] & d_mask[None, :],
|
||||
other=0.0,
|
||||
)
|
||||
score_state = tl.load(
|
||||
buffer_ptr + slot_bases[:, None] + HEAD_DIM + d_offs[None, :],
|
||||
mask=state_valid[:, None] & d_mask[None, :],
|
||||
other=NEG_BIG,
|
||||
)
|
||||
kv_input = tl.load(
|
||||
kv_in_ptr + in_bases[:, None] + d_offs[None, :],
|
||||
mask=(~is_state)[:, None] & d_mask[None, :],
|
||||
other=0.0,
|
||||
)
|
||||
score_input = tl.load(
|
||||
kv_in_ptr + in_bases[:, None] + HEAD_DIM + d_offs[None, :],
|
||||
mask=(~is_state)[:, None] & d_mask[None, :],
|
||||
other=NEG_BIG,
|
||||
)
|
||||
kv_tile = tl.where(is_state[:, None], kv_state, kv_input)
|
||||
score_tile = tl.where(is_state[:, None], score_state, score_input)
|
||||
bias_tile = tl.load(
|
||||
ape_ptr + slot_offs[:, None] * ape_row_stride + d_offs[None, :],
|
||||
mask=d_mask[None, :],
|
||||
other=0.0,
|
||||
)
|
||||
|
||||
s = score_tile + bias_tile
|
||||
local_max = tl.max(s, axis=0)
|
||||
new_max = tl.maximum(running_max, local_max)
|
||||
exp_s = tl.exp(s - new_max[None, :])
|
||||
# Keep input-path entries valid; only state-path entries need src_pos guard.
|
||||
valid = state_valid | (~is_state)
|
||||
exp_s = tl.where(valid[:, None], exp_s, 0.0)
|
||||
factor = tl.exp(running_max - new_max)
|
||||
running_sum = running_sum * factor + tl.sum(exp_s, axis=0)
|
||||
weighted = weighted * factor + tl.sum(kv_tile * exp_s, axis=0)
|
||||
running_max = new_max
|
||||
|
||||
tl.store(
|
||||
out_ptr + ragged_id.to(tl.int64) * out_row_stride + d_offs,
|
||||
weighted / running_sum,
|
||||
mask=d_mask,
|
||||
)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _c128_prefill_write_kernel(
|
||||
kv_in_ptr,
|
||||
buffer_ptr,
|
||||
indices_ptr,
|
||||
plan_ptr,
|
||||
kv_in_row_stride,
|
||||
buffer_page_stride,
|
||||
buffer_slot_stride,
|
||||
plan_row_stride,
|
||||
HEAD_DIM: tl.constexpr,
|
||||
BLOCK_D: tl.constexpr,
|
||||
):
|
||||
pid_p = tl.program_id(0)
|
||||
pid_d = tl.program_id(1)
|
||||
d_offs = pid_d * BLOCK_D + tl.arange(0, BLOCK_D)
|
||||
d_mask = d_offs < HEAD_DIM
|
||||
|
||||
plan_base = plan_ptr + pid_p * plan_row_stride
|
||||
ragged_id = tl.load(plan_base + 0).to(tl.int32)
|
||||
batch_id = tl.load(plan_base + 1).to(tl.int32)
|
||||
position = tl.load(plan_base + 2).to(tl.int32)
|
||||
if ragged_id < 0:
|
||||
return
|
||||
|
||||
index = tl.load(indices_ptr + batch_id).to(tl.int64)
|
||||
slot = position % 128
|
||||
in_base = ragged_id.to(tl.int64) * kv_in_row_stride
|
||||
dst_base = index * buffer_page_stride + slot.to(tl.int64) * buffer_slot_stride
|
||||
|
||||
for ch in tl.static_range(2):
|
||||
ch_off = ch * HEAD_DIM
|
||||
val = tl.load(kv_in_ptr + in_base + ch_off + d_offs, mask=d_mask, other=0.0)
|
||||
tl.store(buffer_ptr + dst_base + ch_off + d_offs, val, mask=d_mask)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _compress_norm_rope_kernel(
|
||||
kv_ptr,
|
||||
weight_ptr,
|
||||
freqs_ptr,
|
||||
handle_ptr,
|
||||
eps,
|
||||
kv_row_stride,
|
||||
freqs_row_stride,
|
||||
plan_row_stride,
|
||||
HEAD_DIM: tl.constexpr,
|
||||
ROPE_DIM: tl.constexpr,
|
||||
HEAD_BLOCK: tl.constexpr,
|
||||
ROPE_PAIR_BLOCK: tl.constexpr,
|
||||
COMPRESS_RATIO: tl.constexpr,
|
||||
IS_DECODE: tl.constexpr,
|
||||
):
|
||||
work_id = tl.program_id(0)
|
||||
|
||||
if IS_DECODE:
|
||||
row = work_id
|
||||
seq_len = tl.load(handle_ptr + work_id).to(tl.int32)
|
||||
position = ((seq_len - 1) // COMPRESS_RATIO) * COMPRESS_RATIO
|
||||
else:
|
||||
plan_base = handle_ptr + work_id * plan_row_stride
|
||||
row = tl.load(plan_base + 0).to(tl.int32)
|
||||
plan_position = tl.load(plan_base + 2).to(tl.int32)
|
||||
if row < 0:
|
||||
return
|
||||
position = plan_position + 1 - COMPRESS_RATIO
|
||||
|
||||
base = row.to(tl.int64) * kv_row_stride
|
||||
offs = tl.arange(0, HEAD_BLOCK)
|
||||
mask = offs < HEAD_DIM
|
||||
x = tl.load(kv_ptr + base + offs, mask=mask, other=0.0).to(tl.float32)
|
||||
w = tl.load(weight_ptr + offs, mask=mask, other=0.0).to(tl.float32)
|
||||
rms_inv = tl.rsqrt(tl.sum(x * x, axis=0) / HEAD_DIM + eps)
|
||||
x_normed = x * rms_inv * w
|
||||
|
||||
rope_start: tl.constexpr = HEAD_DIM - ROPE_DIM
|
||||
pair_offs = tl.arange(0, ROPE_PAIR_BLOCK)
|
||||
pair_mask = pair_offs < (ROPE_DIM // 2)
|
||||
x_real = tl.load(
|
||||
kv_ptr + base + rope_start + 2 * pair_offs,
|
||||
mask=pair_mask,
|
||||
other=0.0,
|
||||
).to(tl.float32)
|
||||
x_imag = tl.load(
|
||||
kv_ptr + base + rope_start + 2 * pair_offs + 1,
|
||||
mask=pair_mask,
|
||||
other=0.0,
|
||||
).to(tl.float32)
|
||||
w_real = tl.load(
|
||||
weight_ptr + rope_start + 2 * pair_offs,
|
||||
mask=pair_mask,
|
||||
other=1.0,
|
||||
).to(tl.float32)
|
||||
w_imag = tl.load(
|
||||
weight_ptr + rope_start + 2 * pair_offs + 1,
|
||||
mask=pair_mask,
|
||||
other=1.0,
|
||||
).to(tl.float32)
|
||||
x_real = x_real * rms_inv * w_real
|
||||
x_imag = x_imag * rms_inv * w_imag
|
||||
|
||||
freq_base = position.to(tl.int64) * freqs_row_stride
|
||||
f_real = tl.load(freqs_ptr + freq_base + 2 * pair_offs, mask=pair_mask, other=0.0)
|
||||
f_imag = tl.load(
|
||||
freqs_ptr + freq_base + 2 * pair_offs + 1,
|
||||
mask=pair_mask,
|
||||
other=0.0,
|
||||
)
|
||||
out_real = x_real * f_real - x_imag * f_imag
|
||||
out_imag = x_real * f_imag + x_imag * f_real
|
||||
|
||||
tl.store(kv_ptr + base + offs, x_normed, mask=mask & (offs < rope_start))
|
||||
tl.store(kv_ptr + base + rope_start + 2 * pair_offs, out_real, mask=pair_mask)
|
||||
tl.store(kv_ptr + base + rope_start + 2 * pair_offs + 1, out_imag, mask=pair_mask)
|
||||
|
||||
|
||||
def _plan_as_i32(plan: torch.Tensor) -> torch.Tensor:
|
||||
assert plan.dtype == torch.uint8 and plan.dim() == 2 and plan.shape[1] == 16
|
||||
return plan.view(torch.int32).view(-1, 4)
|
||||
|
||||
|
||||
def _block_d(head_dim: int) -> int:
|
||||
return min(32, triton.next_power_of_2(head_dim))
|
||||
|
||||
|
||||
def _check_common(
|
||||
kv_score_buffer: torch.Tensor,
|
||||
kv_score_input: torch.Tensor,
|
||||
out: torch.Tensor,
|
||||
ape: torch.Tensor,
|
||||
indices: torch.Tensor,
|
||||
head_dim: int,
|
||||
compress_ratio: int,
|
||||
) -> None:
|
||||
coff = 2 if compress_ratio == 4 else 1
|
||||
assert kv_score_input.is_cuda and kv_score_buffer.is_cuda
|
||||
assert kv_score_input.dim() == 2 and kv_score_input.dtype == torch.float32
|
||||
assert kv_score_input.shape[1] == 2 * coff * head_dim
|
||||
assert kv_score_buffer.dim() == 3 and kv_score_buffer.dtype == torch.float32
|
||||
assert kv_score_buffer.shape[1:] == (compress_ratio, 2 * coff * head_dim)
|
||||
assert out.shape == (kv_score_input.shape[0], head_dim)
|
||||
assert out.dtype == torch.float32 and out.is_cuda
|
||||
assert ape.shape == (compress_ratio * coff, head_dim)
|
||||
assert ape.dtype == torch.float32 and ape.is_cuda
|
||||
assert indices.dtype == torch.int32 and indices.is_cuda
|
||||
|
||||
|
||||
def _is_decode_plan(plan: Union[CompressorDecodePlan, CompressorPrefillPlan]) -> bool:
|
||||
return isinstance(plan, CompressorDecodePlan)
|
||||
|
||||
|
||||
def hip_compress_forward(
|
||||
*,
|
||||
kv_score_buffer: torch.Tensor,
|
||||
kv_score_input: torch.Tensor,
|
||||
ape: torch.Tensor,
|
||||
indices: torch.Tensor,
|
||||
plan: Union[CompressorDecodePlan, CompressorPrefillPlan],
|
||||
extra_data: Optional[torch.Tensor],
|
||||
head_dim: int,
|
||||
compress_ratio: int,
|
||||
out: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
if compress_ratio not in (4, 128):
|
||||
raise ValueError(f"unsupported {compress_ratio=}")
|
||||
if out is None:
|
||||
out = kv_score_input.new_empty((kv_score_input.shape[0], head_dim))
|
||||
is_decode = _is_decode_plan(plan)
|
||||
if not is_decode:
|
||||
out.fill_(10000.0)
|
||||
|
||||
_check_common(
|
||||
kv_score_buffer,
|
||||
kv_score_input,
|
||||
out,
|
||||
ape,
|
||||
indices,
|
||||
head_dim,
|
||||
compress_ratio,
|
||||
)
|
||||
|
||||
BLOCK_D = _block_d(head_dim)
|
||||
num_d_chunks = triton.cdiv(head_dim, BLOCK_D)
|
||||
|
||||
if is_decode:
|
||||
seq_lens = plan.seq_lens
|
||||
assert seq_lens.dtype == torch.int32 and seq_lens.is_cuda
|
||||
assert seq_lens.shape == indices.shape
|
||||
grid = (seq_lens.numel(), num_d_chunks)
|
||||
if compress_ratio == 4:
|
||||
assert extra_data is not None
|
||||
assert extra_data.shape == (seq_lens.numel(), 1)
|
||||
_c4_decode_kernel[grid](
|
||||
kv_score_input,
|
||||
out,
|
||||
kv_score_buffer,
|
||||
ape,
|
||||
indices,
|
||||
seq_lens,
|
||||
extra_data,
|
||||
kv_score_input.stride(0),
|
||||
out.stride(0),
|
||||
kv_score_buffer.stride(0),
|
||||
kv_score_buffer.stride(1),
|
||||
ape.stride(0),
|
||||
HEAD_DIM=head_dim,
|
||||
BLOCK_D=BLOCK_D,
|
||||
)
|
||||
else:
|
||||
assert extra_data is not None
|
||||
assert extra_data.shape == seq_lens.shape
|
||||
_c128_decode_kernel[grid](
|
||||
kv_score_input,
|
||||
out,
|
||||
kv_score_buffer,
|
||||
ape,
|
||||
indices,
|
||||
seq_lens,
|
||||
extra_data,
|
||||
kv_score_input.stride(0),
|
||||
out.stride(0),
|
||||
kv_score_buffer.stride(0),
|
||||
kv_score_buffer.stride(1),
|
||||
ape.stride(0),
|
||||
HEAD_DIM=head_dim,
|
||||
BLOCK_D=BLOCK_D,
|
||||
BLOCK_S=64,
|
||||
)
|
||||
return out
|
||||
|
||||
compress_plan = _plan_as_i32(plan.compress_plan)
|
||||
write_plan = _plan_as_i32(plan.write_plan)
|
||||
if compress_ratio == 4:
|
||||
assert extra_data is not None
|
||||
assert extra_data.dim() == 2 and extra_data.shape[1] == 4
|
||||
compress_grid = (compress_plan.shape[0], num_d_chunks)
|
||||
write_grid = (write_plan.shape[0], num_d_chunks)
|
||||
_c4_prefill_compress_kernel[compress_grid](
|
||||
kv_score_input,
|
||||
out,
|
||||
kv_score_buffer,
|
||||
ape,
|
||||
indices,
|
||||
extra_data,
|
||||
compress_plan,
|
||||
kv_score_input.stride(0),
|
||||
out.stride(0),
|
||||
kv_score_buffer.stride(0),
|
||||
kv_score_buffer.stride(1),
|
||||
ape.stride(0),
|
||||
compress_plan.stride(0),
|
||||
HEAD_DIM=head_dim,
|
||||
BLOCK_D=BLOCK_D,
|
||||
)
|
||||
_c4_prefill_write_kernel[write_grid](
|
||||
kv_score_input,
|
||||
kv_score_buffer,
|
||||
indices,
|
||||
extra_data,
|
||||
write_plan,
|
||||
kv_score_input.stride(0),
|
||||
kv_score_buffer.stride(0),
|
||||
kv_score_buffer.stride(1),
|
||||
write_plan.stride(0),
|
||||
HEAD_DIM=head_dim,
|
||||
BLOCK_D=BLOCK_D,
|
||||
)
|
||||
else:
|
||||
load_indices = indices if extra_data is None else extra_data
|
||||
assert load_indices.dim() == 1 and load_indices.dtype == torch.int32
|
||||
compress_grid = (compress_plan.shape[0], num_d_chunks)
|
||||
write_grid = (write_plan.shape[0], num_d_chunks)
|
||||
_c128_prefill_compress_kernel[compress_grid](
|
||||
kv_score_input,
|
||||
out,
|
||||
kv_score_buffer,
|
||||
ape,
|
||||
load_indices,
|
||||
compress_plan,
|
||||
kv_score_input.stride(0),
|
||||
out.stride(0),
|
||||
kv_score_buffer.stride(0),
|
||||
kv_score_buffer.stride(1),
|
||||
ape.stride(0),
|
||||
compress_plan.stride(0),
|
||||
HEAD_DIM=head_dim,
|
||||
BLOCK_D=BLOCK_D,
|
||||
BLOCK_S=64,
|
||||
)
|
||||
_c128_prefill_write_kernel[write_grid](
|
||||
kv_score_input,
|
||||
kv_score_buffer,
|
||||
indices,
|
||||
write_plan,
|
||||
kv_score_input.stride(0),
|
||||
kv_score_buffer.stride(0),
|
||||
kv_score_buffer.stride(1),
|
||||
write_plan.stride(0),
|
||||
HEAD_DIM=head_dim,
|
||||
BLOCK_D=BLOCK_D,
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
def hip_compress_fused_norm_rope_inplace(
|
||||
kv: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
eps: float,
|
||||
freqs_cis: torch.Tensor,
|
||||
plan: Union[CompressorDecodePlan, CompressorPrefillPlan],
|
||||
) -> None:
|
||||
assert kv.dim() == 2 and kv.stride(-1) == 1
|
||||
assert weight.shape == (kv.shape[1],)
|
||||
freqs_real = torch.view_as_real(freqs_cis).flatten(-2)
|
||||
head_dim = kv.shape[1]
|
||||
rope_dim = freqs_real.shape[-1]
|
||||
assert head_dim >= rope_dim and rope_dim % 2 == 0
|
||||
|
||||
is_decode = _is_decode_plan(plan)
|
||||
if is_decode:
|
||||
handle = plan.seq_lens
|
||||
else:
|
||||
handle = _plan_as_i32(plan.compress_plan)
|
||||
|
||||
if handle.numel() == 0:
|
||||
return
|
||||
|
||||
HEAD_BLOCK = triton.next_power_of_2(head_dim)
|
||||
ROPE_PAIR_BLOCK = max(triton.next_power_of_2(rope_dim // 2), 1)
|
||||
_compress_norm_rope_kernel[(handle.shape[0],)](
|
||||
kv,
|
||||
weight,
|
||||
freqs_real,
|
||||
handle,
|
||||
eps,
|
||||
kv.stride(0),
|
||||
freqs_real.stride(0),
|
||||
handle.stride(0) if not is_decode else 0,
|
||||
HEAD_DIM=head_dim,
|
||||
ROPE_DIM=rope_dim,
|
||||
HEAD_BLOCK=HEAD_BLOCK,
|
||||
ROPE_PAIR_BLOCK=ROPE_PAIR_BLOCK,
|
||||
COMPRESS_RATIO=plan.compress_ratio,
|
||||
IS_DECODE=is_decode,
|
||||
)
|
||||
@@ -39,6 +39,9 @@ else:
|
||||
FP8_MAX = torch.finfo(FP8_DTYPE).max
|
||||
|
||||
|
||||
_arange_cache = {}
|
||||
|
||||
|
||||
def fp8_paged_mqa_logits_torch(
|
||||
q_fp8: torch.Tensor,
|
||||
kvcache_fp8: torch.Tensor,
|
||||
@@ -49,12 +52,13 @@ def fp8_paged_mqa_logits_torch(
|
||||
max_seq_len: int,
|
||||
clean_logits: bool = True,
|
||||
) -> torch.Tensor:
|
||||
"""Vectorized implementation compatible with CUDA graph capture."""
|
||||
_ = deep_gemm_metadata
|
||||
batch_size, _, num_heads, head_dim = q_fp8.shape
|
||||
block_size = kvcache_fp8.shape[1]
|
||||
|
||||
assert head_dim == 128, "torch reference impl hardcodes DSV4 indexer head_dim=128"
|
||||
assert block_size == 64, "torch reference impl hardcodes block_size=64 cache layout"
|
||||
assert head_dim == 128
|
||||
assert block_size == 64
|
||||
assert q_fp8.shape == (batch_size, 1, num_heads, head_dim)
|
||||
assert kvcache_fp8.shape[1:] == (block_size, 1, head_dim + 4)
|
||||
assert weight.shape == (batch_size, num_heads)
|
||||
@@ -62,32 +66,85 @@ def fp8_paged_mqa_logits_torch(
|
||||
assert page_table.shape[0] == batch_size
|
||||
assert clean_logits == False
|
||||
|
||||
logits = page_table.new_empty((batch_size, max_seq_len), dtype=torch.float32)
|
||||
for i in range(batch_size):
|
||||
q = q_fp8[i, 0]
|
||||
q = q.to(torch.float32)
|
||||
q_scale = weight[i]
|
||||
seq_len = int(seq_lens[i].item())
|
||||
assert seq_len <= max_seq_len
|
||||
num_pages = (seq_len + block_size - 1) // block_size
|
||||
padded_seq_len = num_pages * block_size
|
||||
pages = page_table[i, :num_pages]
|
||||
kvcache_fp8 = kvcache_fp8.view(-1, block_size * (head_dim + 4))
|
||||
kvcache = kvcache_fp8[pages]
|
||||
SCALE_OFFSET = block_size * head_dim
|
||||
kvcache_value = kvcache[..., :SCALE_OFFSET].view(dtype=FP8_DTYPE)
|
||||
kvcache_scale = kvcache[..., SCALE_OFFSET:].view(dtype=torch.float32)
|
||||
kvcache_value = kvcache_value.to(torch.float32)
|
||||
kvcache_scale = kvcache_scale.contiguous()
|
||||
kvcache_value = kvcache_value.view(padded_seq_len, head_dim)
|
||||
kvcache_scale = kvcache_scale.view(padded_seq_len)
|
||||
score = F.linear(kvcache_value, q)
|
||||
score = F.relu(score)
|
||||
score *= q_scale[None, :]
|
||||
score = score.sum(dim=1)
|
||||
score *= kvcache_scale
|
||||
logits[i, :seq_len] = score[:seq_len]
|
||||
max_num_pages = page_table.shape[1]
|
||||
SCALE_OFFSET = block_size * head_dim
|
||||
total_dim = block_size * (head_dim + 4)
|
||||
|
||||
kvcache_flat = kvcache_fp8.view(-1, total_dim)
|
||||
|
||||
pages_clamped = page_table.clamp(min=0)
|
||||
kvcache_gathered = kvcache_flat[pages_clamped]
|
||||
|
||||
kv_values_raw = kvcache_gathered[..., :SCALE_OFFSET].contiguous()
|
||||
kv_values_fp8 = kv_values_raw.view(dtype=FP8_DTYPE)
|
||||
kv_values = kv_values_fp8.to(torch.float32)
|
||||
kv_values = kv_values.reshape(batch_size, max_num_pages * block_size, head_dim)
|
||||
|
||||
kv_scales_raw = kvcache_gathered[..., SCALE_OFFSET:].contiguous()
|
||||
kv_scales = kv_scales_raw.view(dtype=torch.float32)
|
||||
kv_scales = kv_scales.reshape(batch_size, max_num_pages * block_size)
|
||||
|
||||
q_float = q_fp8[:, 0].to(torch.float32)
|
||||
scores = torch.bmm(kv_values, q_float.transpose(1, 2))
|
||||
scores = F.relu(scores)
|
||||
scores = scores * weight.unsqueeze(1)
|
||||
scores = scores.sum(dim=2)
|
||||
scores = scores * kv_scales
|
||||
|
||||
padded_seq_len = max_num_pages * block_size
|
||||
cache = _arange_cache
|
||||
arange_key = f"arange_{padded_seq_len}_{scores.device}"
|
||||
if arange_key not in cache:
|
||||
cache[arange_key] = torch.arange(padded_seq_len, device=scores.device)
|
||||
positions = cache[arange_key].unsqueeze(0)
|
||||
valid_mask = positions < seq_lens.unsqueeze(1)
|
||||
scores = scores.masked_fill(~valid_mask, 0.0)
|
||||
|
||||
if padded_seq_len < max_seq_len:
|
||||
scores = F.pad(scores, (0, max_seq_len - padded_seq_len), value=0.0)
|
||||
else:
|
||||
scores = scores[:, :max_seq_len]
|
||||
|
||||
return scores
|
||||
|
||||
|
||||
def _aiter_fp8_paged_mqa_logits(
|
||||
q_fp8: torch.Tensor,
|
||||
kvcache_fp8: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
seq_lens: torch.Tensor,
|
||||
page_table: torch.Tensor,
|
||||
deep_gemm_metadata: Any,
|
||||
max_seq_len: int,
|
||||
clean_logits: bool = False,
|
||||
) -> torch.Tensor:
|
||||
"""Wrapper adapting aiter's deepgemm_fp8_paged_mqa_logits to SGLang's interface."""
|
||||
from aiter.ops.triton.attention.pa_mqa_logits import (
|
||||
deepgemm_fp8_paged_mqa_logits,
|
||||
)
|
||||
|
||||
batch_size = q_fp8.shape[0]
|
||||
next_n = q_fp8.shape[1]
|
||||
total_tokens = batch_size * next_n
|
||||
_sl = seq_lens.squeeze(-1) if seq_lens.dim() == 2 else seq_lens
|
||||
kv_block_size = kvcache_fp8.shape[1]
|
||||
logits = torch.empty(
|
||||
total_tokens,
|
||||
max_seq_len,
|
||||
dtype=torch.float32,
|
||||
device=q_fp8.device,
|
||||
)
|
||||
deepgemm_fp8_paged_mqa_logits(
|
||||
q_fp8,
|
||||
kvcache_fp8,
|
||||
weight,
|
||||
logits,
|
||||
_sl.to(torch.int32),
|
||||
page_table.to(torch.int32),
|
||||
max_seq_len,
|
||||
KVBlockSize=kv_block_size,
|
||||
Preshuffle=True,
|
||||
)
|
||||
return logits
|
||||
|
||||
|
||||
@@ -99,6 +156,9 @@ def topk_transform_512_pytorch_vectorized(
|
||||
page_size: int,
|
||||
out_raw_indices: Optional[torch.Tensor] = None,
|
||||
) -> None:
|
||||
"""Vectorized PyTorch fallback for topk_transform_512.
|
||||
All helper tensors (arange, zeros) are cached to avoid device-tensor
|
||||
creation during HIP/CUDA graph capture."""
|
||||
|
||||
TOPK = out_page_indices.shape[1]
|
||||
batch_size = scores.shape[0]
|
||||
@@ -108,13 +168,22 @@ def topk_transform_512_pytorch_vectorized(
|
||||
page_bits = (page_size - 1).bit_length() if page_size > 1 else 0
|
||||
page_mask = page_size - 1
|
||||
|
||||
positions = (
|
||||
torch.arange(max_seq_len, device=device).unsqueeze(0).expand(batch_size, -1)
|
||||
)
|
||||
cache = _arange_cache
|
||||
key_seq = f"arange_{max_seq_len}_{device}"
|
||||
key_topk = f"arange_{TOPK}_{device}"
|
||||
key_bs = f"arange_{batch_size}_{device}"
|
||||
if key_seq not in cache:
|
||||
cache[key_seq] = torch.arange(max_seq_len, device=device)
|
||||
if key_topk not in cache:
|
||||
cache[key_topk] = torch.arange(TOPK, device=device, dtype=torch.int32)
|
||||
if key_bs not in cache:
|
||||
cache[key_bs] = torch.arange(batch_size, device=device)
|
||||
|
||||
positions = cache[key_seq].unsqueeze(0).expand(batch_size, -1)
|
||||
valid_mask = positions < seq_lens.unsqueeze(1)
|
||||
|
||||
masked_scores = scores.clone()
|
||||
masked_scores[~valid_mask] = float("-inf")
|
||||
masked_scores.masked_fill_(~valid_mask, float("-inf"))
|
||||
|
||||
actual_k = min(TOPK, max_seq_len)
|
||||
_, raw_indices = torch.topk(
|
||||
@@ -123,44 +192,28 @@ def topk_transform_512_pytorch_vectorized(
|
||||
raw_indices = raw_indices.to(torch.int32)
|
||||
|
||||
if actual_k < TOPK:
|
||||
padding = torch.zeros(
|
||||
(batch_size, TOPK - actual_k), dtype=torch.int32, device=device
|
||||
)
|
||||
raw_indices = torch.cat([raw_indices, padding], dim=1)
|
||||
raw_indices = F.pad(raw_indices, (0, TOPK - actual_k), value=0)
|
||||
|
||||
batch_indices = (
|
||||
torch.arange(batch_size, device=device).unsqueeze(1).expand(-1, TOPK)
|
||||
)
|
||||
batch_indices = cache[key_bs].unsqueeze(1).expand(-1, TOPK)
|
||||
gathered_scores = scores[
|
||||
batch_indices.flatten(), raw_indices.clamp(min=0).flatten()
|
||||
].view(batch_size, TOPK)
|
||||
|
||||
valid_topk = gathered_scores != float("-inf")
|
||||
if actual_k < TOPK:
|
||||
pad_mask = torch.arange(TOPK, device=device).unsqueeze(0) >= actual_k
|
||||
pad_mask = cache[key_topk].unsqueeze(0) >= actual_k
|
||||
valid_topk = valid_topk & ~pad_mask
|
||||
|
||||
needs_sequential = seq_lens <= TOPK
|
||||
if needs_sequential.any():
|
||||
sequential_indices = (
|
||||
torch.arange(TOPK, device=device, dtype=torch.int32)
|
||||
.unsqueeze(0)
|
||||
.expand(batch_size, -1)
|
||||
)
|
||||
sequential_valid = sequential_indices < seq_lens.unsqueeze(1)
|
||||
sequential_indices = cache[key_topk].unsqueeze(0).expand(batch_size, -1)
|
||||
sequential_valid = sequential_indices < seq_lens.unsqueeze(1)
|
||||
|
||||
raw_indices = torch.where(
|
||||
needs_sequential.unsqueeze(1).expand(-1, TOPK),
|
||||
torch.where(
|
||||
sequential_valid,
|
||||
sequential_indices,
|
||||
torch.tensor(-1, device=device, dtype=torch.int32),
|
||||
),
|
||||
raw_indices,
|
||||
)
|
||||
valid_topk = torch.where(
|
||||
needs_sequential.unsqueeze(1).expand(-1, TOPK), sequential_valid, valid_topk
|
||||
)
|
||||
seq_indices_or_neg1 = sequential_indices.clone()
|
||||
seq_indices_or_neg1.masked_fill_(~sequential_valid, -1)
|
||||
|
||||
needs_seq_mask = needs_sequential.unsqueeze(1).expand(-1, TOPK)
|
||||
raw_indices = torch.where(needs_seq_mask, seq_indices_or_neg1, raw_indices)
|
||||
valid_topk = torch.where(needs_seq_mask, sequential_valid, valid_topk)
|
||||
|
||||
page_idx = raw_indices >> page_bits
|
||||
offset_in_page = raw_indices & page_mask
|
||||
@@ -170,17 +223,13 @@ def topk_transform_512_pytorch_vectorized(
|
||||
|
||||
page_indices = (physical_pages << page_bits) | offset_in_page
|
||||
page_indices = page_indices.to(torch.int32)
|
||||
|
||||
page_indices = torch.where(
|
||||
valid_topk, page_indices, torch.tensor(-1, device=device, dtype=torch.int32)
|
||||
)
|
||||
page_indices.masked_fill_(~valid_topk, -1)
|
||||
|
||||
out_page_indices.copy_(page_indices)
|
||||
|
||||
if out_raw_indices is not None:
|
||||
raw_indices = torch.where(
|
||||
valid_topk, raw_indices, torch.tensor(-1, device=device, dtype=torch.int32)
|
||||
)
|
||||
raw_indices = raw_indices.clone()
|
||||
raw_indices.masked_fill_(~valid_topk, -1)
|
||||
out_raw_indices.copy_(raw_indices)
|
||||
|
||||
|
||||
@@ -290,18 +339,20 @@ class C4IndexerBackendMixin:
|
||||
positions: torch.Tensor,
|
||||
forward_batch: ForwardBatch,
|
||||
token_to_kv_pool: DeepSeekV4TokenToKVPool,
|
||||
skip_compressor: bool = False,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
if TYPE_CHECKING:
|
||||
assert isinstance(self, CompressorBackendMixin)
|
||||
|
||||
weights = c4_indexer.compute_weights(x, skip_scale=True)
|
||||
q_fp8, weights = c4_indexer.compute_q(q_lora, positions, weights)
|
||||
self.forward_indexer_compressor(
|
||||
x=x,
|
||||
forward_batch=forward_batch,
|
||||
layer_id=c4_indexer.layer_id,
|
||||
compressor=c4_indexer.compressor,
|
||||
)
|
||||
if not skip_compressor:
|
||||
self.forward_indexer_compressor(
|
||||
x=x,
|
||||
forward_batch=forward_batch,
|
||||
layer_id=c4_indexer.layer_id,
|
||||
compressor=c4_indexer.compressor,
|
||||
)
|
||||
c4_indexer_kv_cache = token_to_kv_pool.get_index_k_with_scale_buffer(
|
||||
layer_id=c4_indexer.layer_id,
|
||||
)
|
||||
@@ -316,6 +367,7 @@ class C4IndexerBackendMixin:
|
||||
alt_streams: Optional[List[torch.cuda.Stream]] = None,
|
||||
enable_multi_stream: bool = False,
|
||||
q_lora_ready: Optional[torch.cuda.Event] = None,
|
||||
skip_compressor: bool = False,
|
||||
) -> None:
|
||||
if forward_batch.forward_mode.is_idle():
|
||||
return
|
||||
@@ -354,6 +406,7 @@ class C4IndexerBackendMixin:
|
||||
positions=core_metadata.positions,
|
||||
forward_batch=forward_batch,
|
||||
token_to_kv_pool=token_to_kv_pool,
|
||||
skip_compressor=skip_compressor,
|
||||
)
|
||||
|
||||
assert len(q_fp8.shape) == 3
|
||||
@@ -372,6 +425,8 @@ class C4IndexerBackendMixin:
|
||||
from sglang.srt.layers.attention.dsa.tilelang_kernel import (
|
||||
tilelang_fp8_paged_mqa_logits as fn,
|
||||
)
|
||||
elif envs.SGLANG_OPT_USE_AITER_INDEXER.get():
|
||||
fn = _aiter_fp8_paged_mqa_logits
|
||||
elif envs.SGLANG_FP8_PAGED_MQA_LOGITS_TORCH.get():
|
||||
fn = fp8_paged_mqa_logits_torch
|
||||
else:
|
||||
@@ -379,7 +434,8 @@ class C4IndexerBackendMixin:
|
||||
|
||||
_c4sl = indexer_metadata.c4_seq_lens
|
||||
_use_tilelang = envs.SGLANG_OPT_USE_TILELANG_INDEXER.get()
|
||||
if _c4sl.dim() == 1 and not _use_tilelang:
|
||||
_use_aiter = envs.SGLANG_OPT_USE_AITER_INDEXER.get()
|
||||
if _c4sl.dim() == 1 and not _use_tilelang and not _use_aiter:
|
||||
_c4sl = _c4sl.unsqueeze(-1)
|
||||
logits = fn(
|
||||
q_fp8,
|
||||
@@ -545,6 +601,7 @@ class C4Indexer(nn.Module):
|
||||
attn_backend: AttentionBackend,
|
||||
enable_multi_stream: bool = False,
|
||||
q_lora_ready: Optional[torch.cuda.Event] = None,
|
||||
skip_compressor: bool = False,
|
||||
) -> None:
|
||||
return attn_backend.forward_c4_indexer(
|
||||
x=x,
|
||||
@@ -554,4 +611,5 @@ class C4Indexer(nn.Module):
|
||||
alt_streams=self.alt_streams,
|
||||
enable_multi_stream=enable_multi_stream,
|
||||
q_lora_ready=q_lora_ready,
|
||||
skip_compressor=skip_compressor,
|
||||
)
|
||||
|
||||
@@ -103,7 +103,10 @@ class PagedIndexerMetadata:
|
||||
topk_metadata: torch.Tensor = field(init=False, repr=False)
|
||||
|
||||
def __post_init__(self):
|
||||
if envs.SGLANG_FP8_PAGED_MQA_LOGITS_TORCH.get():
|
||||
if (
|
||||
envs.SGLANG_FP8_PAGED_MQA_LOGITS_TORCH.get()
|
||||
or envs.SGLANG_OPT_USE_AITER_INDEXER.get()
|
||||
):
|
||||
self.deep_gemm_metadata = None
|
||||
else:
|
||||
import deep_gemm
|
||||
@@ -148,14 +151,17 @@ class PagedIndexerMetadata:
|
||||
def copy_(self, other: "PagedIndexerMetadata"):
|
||||
if is_hip():
|
||||
copy_fields = ["page_table", "c4_seq_lens"]
|
||||
assign_fields = ["deep_gemm_metadata"]
|
||||
else:
|
||||
copy_fields = ["page_table", "c4_seq_lens", "deep_gemm_metadata"]
|
||||
assign_fields = []
|
||||
copy_fields += ["topk_metadata"]
|
||||
copy_metadata(
|
||||
src=other,
|
||||
dst=self,
|
||||
check_eq_fields=["page_size"],
|
||||
copy_fields=copy_fields,
|
||||
assign_fields=assign_fields,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -12,10 +12,6 @@ def flash_mla_with_kvcache_entrypoint(backend: str, **kwargs):
|
||||
if is_hip():
|
||||
import os
|
||||
|
||||
from sglang.srt.layers.attention.dsa.tilelang_kernel import (
|
||||
dpsk_v4_fp8_attention_fwd,
|
||||
)
|
||||
|
||||
backend = os.environ.get("SGLANG_HACK_FLASHMLA_BACKEND", "tilelang")
|
||||
else:
|
||||
import flash_mla
|
||||
@@ -36,8 +32,19 @@ def flash_mla_with_kvcache_entrypoint(backend: str, **kwargs):
|
||||
return flash_mla_with_kvcache_torch(**kwargs)
|
||||
|
||||
if backend == "tilelang":
|
||||
from sglang.srt.layers.attention.dsa.tilelang_kernel import (
|
||||
dpsk_v4_fp8_attention_fwd,
|
||||
)
|
||||
|
||||
return dpsk_v4_fp8_attention_fwd(**kwargs)
|
||||
|
||||
if backend == "triton":
|
||||
from sglang.srt.layers.attention.nsa.triton_decode import (
|
||||
triton_fp8_attention_fwd,
|
||||
)
|
||||
|
||||
return triton_fp8_attention_fwd(**kwargs)
|
||||
|
||||
if backend == "kernel":
|
||||
return flash_mla.flash_mla_with_kvcache(**kwargs)
|
||||
|
||||
|
||||
@@ -0,0 +1,98 @@
|
||||
"""
|
||||
Triton-based sparse attention decode kernels for DeepSeek V4.
|
||||
|
||||
This package provides an alternative to the tilelang implementation,
|
||||
controlled by the environment variable SGLANG_HACK_FLASHMLA_BACKEND=triton.
|
||||
"""
|
||||
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.attention.nsa.triton_decode.triton_mla_kernels_decode_optimized import (
|
||||
triton_sparse_attn_decode,
|
||||
)
|
||||
|
||||
|
||||
class _KVScopeAdapter:
|
||||
"""Lightweight adapter providing the kv_scope interface expected by
|
||||
``triton_sparse_attn_decode``.
|
||||
|
||||
The Triton kernels access four fields:
|
||||
* ``blocked_k_quantized`` – the raw FP8 KV cache tensor.
|
||||
* ``blocked_k`` – only ``blocked_k.shape[1]`` (block size)
|
||||
is read, so we reuse the same tensor.
|
||||
* ``indices_in_kvcache`` – sparse top-k page indices.
|
||||
* ``topk_length`` – valid length per batch element.
|
||||
"""
|
||||
|
||||
__slots__ = [
|
||||
"blocked_k",
|
||||
"blocked_k_quantized",
|
||||
"indices_in_kvcache",
|
||||
"topk_length",
|
||||
]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
k_cache: torch.Tensor,
|
||||
indices: torch.Tensor,
|
||||
topk_length: Optional[torch.Tensor],
|
||||
):
|
||||
self.blocked_k_quantized = k_cache
|
||||
self.blocked_k = k_cache
|
||||
self.indices_in_kvcache = indices
|
||||
self.topk_length = topk_length
|
||||
|
||||
|
||||
def triton_fp8_attention_fwd(
|
||||
q: torch.Tensor,
|
||||
k_cache: torch.Tensor,
|
||||
head_dim_v: int,
|
||||
softmax_scale: float,
|
||||
indices: torch.Tensor,
|
||||
attn_sink: Optional[torch.Tensor] = None,
|
||||
extra_k_cache: Optional[torch.Tensor] = None,
|
||||
extra_indices_in_kvcache: Optional[torch.Tensor] = None,
|
||||
topk_length: Optional[torch.Tensor] = None,
|
||||
extra_topk_length: Optional[torch.Tensor] = None,
|
||||
**_unused,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Sparse MLA decode via Triton kernels.
|
||||
|
||||
Accepts the same ``**kwargs`` dict that the caller builds for
|
||||
``flash_mla_with_kvcache`` / ``dpsk_v4_fp8_attention_fwd``, but only
|
||||
uses the subset of arguments relevant to the Triton implementation.
|
||||
Unused keys (``block_table``, ``cache_seqlens``,
|
||||
``tile_scheduler_metadata``, ``num_splits``, ``causal``,
|
||||
``is_fp8_kvcache``) are silently ignored via ``**_unused``.
|
||||
|
||||
Returns:
|
||||
``(output, lse)`` where *output* has shape
|
||||
``[batch, seq_len, num_heads, head_dim_v]`` and *lse* has shape
|
||||
``[batch, seq_len, num_heads]``.
|
||||
"""
|
||||
kv_scope = _KVScopeAdapter(k_cache, indices, topk_length)
|
||||
|
||||
extra_kv_scope = None
|
||||
if extra_k_cache is not None:
|
||||
extra_kv_scope = _KVScopeAdapter(
|
||||
extra_k_cache,
|
||||
extra_indices_in_kvcache,
|
||||
extra_topk_length,
|
||||
)
|
||||
|
||||
output, lse = triton_sparse_attn_decode(
|
||||
q=q,
|
||||
kv_scope=kv_scope,
|
||||
extra_kv_scope=extra_kv_scope,
|
||||
sm_scale=softmax_scale,
|
||||
d_v=head_dim_v,
|
||||
attn_sink=attn_sink,
|
||||
)
|
||||
|
||||
# Triton kernel returns lse as (b, h_q, s_q); transpose to
|
||||
# (b, s_q, h_q) to match the tilelang / flash_mla convention.
|
||||
lse = lse.transpose(1, 2)
|
||||
|
||||
return output, lse
|
||||
+585
@@ -0,0 +1,585 @@
|
||||
"""
|
||||
Common utilities and attention kernels for Triton MLA Decode.
|
||||
|
||||
This module contains shared code for the DeepSeek V4 Triton decode implementation:
|
||||
- Attention kernels (unified sparse decode)
|
||||
- Helper functions for chunked attention
|
||||
- Token range computation for memory-based chunking
|
||||
"""
|
||||
|
||||
from typing import List, Tuple
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
LOG2E = tl.constexpr(1.4426950408889634)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Bucketing for autotune keys to avoid recompilation per unique batch size
|
||||
# ============================================================================
|
||||
def _bucket_total_tokens(total_tokens: int) -> int:
|
||||
"""Round total_tokens up to the nearest power of 2 for autotune key stability.
|
||||
|
||||
In serving, total_tokens (= batch_size * seq_len) varies with every batch.
|
||||
Using the exact value as an autotune key causes recompilation for each unique
|
||||
value. Bucketing to powers of 2 limits the number of unique keys to ~15,
|
||||
dramatically reducing autotuning overhead.
|
||||
|
||||
Returns:
|
||||
Power-of-2 bucket: 1, 2, 4, 8, ..., up to the next power of 2.
|
||||
"""
|
||||
if total_tokens <= 0:
|
||||
return 1
|
||||
# Round up to next power of 2
|
||||
n = 1
|
||||
while n < total_tokens:
|
||||
n <<= 1
|
||||
return n
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Helper function to compute workload size category for autotune
|
||||
# ============================================================================
|
||||
def _get_workload_size_category(total_tokens: int, topk: int) -> int:
|
||||
"""
|
||||
Compute workload size category for autotune key.
|
||||
Returns:
|
||||
0: small (< 10K elements)
|
||||
1: medium (10K - 100K elements)
|
||||
2: large (100K - 1M elements)
|
||||
3: very large (> 1M elements)
|
||||
"""
|
||||
total_elements = total_tokens * topk
|
||||
if total_elements < 10000:
|
||||
return 0
|
||||
elif total_elements < 100000:
|
||||
return 1
|
||||
elif total_elements < 1000000:
|
||||
return 2
|
||||
else:
|
||||
return 3
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Unified Attention Kernels
|
||||
# ============================================================================
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# CDNA4 (gfx950) Optimized: Added high-performance configs for MI355X
|
||||
# Best config for h_q=128, large topk: BLOCK_H=64, BLOCK_N=256, num_warps=8
|
||||
# ============================================================================
|
||||
@triton.autotune(
|
||||
configs=[
|
||||
# Selected based on CDNA4 architecture analysis:
|
||||
# - BLOCK_D=128 is fixed (matches KV tile structure for d_qk=512).
|
||||
# - BLOCK_N=256: best for amortizing memory access over topk dimension.
|
||||
# (decode attention is memory-bound; larger BLOCK_N = fewer iterations)
|
||||
# - num_warps=8: memory-bound decode benefits from more warps for latency hiding.
|
||||
# - BLOCK_H varies to cover different batch sizes:
|
||||
# * BLOCK_H=16: cdiv(128,16)=8 H-blocks, best for small batches (bs=1-8)
|
||||
# * BLOCK_H=32: cdiv(128,32)=4 H-blocks, good for medium batches (bs=8-32)
|
||||
# * BLOCK_H=64: cdiv(128,64)=2 H-blocks, best for large batches (bs=32+)
|
||||
# (original comment: "Best for h_q=128, large topk")
|
||||
# * BLOCK_H=128: cdiv(128,128)=1 H-block, for very large batches (bs=128+)
|
||||
triton.Config(
|
||||
{"BLOCK_H": 16, "BLOCK_N": 256, "BLOCK_D": 128}, num_warps=8, num_stages=1
|
||||
),
|
||||
triton.Config(
|
||||
{"BLOCK_H": 32, "BLOCK_N": 256, "BLOCK_D": 128}, num_warps=8, num_stages=1
|
||||
),
|
||||
triton.Config(
|
||||
{"BLOCK_H": 64, "BLOCK_N": 256, "BLOCK_D": 128}, num_warps=8, num_stages=1
|
||||
),
|
||||
triton.Config(
|
||||
{"BLOCK_H": 128, "BLOCK_N": 256, "BLOCK_D": 128}, num_warps=8, num_stages=1
|
||||
),
|
||||
],
|
||||
key=["total_tokens_bucket", "h_q", "total_topk", "d_qk"],
|
||||
)
|
||||
@triton.jit
|
||||
def _unified_sparse_decode_kernel(
|
||||
Q,
|
||||
KV,
|
||||
Mask,
|
||||
AttnSink,
|
||||
Output,
|
||||
LSE,
|
||||
sm_scale,
|
||||
total_tokens,
|
||||
total_tokens_bucket,
|
||||
h_q,
|
||||
total_topk,
|
||||
d_qk,
|
||||
d_v,
|
||||
stride_q_t,
|
||||
stride_q_h,
|
||||
stride_q_d,
|
||||
stride_kv_t,
|
||||
stride_kv_k,
|
||||
stride_kv_d,
|
||||
stride_mask_t,
|
||||
stride_mask_k,
|
||||
stride_o_t,
|
||||
stride_o_h,
|
||||
stride_o_d,
|
||||
stride_lse_t,
|
||||
stride_lse_h,
|
||||
HAS_ATTN_SINK: tl.constexpr,
|
||||
BLOCK_H: tl.constexpr,
|
||||
BLOCK_N: tl.constexpr,
|
||||
BLOCK_D: tl.constexpr,
|
||||
):
|
||||
"""Unified attention kernel with single KV buffer (int64 safe)."""
|
||||
pid_t = tl.program_id(0)
|
||||
pid_h = tl.program_id(1)
|
||||
pid_t_64 = pid_t.to(tl.int64)
|
||||
|
||||
NEG_INF = float("-inf")
|
||||
POS_INF = float("+inf")
|
||||
|
||||
offs_h = pid_h * BLOCK_H + tl.arange(0, BLOCK_H)
|
||||
mask_h = offs_h < h_q
|
||||
|
||||
m_i = tl.full([BLOCK_H], NEG_INF, dtype=tl.float32)
|
||||
l_i = tl.zeros([BLOCK_H], dtype=tl.float32)
|
||||
|
||||
acc_0 = tl.zeros([BLOCK_H, BLOCK_D], dtype=tl.float32)
|
||||
acc_1 = tl.zeros([BLOCK_H, BLOCK_D], dtype=tl.float32)
|
||||
acc_2 = tl.zeros([BLOCK_H, BLOCK_D], dtype=tl.float32)
|
||||
acc_3 = tl.zeros([BLOCK_H, BLOCK_D], dtype=tl.float32)
|
||||
|
||||
stride_q_t_64 = tl.cast(stride_q_t, tl.int64)
|
||||
stride_kv_t_64 = tl.cast(stride_kv_t, tl.int64)
|
||||
stride_mask_t_64 = tl.cast(stride_mask_t, tl.int64)
|
||||
q_base = Q + pid_t_64 * stride_q_t_64
|
||||
kv_base = KV + pid_t_64 * stride_kv_t_64
|
||||
mask_base = Mask + pid_t_64 * stride_mask_t_64
|
||||
|
||||
for n_start in range(0, total_topk, BLOCK_N):
|
||||
offs_n = n_start + tl.arange(0, BLOCK_N)
|
||||
mask_n = offs_n < total_topk
|
||||
|
||||
mask_ptrs = mask_base + offs_n * stride_mask_k
|
||||
invalid = tl.load(mask_ptrs, mask=mask_n, other=True)
|
||||
valid = mask_n & ~invalid
|
||||
|
||||
qk = tl.zeros([BLOCK_H, BLOCK_N], dtype=tl.float32)
|
||||
|
||||
for d_start in range(0, d_qk, BLOCK_D):
|
||||
offs_d = d_start + tl.arange(0, BLOCK_D)
|
||||
mask_d = offs_d < d_qk
|
||||
|
||||
q_ptrs = (
|
||||
q_base + offs_h[:, None] * stride_q_h + offs_d[None, :] * stride_q_d
|
||||
)
|
||||
q_chunk = tl.load(
|
||||
q_ptrs, mask=mask_h[:, None] & mask_d[None, :], other=0.0
|
||||
).to(tl.bfloat16)
|
||||
|
||||
k_ptrs = (
|
||||
kv_base + offs_n[:, None] * stride_kv_k + offs_d[None, :] * stride_kv_d
|
||||
)
|
||||
k_chunk = tl.load(
|
||||
k_ptrs, mask=valid[:, None] & mask_d[None, :], other=0.0
|
||||
).to(tl.bfloat16)
|
||||
|
||||
qk += tl.dot(q_chunk, tl.trans(k_chunk))
|
||||
|
||||
qk = qk * sm_scale
|
||||
qk = tl.where(valid[None, :], qk, NEG_INF)
|
||||
|
||||
m_ij = tl.max(qk, axis=1)
|
||||
m_new = tl.maximum(m_i, m_ij)
|
||||
alpha = tl.where(m_i == NEG_INF, 0.0, tl.math.exp2((m_i - m_new) * LOG2E))
|
||||
p = tl.where(qk == NEG_INF, 0.0, tl.math.exp2((qk - m_new[:, None]) * LOG2E))
|
||||
l_new = alpha * l_i + tl.sum(p, axis=1)
|
||||
p_bf16 = p.to(tl.bfloat16)
|
||||
|
||||
offs_v = tl.arange(0, BLOCK_D)
|
||||
v_ptrs = kv_base + offs_n[:, None] * stride_kv_k + offs_v[None, :] * stride_kv_d
|
||||
v = tl.load(v_ptrs, mask=valid[:, None], other=0.0).to(tl.bfloat16)
|
||||
acc_0 = acc_0 * alpha[:, None] + tl.dot(p_bf16, v)
|
||||
|
||||
offs_v = BLOCK_D + tl.arange(0, BLOCK_D)
|
||||
v_ptrs = kv_base + offs_n[:, None] * stride_kv_k + offs_v[None, :] * stride_kv_d
|
||||
v = tl.load(
|
||||
v_ptrs, mask=valid[:, None] & (offs_v[None, :] < d_v), other=0.0
|
||||
).to(tl.bfloat16)
|
||||
acc_1 = acc_1 * alpha[:, None] + tl.dot(p_bf16, v)
|
||||
|
||||
offs_v = 2 * BLOCK_D + tl.arange(0, BLOCK_D)
|
||||
v_ptrs = kv_base + offs_n[:, None] * stride_kv_k + offs_v[None, :] * stride_kv_d
|
||||
v = tl.load(
|
||||
v_ptrs, mask=valid[:, None] & (offs_v[None, :] < d_v), other=0.0
|
||||
).to(tl.bfloat16)
|
||||
acc_2 = acc_2 * alpha[:, None] + tl.dot(p_bf16, v)
|
||||
|
||||
offs_v = 3 * BLOCK_D + tl.arange(0, BLOCK_D)
|
||||
v_ptrs = kv_base + offs_n[:, None] * stride_kv_k + offs_v[None, :] * stride_kv_d
|
||||
v = tl.load(
|
||||
v_ptrs, mask=valid[:, None] & (offs_v[None, :] < d_v), other=0.0
|
||||
).to(tl.bfloat16)
|
||||
acc_3 = acc_3 * alpha[:, None] + tl.dot(p_bf16, v)
|
||||
|
||||
m_i = m_new
|
||||
l_i = l_new
|
||||
|
||||
lse = m_i + tl.math.log2(tl.where(l_i == 0.0, 1.0, l_i)) / LOG2E
|
||||
is_lonely_q = l_i == 0.0
|
||||
|
||||
if HAS_ATTN_SINK:
|
||||
attn_sink_vals = tl.load(AttnSink + offs_h, mask=mask_h, other=0.0)
|
||||
exp_attn_sink_minus_m = tl.math.exp2((attn_sink_vals - m_i) * LOG2E)
|
||||
denominator = l_i + exp_attn_sink_minus_m
|
||||
denominator = tl.where(denominator == 0.0, 1.0, denominator)
|
||||
output_scale = 1.0 / denominator
|
||||
else:
|
||||
output_scale = tl.where(l_i == 0.0, 0.0, 1.0 / l_i)
|
||||
|
||||
# Pre-compute 2D versions for efficiency
|
||||
is_lonely_q_2d = is_lonely_q[:, None]
|
||||
output_scale_2d = output_scale[:, None]
|
||||
acc_0 = tl.where(is_lonely_q_2d, 0.0, acc_0 * output_scale_2d)
|
||||
acc_1 = tl.where(is_lonely_q_2d, 0.0, acc_1 * output_scale_2d)
|
||||
acc_2 = tl.where(is_lonely_q_2d, 0.0, acc_2 * output_scale_2d)
|
||||
acc_3 = tl.where(is_lonely_q_2d, 0.0, acc_3 * output_scale_2d)
|
||||
lse = tl.where(is_lonely_q, POS_INF, lse)
|
||||
|
||||
stride_lse_t_64 = tl.cast(stride_lse_t, tl.int64)
|
||||
tl.store(LSE + pid_t_64 * stride_lse_t_64 + offs_h * stride_lse_h, lse, mask=mask_h)
|
||||
|
||||
stride_o_t_64 = tl.cast(stride_o_t, tl.int64)
|
||||
o_base = Output + pid_t_64 * stride_o_t_64
|
||||
# Pre-compute 2D versions
|
||||
offs_h_2d = offs_h[:, None]
|
||||
mask_h_2d = mask_h[:, None]
|
||||
offs_v_0 = tl.arange(0, BLOCK_D)
|
||||
offs_v_1 = BLOCK_D + tl.arange(0, BLOCK_D)
|
||||
offs_v_2 = 2 * BLOCK_D + tl.arange(0, BLOCK_D)
|
||||
offs_v_3 = 3 * BLOCK_D + tl.arange(0, BLOCK_D)
|
||||
tl.store(
|
||||
o_base + offs_h_2d * stride_o_h + offs_v_0[None, :] * stride_o_d,
|
||||
acc_0.to(tl.bfloat16),
|
||||
mask=mask_h_2d,
|
||||
)
|
||||
tl.store(
|
||||
o_base + offs_h_2d * stride_o_h + offs_v_1[None, :] * stride_o_d,
|
||||
acc_1.to(tl.bfloat16),
|
||||
mask=mask_h_2d & (offs_v_1[None, :] < d_v),
|
||||
)
|
||||
tl.store(
|
||||
o_base + offs_h_2d * stride_o_h + offs_v_2[None, :] * stride_o_d,
|
||||
acc_2.to(tl.bfloat16),
|
||||
mask=mask_h_2d & (offs_v_2[None, :] < d_v),
|
||||
)
|
||||
tl.store(
|
||||
o_base + offs_h_2d * stride_o_h + offs_v_3[None, :] * stride_o_d,
|
||||
acc_3.to(tl.bfloat16),
|
||||
mask=mask_h_2d & (offs_v_3[None, :] < d_v),
|
||||
)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Attention Runner Functions
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def run_unified_attention(
|
||||
q_reshaped,
|
||||
gathered_kv,
|
||||
invalid_mask,
|
||||
d_v,
|
||||
sm_scale,
|
||||
total_tokens,
|
||||
h_q,
|
||||
total_topk,
|
||||
d_qk,
|
||||
attn_sink=None,
|
||||
):
|
||||
"""Run unified attention with single KV buffer.
|
||||
|
||||
Run unified sparse decode attention kernel.
|
||||
"""
|
||||
output = torch.empty(
|
||||
(total_tokens, h_q, d_v), dtype=torch.bfloat16, device=q_reshaped.device
|
||||
)
|
||||
lse = torch.empty(
|
||||
(total_tokens, h_q), dtype=torch.float32, device=q_reshaped.device
|
||||
)
|
||||
|
||||
HAS_ATTN_SINK = attn_sink is not None
|
||||
attn_sink_tensor = attn_sink if HAS_ATTN_SINK else lse[:1]
|
||||
|
||||
grid = lambda meta: (total_tokens, triton.cdiv(h_q, meta["BLOCK_H"]))
|
||||
_unified_sparse_decode_kernel[grid](
|
||||
q_reshaped,
|
||||
gathered_kv,
|
||||
invalid_mask,
|
||||
attn_sink_tensor,
|
||||
output,
|
||||
lse,
|
||||
sm_scale,
|
||||
total_tokens,
|
||||
_bucket_total_tokens(total_tokens),
|
||||
h_q,
|
||||
total_topk,
|
||||
d_qk,
|
||||
d_v,
|
||||
q_reshaped.stride(0),
|
||||
q_reshaped.stride(1),
|
||||
q_reshaped.stride(2),
|
||||
gathered_kv.stride(0),
|
||||
gathered_kv.stride(1),
|
||||
gathered_kv.stride(2),
|
||||
invalid_mask.stride(0),
|
||||
invalid_mask.stride(1),
|
||||
output.stride(0),
|
||||
output.stride(1),
|
||||
output.stride(2),
|
||||
lse.stride(0),
|
||||
lse.stride(1),
|
||||
HAS_ATTN_SINK=HAS_ATTN_SINK,
|
||||
)
|
||||
return output, lse
|
||||
|
||||
|
||||
def run_chunked_attention_triton(
|
||||
q_reshaped,
|
||||
gathered_kv,
|
||||
invalid_mask,
|
||||
d_v,
|
||||
sm_scale,
|
||||
total_tokens,
|
||||
h_q,
|
||||
total_topk,
|
||||
d_qk,
|
||||
attn_sink=None,
|
||||
chunk_size=8192,
|
||||
):
|
||||
"""Chunked attention using Triton kernels with cross-chunk softmax merging."""
|
||||
device = q_reshaped.device
|
||||
|
||||
num_chunks = (total_topk + chunk_size - 1) // chunk_size
|
||||
|
||||
kv_chunks = []
|
||||
mask_chunks = []
|
||||
chunk_sizes = []
|
||||
|
||||
for chunk_idx in range(num_chunks):
|
||||
start_k = chunk_idx * chunk_size
|
||||
end_k = min(start_k + chunk_size, total_topk)
|
||||
chunk_topk = end_k - start_k
|
||||
chunk_sizes.append(chunk_topk)
|
||||
kv_chunks.append(gathered_kv[:, start_k:end_k, :].contiguous())
|
||||
mask_chunks.append(invalid_mask[:, start_k:end_k].contiguous())
|
||||
|
||||
lse_acc = torch.full(
|
||||
(total_tokens, h_q), float("-inf"), dtype=torch.float32, device=device
|
||||
)
|
||||
acc = torch.zeros((total_tokens, h_q, d_v), dtype=torch.float32, device=device)
|
||||
|
||||
for chunk_idx in range(num_chunks):
|
||||
kv_chunk = kv_chunks[chunk_idx]
|
||||
mask_chunk = mask_chunks[chunk_idx]
|
||||
chunk_topk = chunk_sizes[chunk_idx]
|
||||
|
||||
chunk_output, chunk_lse = run_unified_attention(
|
||||
q_reshaped,
|
||||
kv_chunk,
|
||||
mask_chunk,
|
||||
d_v,
|
||||
sm_scale,
|
||||
total_tokens,
|
||||
h_q,
|
||||
chunk_topk,
|
||||
d_qk,
|
||||
attn_sink=None,
|
||||
)
|
||||
|
||||
is_chunk_lonely = torch.isinf(chunk_lse) & (chunk_lse > 0)
|
||||
|
||||
chunk_lse_for_merge = torch.where(
|
||||
is_chunk_lonely, torch.full_like(chunk_lse, float("-inf")), chunk_lse
|
||||
)
|
||||
|
||||
lse_max = torch.maximum(lse_acc, chunk_lse_for_merge)
|
||||
|
||||
exp_acc = torch.exp(lse_acc - lse_max)
|
||||
exp_acc = torch.where(torch.isnan(exp_acc), torch.zeros_like(exp_acc), exp_acc)
|
||||
|
||||
exp_chunk = torch.exp(chunk_lse_for_merge - lse_max)
|
||||
exp_chunk = torch.where(
|
||||
torch.isnan(exp_chunk) | is_chunk_lonely,
|
||||
torch.zeros_like(exp_chunk),
|
||||
exp_chunk,
|
||||
)
|
||||
|
||||
sum_exp = exp_acc + exp_chunk
|
||||
lse_new = lse_max + torch.log(
|
||||
torch.where(sum_exp == 0, torch.ones_like(sum_exp), sum_exp)
|
||||
)
|
||||
|
||||
both_empty = (lse_acc == float("-inf")) & (chunk_lse_for_merge == float("-inf"))
|
||||
lse_new = torch.where(
|
||||
both_empty, torch.full_like(lse_new, float("-inf")), lse_new
|
||||
)
|
||||
|
||||
weight_acc = torch.exp(lse_acc - lse_new)
|
||||
weight_acc = torch.where(
|
||||
torch.isnan(weight_acc) | torch.isinf(weight_acc),
|
||||
torch.zeros_like(weight_acc),
|
||||
weight_acc,
|
||||
)
|
||||
|
||||
weight_chunk = torch.exp(chunk_lse_for_merge - lse_new)
|
||||
weight_chunk = torch.where(
|
||||
torch.isnan(weight_chunk) | torch.isinf(weight_chunk) | is_chunk_lonely,
|
||||
torch.zeros_like(weight_chunk),
|
||||
weight_chunk,
|
||||
)
|
||||
|
||||
acc = (
|
||||
weight_acc.unsqueeze(-1) * acc
|
||||
+ weight_chunk.unsqueeze(-1) * chunk_output.float()
|
||||
)
|
||||
|
||||
lse_acc = lse_new
|
||||
|
||||
output = acc
|
||||
lse = lse_acc
|
||||
|
||||
is_lonely_final = lse == float("-inf")
|
||||
|
||||
lse = torch.where(is_lonely_final, torch.full_like(lse, float("+inf")), lse)
|
||||
|
||||
if attn_sink is not None:
|
||||
attn_sink_expanded = attn_sink.view(1, h_q)
|
||||
exp_diff = torch.exp(attn_sink_expanded - lse)
|
||||
exp_diff = torch.where(
|
||||
is_lonely_final, torch.full_like(exp_diff, float("inf")), exp_diff
|
||||
)
|
||||
scale = 1.0 / (1.0 + exp_diff)
|
||||
output = output * scale.unsqueeze(-1)
|
||||
|
||||
output = torch.where(
|
||||
is_lonely_final.unsqueeze(-1), torch.zeros_like(output), output
|
||||
)
|
||||
|
||||
return output.to(torch.bfloat16), lse
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Helper class and functions for token-range based chunking
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class SlicedKVScope:
|
||||
"""A sliced view of KV scope for a specific token range."""
|
||||
|
||||
__slots__ = [
|
||||
"blocked_k",
|
||||
"blocked_k_quantized",
|
||||
"indices_in_kvcache",
|
||||
"topk_length",
|
||||
]
|
||||
|
||||
def __init__(self, blocked_k, blocked_k_quantized, indices_in_kvcache, topk_length):
|
||||
self.blocked_k = blocked_k
|
||||
self.blocked_k_quantized = blocked_k_quantized
|
||||
self.indices_in_kvcache = indices_in_kvcache
|
||||
self.topk_length = topk_length
|
||||
|
||||
|
||||
def slice_kv_scope_for_tokens(orig_scope, start_t: int, end_t: int, s_q: int):
|
||||
"""Slice a KV scope to only include tokens in range [start_t, end_t)."""
|
||||
if orig_scope is None:
|
||||
return None
|
||||
|
||||
orig_indices = orig_scope.indices_in_kvcache.reshape(
|
||||
-1, orig_scope.indices_in_kvcache.size(-1)
|
||||
)
|
||||
sliced_indices = orig_indices[start_t:end_t]
|
||||
|
||||
sliced_topk_length = None
|
||||
if orig_scope.topk_length is not None:
|
||||
batch_start = start_t // s_q
|
||||
batch_end = (end_t + s_q - 1) // s_q
|
||||
batch_topk_length = orig_scope.topk_length[batch_start:batch_end]
|
||||
if s_q > 1:
|
||||
chunk_tokens = end_t - start_t
|
||||
expanded = batch_topk_length.unsqueeze(1).expand(-1, s_q).reshape(-1)
|
||||
offset_in_first_batch = start_t % s_q
|
||||
sliced_topk_length = expanded[
|
||||
offset_in_first_batch : offset_in_first_batch + chunk_tokens
|
||||
]
|
||||
else:
|
||||
sliced_topk_length = batch_topk_length
|
||||
|
||||
return SlicedKVScope(
|
||||
blocked_k=orig_scope.blocked_k,
|
||||
blocked_k_quantized=orig_scope.blocked_k_quantized,
|
||||
indices_in_kvcache=sliced_indices,
|
||||
topk_length=sliced_topk_length,
|
||||
)
|
||||
|
||||
|
||||
def compute_token_ranges(
|
||||
total_tokens: int,
|
||||
total_topk: int,
|
||||
d_qk: int,
|
||||
max_buffer_bytes: int = 2 * 1024 * 1024 * 1024,
|
||||
) -> List[Tuple[int, int]]:
|
||||
"""Compute token ranges for processing, chunking if buffer would exceed limit."""
|
||||
buffer_size_bytes = total_tokens * total_topk * d_qk * 2
|
||||
|
||||
if buffer_size_bytes <= max_buffer_bytes:
|
||||
return [(0, total_tokens)]
|
||||
|
||||
max_tokens_per_chunk = max_buffer_bytes // (total_topk * d_qk * 2)
|
||||
chunk_size = max(1, max_tokens_per_chunk)
|
||||
|
||||
token_ranges = []
|
||||
start_t = 0
|
||||
while start_t < total_tokens:
|
||||
end_t = min(start_t + chunk_size, total_tokens)
|
||||
token_ranges.append((start_t, end_t))
|
||||
start_t = end_t
|
||||
|
||||
return token_ranges
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Split-K Attention for Large TopK
|
||||
# ============================================================================
|
||||
def run_splitk_unified_attention(
|
||||
q_reshaped,
|
||||
gathered_kv,
|
||||
invalid_mask,
|
||||
d_v,
|
||||
sm_scale,
|
||||
total_tokens,
|
||||
h_q,
|
||||
total_topk,
|
||||
d_qk,
|
||||
attn_sink=None,
|
||||
split_k=4,
|
||||
):
|
||||
"""Run split-K attention for large topk cases."""
|
||||
from .triton_mla_kernels_decode_splitk import run_splitk_attention
|
||||
|
||||
return run_splitk_attention(
|
||||
q_reshaped,
|
||||
gathered_kv,
|
||||
invalid_mask,
|
||||
d_v,
|
||||
sm_scale,
|
||||
total_tokens,
|
||||
h_q,
|
||||
total_topk,
|
||||
d_qk,
|
||||
attn_sink=attn_sink,
|
||||
split_k=split_k,
|
||||
)
|
||||
+1355
File diff suppressed because it is too large
Load Diff
+3089
File diff suppressed because it is too large
Load Diff
+289
@@ -0,0 +1,289 @@
|
||||
"""
|
||||
Optimized Triton MLA Decode Kernels for DeepSeek V4.
|
||||
|
||||
This module provides optimized sparse attention decode with reduced Python overhead.
|
||||
|
||||
Key optimizations:
|
||||
1. Fused gather+dequant+attention kernels (eliminates intermediate buffers)
|
||||
2. Split-K for better GPU parallelism on small batches
|
||||
3. Pre-allocated buffer pool for splitk intermediate results
|
||||
4. Pre-computed strides to reduce tensor metadata operations
|
||||
|
||||
Note: This implementation assumes KV cache is always FP8 quantized.
|
||||
"""
|
||||
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import torch
|
||||
import triton
|
||||
|
||||
from .triton_mla_kernels_decode_common import (
|
||||
_bucket_total_tokens,
|
||||
_unified_sparse_decode_kernel,
|
||||
compute_token_ranges,
|
||||
)
|
||||
from .triton_mla_kernels_decode_dsv4 import (
|
||||
DSV4_D_QK,
|
||||
fused_gather_dequant_fp8_dsv4,
|
||||
)
|
||||
from .triton_mla_kernels_decode_fused import (
|
||||
fused_gather_attn_decode_dsv4,
|
||||
fused_gather_attn_decode_dsv4_dual_scope_low_overhead,
|
||||
)
|
||||
|
||||
|
||||
def triton_sparse_attn_decode(
|
||||
q: torch.Tensor,
|
||||
kv_scope,
|
||||
extra_kv_scope,
|
||||
sm_scale: float,
|
||||
d_v: int = 512,
|
||||
attn_sink: Optional[torch.Tensor] = None,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Optimized sparse attention decode for DeepSeek V4 (d_qk=512)."""
|
||||
d_qk = q.shape[-1]
|
||||
|
||||
if d_qk != DSV4_D_QK:
|
||||
raise ValueError(
|
||||
f"Unsupported d_qk: {d_qk}. Expected {DSV4_D_QK} (DeepSeek V4)"
|
||||
)
|
||||
|
||||
return _triton_sparse_attn_decode_dsv4(
|
||||
q, kv_scope, extra_kv_scope, sm_scale, d_v, attn_sink
|
||||
)
|
||||
|
||||
|
||||
def _should_use_fused_dual_scope(total_tokens: int, h_q: int, total_topk: int) -> bool:
|
||||
"""Determine whether to use fused kernel for dual-scope cases.
|
||||
|
||||
The fused kernel avoids allocating a large intermediate gathered_kv
|
||||
buffer and eliminates a separate gather kernel launch. However, for
|
||||
h_q > 64 with medium-to-large batch sizes and larger topk, the
|
||||
non-splitk fused kernel suffers from low GPU utilization (the grid
|
||||
has only cdiv(h_q, BLOCK_H) blocks in the H dimension). In those
|
||||
cases the fallback (separate gather + attention) can be faster on
|
||||
the GPU, though it incurs extra torch.empty() overhead in CUDA
|
||||
graphs.
|
||||
|
||||
The thresholds below were determined empirically on MI355X (256 CUs).
|
||||
"""
|
||||
if total_tokens <= 4:
|
||||
return True
|
||||
if h_q <= 64 and total_topk <= 800:
|
||||
return total_tokens <= 256
|
||||
if h_q <= 64 and total_topk >= 1024:
|
||||
return total_tokens <= 128
|
||||
# h_q > 64 (e.g. h_q=128 when q is padded to full n_heads).
|
||||
# For small topk (c128 layers, topk~192), fused always wins.
|
||||
# For larger topk (c4 layers, topk~640), fused wins at small bs
|
||||
# but the fallback catches up at bs>=16 due to better GPU utilization.
|
||||
# However, the fallback has 4 extra torch.empty() calls that add
|
||||
# ~30us CUDA-graph replay overhead, roughly cancelling the GPU gain.
|
||||
# So we route to fused for all practical batch sizes.
|
||||
if h_q > 64:
|
||||
return total_tokens <= 256
|
||||
return True
|
||||
|
||||
|
||||
def _triton_sparse_attn_decode_dsv4(
|
||||
q: torch.Tensor,
|
||||
kv_scope,
|
||||
extra_kv_scope,
|
||||
sm_scale: float,
|
||||
d_v: int,
|
||||
attn_sink: Optional[torch.Tensor],
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Optimized sparse attention decode for DeepSeek V4 (d_qk=512)."""
|
||||
b, s_q, h_q, d_qk = q.shape
|
||||
total_tokens = b * s_q
|
||||
device = q.device
|
||||
|
||||
topk_main = kv_scope.indices_in_kvcache.shape[-1]
|
||||
kv_quantized_main = kv_scope.blocked_k_quantized
|
||||
block_size_main = kv_scope.blocked_k.shape[1]
|
||||
|
||||
# Single scope case
|
||||
if extra_kv_scope is None:
|
||||
if topk_main < 8192:
|
||||
q_reshaped = q.reshape(total_tokens, h_q, d_qk)
|
||||
if not q_reshaped.is_contiguous():
|
||||
q_reshaped = q_reshaped.contiguous()
|
||||
|
||||
indices_main = kv_scope.indices_in_kvcache.reshape(total_tokens, topk_main)
|
||||
if not indices_main.is_contiguous():
|
||||
indices_main = indices_main.contiguous()
|
||||
|
||||
output, lse = fused_gather_attn_decode_dsv4(
|
||||
q_reshaped,
|
||||
kv_quantized_main,
|
||||
indices_main,
|
||||
block_size_main,
|
||||
sm_scale,
|
||||
topk_length=kv_scope.topk_length,
|
||||
attn_sink=attn_sink,
|
||||
s_q=s_q,
|
||||
)
|
||||
return output.view(b, s_q, h_q, d_v), lse.view(b, s_q, h_q).transpose(1, 2)
|
||||
else:
|
||||
from .triton_mla_kernels_decode_dsv4 import triton_sparse_attn_decode_dsv4
|
||||
|
||||
return triton_sparse_attn_decode_dsv4(
|
||||
q, kv_scope, extra_kv_scope, sm_scale, d_v, attn_sink
|
||||
)
|
||||
|
||||
# Dual scope case
|
||||
topk_extra = extra_kv_scope.indices_in_kvcache.shape[-1]
|
||||
total_topk = topk_main + topk_extra
|
||||
|
||||
# Check if chunking needed (fall back to original implementation)
|
||||
token_ranges = compute_token_ranges(total_tokens, total_topk, d_qk)
|
||||
if len(token_ranges) > 1:
|
||||
from .triton_mla_kernels_decode_dsv4 import triton_sparse_attn_decode_dsv4
|
||||
|
||||
return triton_sparse_attn_decode_dsv4(
|
||||
q, kv_scope, extra_kv_scope, sm_scale, d_v, attn_sink
|
||||
)
|
||||
|
||||
# Use fused dual-scope kernel with low-overhead buffer pool
|
||||
if _should_use_fused_dual_scope(total_tokens, h_q, total_topk):
|
||||
q_reshaped = q.reshape(total_tokens, h_q, d_qk)
|
||||
if not q_reshaped.is_contiguous():
|
||||
q_reshaped = q_reshaped.contiguous()
|
||||
|
||||
indices_main = kv_scope.indices_in_kvcache.reshape(total_tokens, topk_main)
|
||||
if not indices_main.is_contiguous():
|
||||
indices_main = indices_main.contiguous()
|
||||
|
||||
block_size_extra = extra_kv_scope.blocked_k.shape[1]
|
||||
indices_extra = extra_kv_scope.indices_in_kvcache.reshape(
|
||||
total_tokens, topk_extra
|
||||
)
|
||||
if not indices_extra.is_contiguous():
|
||||
indices_extra = indices_extra.contiguous()
|
||||
|
||||
output, lse = fused_gather_attn_decode_dsv4_dual_scope_low_overhead(
|
||||
q_reshaped,
|
||||
kv_quantized_main,
|
||||
indices_main,
|
||||
block_size_main,
|
||||
extra_kv_scope.blocked_k_quantized,
|
||||
indices_extra,
|
||||
block_size_extra,
|
||||
sm_scale,
|
||||
topk_length_main=kv_scope.topk_length,
|
||||
topk_length_extra=extra_kv_scope.topk_length,
|
||||
attn_sink=attn_sink,
|
||||
s_q=s_q,
|
||||
)
|
||||
return output.view(b, s_q, h_q, d_v), lse.view(b, s_q, h_q).transpose(1, 2)
|
||||
|
||||
# Fallback: Separate gather + attention path
|
||||
return _fallback_gather_attention(
|
||||
q,
|
||||
kv_scope,
|
||||
extra_kv_scope,
|
||||
sm_scale,
|
||||
d_v,
|
||||
attn_sink,
|
||||
total_tokens,
|
||||
h_q,
|
||||
d_qk,
|
||||
topk_main,
|
||||
topk_extra,
|
||||
block_size_main,
|
||||
kv_quantized_main,
|
||||
fused_gather_dequant_fp8_dsv4,
|
||||
)
|
||||
|
||||
|
||||
def _fallback_gather_attention(
|
||||
q: torch.Tensor,
|
||||
kv_scope,
|
||||
extra_kv_scope,
|
||||
sm_scale: float,
|
||||
d_v: int,
|
||||
attn_sink: Optional[torch.Tensor],
|
||||
total_tokens: int,
|
||||
h_q: int,
|
||||
d_qk: int,
|
||||
topk_main: int,
|
||||
topk_extra: int,
|
||||
block_size_main: int,
|
||||
kv_quantized_main,
|
||||
fused_gather_fn,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Fallback path: separate gather + attention kernels."""
|
||||
b = q.shape[0]
|
||||
s_q = q.shape[1]
|
||||
device = q.device
|
||||
total_topk = topk_main + topk_extra
|
||||
|
||||
gathered_kv = torch.empty(
|
||||
total_tokens, total_topk, d_qk, dtype=torch.bfloat16, device=device
|
||||
)
|
||||
invalid_mask = torch.empty(
|
||||
total_tokens, total_topk, dtype=torch.bool, device=device
|
||||
)
|
||||
output = torch.empty(total_tokens, h_q, d_v, dtype=torch.bfloat16, device=device)
|
||||
lse = torch.empty(total_tokens, h_q, dtype=torch.float32, device=device)
|
||||
|
||||
indices_main = kv_scope.indices_in_kvcache.reshape(total_tokens, topk_main)
|
||||
block_size_extra = extra_kv_scope.blocked_k.shape[1]
|
||||
indices_extra = extra_kv_scope.indices_in_kvcache.reshape(total_tokens, topk_extra)
|
||||
|
||||
fused_gather_fn(
|
||||
kv_quantized_main,
|
||||
indices_main,
|
||||
block_size_main,
|
||||
kv_scope.topk_length,
|
||||
extra_kv_scope.blocked_k_quantized,
|
||||
indices_extra,
|
||||
block_size_extra,
|
||||
extra_kv_scope.topk_length,
|
||||
gathered_kv,
|
||||
invalid_mask,
|
||||
s_q,
|
||||
)
|
||||
|
||||
if q.dtype == torch.bfloat16 and q.is_contiguous():
|
||||
q_reshaped = q.view(total_tokens, h_q, d_qk)
|
||||
else:
|
||||
q_reshaped = q.to(torch.bfloat16).reshape(total_tokens, h_q, d_qk)
|
||||
if not q_reshaped.is_contiguous():
|
||||
q_reshaped = q_reshaped.contiguous()
|
||||
|
||||
HAS_ATTN_SINK = attn_sink is not None
|
||||
attn_sink_tensor = attn_sink if HAS_ATTN_SINK else lse[:1]
|
||||
|
||||
grid = lambda meta: (total_tokens, triton.cdiv(h_q, meta["BLOCK_H"]))
|
||||
_unified_sparse_decode_kernel[grid](
|
||||
q_reshaped,
|
||||
gathered_kv,
|
||||
invalid_mask,
|
||||
attn_sink_tensor,
|
||||
output,
|
||||
lse,
|
||||
sm_scale,
|
||||
total_tokens,
|
||||
_bucket_total_tokens(total_tokens),
|
||||
h_q,
|
||||
total_topk,
|
||||
d_qk,
|
||||
d_v,
|
||||
q_reshaped.stride(0),
|
||||
q_reshaped.stride(1),
|
||||
q_reshaped.stride(2),
|
||||
gathered_kv.stride(0),
|
||||
gathered_kv.stride(1),
|
||||
gathered_kv.stride(2),
|
||||
invalid_mask.stride(0),
|
||||
invalid_mask.stride(1),
|
||||
output.stride(0),
|
||||
output.stride(1),
|
||||
output.stride(2),
|
||||
lse.stride(0),
|
||||
lse.stride(1),
|
||||
HAS_ATTN_SINK=HAS_ATTN_SINK,
|
||||
)
|
||||
|
||||
return output.view(b, s_q, h_q, d_v), lse.view(b, s_q, h_q).transpose(1, 2)
|
||||
+534
@@ -0,0 +1,534 @@
|
||||
"""
|
||||
Split-K Attention Kernel for Large TopK Cases
|
||||
|
||||
This module implements a split-K version of the attention kernel that:
|
||||
1. Splits the K (topk) dimension across multiple kernel instances
|
||||
2. Each instance computes partial results with its own m_i, l_i, and accumulators
|
||||
3. A combine kernel merges the partial results using online softmax
|
||||
|
||||
This reduces register pressure by processing fewer K tokens per kernel instance,
|
||||
improving occupancy and overall performance for large topk cases.
|
||||
"""
|
||||
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
from .triton_mla_kernels_decode_common import _bucket_total_tokens
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Split-K Attention Kernel
|
||||
# ============================================================================
|
||||
@triton.autotune(
|
||||
configs=[
|
||||
# Split-K attention on already-gathered BF16 KV.
|
||||
# - BLOCK_N=256: amortizes memory access over KV tokens (memory-bound kernel).
|
||||
# - BLOCK_D=128: matches KV tile structure.
|
||||
# - num_warps=8, num_stages=2: memory-bound kernel benefits from more warps
|
||||
# and software pipelining (overlaps memory loads with compute).
|
||||
# - BLOCK_H varies for different batch sizes:
|
||||
triton.Config(
|
||||
{"BLOCK_H": 16, "BLOCK_N": 256, "BLOCK_D": 128}, num_warps=8, num_stages=2
|
||||
),
|
||||
triton.Config(
|
||||
{"BLOCK_H": 32, "BLOCK_N": 256, "BLOCK_D": 128}, num_warps=8, num_stages=2
|
||||
),
|
||||
triton.Config(
|
||||
{"BLOCK_H": 64, "BLOCK_N": 256, "BLOCK_D": 128}, num_warps=8, num_stages=2
|
||||
),
|
||||
triton.Config(
|
||||
{"BLOCK_H": 128, "BLOCK_N": 256, "BLOCK_D": 128}, num_warps=8, num_stages=2
|
||||
),
|
||||
],
|
||||
key=["total_tokens_bucket", "h_q", "topk_per_split", "d_qk"],
|
||||
)
|
||||
@triton.jit
|
||||
def _splitk_attention_kernel(
|
||||
Q,
|
||||
KV,
|
||||
Mask,
|
||||
PartialOutput,
|
||||
PartialLSE,
|
||||
PartialM,
|
||||
sm_scale,
|
||||
total_tokens,
|
||||
total_tokens_bucket,
|
||||
h_q,
|
||||
total_topk,
|
||||
d_qk,
|
||||
d_v,
|
||||
topk_per_split,
|
||||
stride_q_t,
|
||||
stride_q_h,
|
||||
stride_q_d,
|
||||
stride_kv_t,
|
||||
stride_kv_k,
|
||||
stride_kv_d,
|
||||
stride_mask_t,
|
||||
stride_mask_k,
|
||||
stride_po_s,
|
||||
stride_po_t,
|
||||
stride_po_h,
|
||||
stride_po_d,
|
||||
stride_plse_s,
|
||||
stride_plse_t,
|
||||
stride_plse_h,
|
||||
stride_pm_s,
|
||||
stride_pm_t,
|
||||
stride_pm_h,
|
||||
BLOCK_H: tl.constexpr,
|
||||
BLOCK_N: tl.constexpr,
|
||||
BLOCK_D: tl.constexpr,
|
||||
):
|
||||
"""Split-K attention kernel that processes a subset of K tokens."""
|
||||
LOG2E: tl.constexpr = 1.4426950408889634
|
||||
|
||||
pid_t = tl.program_id(0)
|
||||
pid_h = tl.program_id(1)
|
||||
pid_k = tl.program_id(2)
|
||||
pid_t_64 = pid_t.to(tl.int64)
|
||||
|
||||
NEG_INF = float("-inf")
|
||||
|
||||
offs_h = pid_h * BLOCK_H + tl.arange(0, BLOCK_H)
|
||||
mask_h = offs_h < h_q
|
||||
|
||||
# Compute K range for this split
|
||||
k_start = pid_k * topk_per_split
|
||||
k_end = tl.minimum(k_start + topk_per_split, total_topk)
|
||||
|
||||
m_i = tl.full([BLOCK_H], NEG_INF, dtype=tl.float32)
|
||||
l_i = tl.zeros([BLOCK_H], dtype=tl.float32)
|
||||
|
||||
acc_0 = tl.zeros([BLOCK_H, BLOCK_D], dtype=tl.float32)
|
||||
acc_1 = tl.zeros([BLOCK_H, BLOCK_D], dtype=tl.float32)
|
||||
acc_2 = tl.zeros([BLOCK_H, BLOCK_D], dtype=tl.float32)
|
||||
acc_3 = tl.zeros([BLOCK_H, BLOCK_D], dtype=tl.float32)
|
||||
|
||||
stride_q_t_64 = tl.cast(stride_q_t, tl.int64)
|
||||
stride_kv_t_64 = tl.cast(stride_kv_t, tl.int64)
|
||||
stride_mask_t_64 = tl.cast(stride_mask_t, tl.int64)
|
||||
q_base = Q + pid_t_64 * stride_q_t_64
|
||||
kv_base = KV + pid_t_64 * stride_kv_t_64
|
||||
mask_base = Mask + pid_t_64 * stride_mask_t_64
|
||||
|
||||
for n_start in range(k_start, k_end, BLOCK_N):
|
||||
offs_n = n_start + tl.arange(0, BLOCK_N)
|
||||
mask_n = offs_n < k_end
|
||||
|
||||
mask_ptrs = mask_base + offs_n * stride_mask_k
|
||||
invalid = tl.load(mask_ptrs, mask=mask_n, other=True)
|
||||
valid = mask_n & ~invalid
|
||||
|
||||
qk = tl.zeros([BLOCK_H, BLOCK_N], dtype=tl.float32)
|
||||
|
||||
for d_start in range(0, d_qk, BLOCK_D):
|
||||
offs_d = d_start + tl.arange(0, BLOCK_D)
|
||||
mask_d = offs_d < d_qk
|
||||
|
||||
q_ptrs = (
|
||||
q_base + offs_h[:, None] * stride_q_h + offs_d[None, :] * stride_q_d
|
||||
)
|
||||
q_chunk = tl.load(
|
||||
q_ptrs, mask=mask_h[:, None] & mask_d[None, :], other=0.0
|
||||
).to(tl.bfloat16)
|
||||
|
||||
k_ptrs = (
|
||||
kv_base + offs_n[:, None] * stride_kv_k + offs_d[None, :] * stride_kv_d
|
||||
)
|
||||
k_chunk = tl.load(
|
||||
k_ptrs, mask=valid[:, None] & mask_d[None, :], other=0.0
|
||||
).to(tl.bfloat16)
|
||||
|
||||
qk += tl.dot(q_chunk, tl.trans(k_chunk))
|
||||
|
||||
qk = qk * sm_scale
|
||||
qk = tl.where(valid[None, :], qk, NEG_INF)
|
||||
|
||||
m_ij = tl.max(qk, axis=1)
|
||||
m_new = tl.maximum(m_i, m_ij)
|
||||
alpha = tl.where(m_i == NEG_INF, 0.0, tl.math.exp2((m_i - m_new) * LOG2E))
|
||||
p = tl.where(qk == NEG_INF, 0.0, tl.math.exp2((qk - m_new[:, None]) * LOG2E))
|
||||
l_new = alpha * l_i + tl.sum(p, axis=1)
|
||||
p_bf16 = p.to(tl.bfloat16)
|
||||
|
||||
offs_v = tl.arange(0, BLOCK_D)
|
||||
v_ptrs = kv_base + offs_n[:, None] * stride_kv_k + offs_v[None, :] * stride_kv_d
|
||||
v = tl.load(v_ptrs, mask=valid[:, None], other=0.0).to(tl.bfloat16)
|
||||
acc_0 = acc_0 * alpha[:, None] + tl.dot(p_bf16, v)
|
||||
|
||||
offs_v = BLOCK_D + tl.arange(0, BLOCK_D)
|
||||
v_ptrs = kv_base + offs_n[:, None] * stride_kv_k + offs_v[None, :] * stride_kv_d
|
||||
v = tl.load(
|
||||
v_ptrs, mask=valid[:, None] & (offs_v[None, :] < d_v), other=0.0
|
||||
).to(tl.bfloat16)
|
||||
acc_1 = acc_1 * alpha[:, None] + tl.dot(p_bf16, v)
|
||||
|
||||
offs_v = 2 * BLOCK_D + tl.arange(0, BLOCK_D)
|
||||
v_ptrs = kv_base + offs_n[:, None] * stride_kv_k + offs_v[None, :] * stride_kv_d
|
||||
v = tl.load(
|
||||
v_ptrs, mask=valid[:, None] & (offs_v[None, :] < d_v), other=0.0
|
||||
).to(tl.bfloat16)
|
||||
acc_2 = acc_2 * alpha[:, None] + tl.dot(p_bf16, v)
|
||||
|
||||
offs_v = 3 * BLOCK_D + tl.arange(0, BLOCK_D)
|
||||
v_ptrs = kv_base + offs_n[:, None] * stride_kv_k + offs_v[None, :] * stride_kv_d
|
||||
v = tl.load(
|
||||
v_ptrs, mask=valid[:, None] & (offs_v[None, :] < d_v), other=0.0
|
||||
).to(tl.bfloat16)
|
||||
acc_3 = acc_3 * alpha[:, None] + tl.dot(p_bf16, v)
|
||||
|
||||
m_i = m_new
|
||||
l_i = l_new
|
||||
|
||||
# Store partial results
|
||||
stride_po_s_64 = tl.cast(stride_po_s, tl.int64)
|
||||
stride_po_t_64 = tl.cast(stride_po_t, tl.int64)
|
||||
po_base = PartialOutput + pid_k * stride_po_s_64 + pid_t_64 * stride_po_t_64
|
||||
|
||||
offs_h_2d = offs_h[:, None]
|
||||
mask_h_2d = mask_h[:, None]
|
||||
offs_v_0 = tl.arange(0, BLOCK_D)
|
||||
offs_v_1 = BLOCK_D + tl.arange(0, BLOCK_D)
|
||||
offs_v_2 = 2 * BLOCK_D + tl.arange(0, BLOCK_D)
|
||||
offs_v_3 = 3 * BLOCK_D + tl.arange(0, BLOCK_D)
|
||||
|
||||
tl.store(
|
||||
po_base + offs_h_2d * stride_po_h + offs_v_0[None, :] * stride_po_d,
|
||||
acc_0,
|
||||
mask=mask_h_2d,
|
||||
)
|
||||
tl.store(
|
||||
po_base + offs_h_2d * stride_po_h + offs_v_1[None, :] * stride_po_d,
|
||||
acc_1,
|
||||
mask=mask_h_2d & (offs_v_1[None, :] < d_v),
|
||||
)
|
||||
tl.store(
|
||||
po_base + offs_h_2d * stride_po_h + offs_v_2[None, :] * stride_po_d,
|
||||
acc_2,
|
||||
mask=mask_h_2d & (offs_v_2[None, :] < d_v),
|
||||
)
|
||||
tl.store(
|
||||
po_base + offs_h_2d * stride_po_h + offs_v_3[None, :] * stride_po_d,
|
||||
acc_3,
|
||||
mask=mask_h_2d & (offs_v_3[None, :] < d_v),
|
||||
)
|
||||
|
||||
stride_plse_s_64 = tl.cast(stride_plse_s, tl.int64)
|
||||
stride_plse_t_64 = tl.cast(stride_plse_t, tl.int64)
|
||||
plse_ptrs = (
|
||||
PartialLSE
|
||||
+ pid_k * stride_plse_s_64
|
||||
+ pid_t_64 * stride_plse_t_64
|
||||
+ offs_h * stride_plse_h
|
||||
)
|
||||
tl.store(plse_ptrs, l_i, mask=mask_h)
|
||||
|
||||
stride_pm_s_64 = tl.cast(stride_pm_s, tl.int64)
|
||||
stride_pm_t_64 = tl.cast(stride_pm_t, tl.int64)
|
||||
pm_ptrs = (
|
||||
PartialM
|
||||
+ pid_k * stride_pm_s_64
|
||||
+ pid_t_64 * stride_pm_t_64
|
||||
+ offs_h * stride_pm_h
|
||||
)
|
||||
tl.store(pm_ptrs, m_i, mask=mask_h)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Combine Kernel for Split-K
|
||||
# ============================================================================
|
||||
@triton.autotune(
|
||||
configs=[
|
||||
# Simple reduce kernel merging split-K results.
|
||||
# - BLOCK_D=128: 4 iterations to cover d_v=512.
|
||||
# - num_warps=4: sufficient for this simple reduce operation.
|
||||
# - BLOCK_H varies for different batch sizes:
|
||||
triton.Config({"BLOCK_H": 16, "BLOCK_D": 128}, num_warps=4, num_stages=1),
|
||||
triton.Config({"BLOCK_H": 32, "BLOCK_D": 128}, num_warps=4, num_stages=1),
|
||||
triton.Config({"BLOCK_H": 64, "BLOCK_D": 128}, num_warps=4, num_stages=1),
|
||||
],
|
||||
key=["total_tokens_bucket", "h_q", "split_k"],
|
||||
)
|
||||
@triton.jit
|
||||
def _combine_splitk_attention_kernel(
|
||||
PartialOutput,
|
||||
PartialLSE,
|
||||
PartialM,
|
||||
AttnSink,
|
||||
Output,
|
||||
LSE,
|
||||
total_tokens,
|
||||
total_tokens_bucket,
|
||||
h_q,
|
||||
d_v,
|
||||
split_k,
|
||||
stride_po_s,
|
||||
stride_po_t,
|
||||
stride_po_h,
|
||||
stride_po_d,
|
||||
stride_plse_s,
|
||||
stride_plse_t,
|
||||
stride_plse_h,
|
||||
stride_pm_s,
|
||||
stride_pm_t,
|
||||
stride_pm_h,
|
||||
stride_o_t,
|
||||
stride_o_h,
|
||||
stride_o_d,
|
||||
stride_lse_t,
|
||||
stride_lse_h,
|
||||
HAS_ATTN_SINK: tl.constexpr,
|
||||
BLOCK_H: tl.constexpr,
|
||||
BLOCK_D: tl.constexpr,
|
||||
):
|
||||
"""Combine partial results from split-K attention kernel."""
|
||||
LOG2E: tl.constexpr = 1.4426950408889634
|
||||
NEG_INF = float("-inf")
|
||||
POS_INF = float("+inf")
|
||||
|
||||
pid_t = tl.program_id(0)
|
||||
pid_h = tl.program_id(1)
|
||||
pid_t_64 = pid_t.to(tl.int64)
|
||||
|
||||
offs_h = pid_h * BLOCK_H + tl.arange(0, BLOCK_H)
|
||||
mask_h = offs_h < h_q
|
||||
|
||||
m_acc = tl.full([BLOCK_H], NEG_INF, dtype=tl.float32)
|
||||
l_acc = tl.zeros([BLOCK_H], dtype=tl.float32)
|
||||
|
||||
acc_0 = tl.zeros([BLOCK_H, BLOCK_D], dtype=tl.float32)
|
||||
acc_1 = tl.zeros([BLOCK_H, BLOCK_D], dtype=tl.float32)
|
||||
acc_2 = tl.zeros([BLOCK_H, BLOCK_D], dtype=tl.float32)
|
||||
acc_3 = tl.zeros([BLOCK_H, BLOCK_D], dtype=tl.float32)
|
||||
|
||||
stride_po_s_64 = tl.cast(stride_po_s, tl.int64)
|
||||
stride_po_t_64 = tl.cast(stride_po_t, tl.int64)
|
||||
stride_plse_s_64 = tl.cast(stride_plse_s, tl.int64)
|
||||
stride_plse_t_64 = tl.cast(stride_plse_t, tl.int64)
|
||||
stride_pm_s_64 = tl.cast(stride_pm_s, tl.int64)
|
||||
stride_pm_t_64 = tl.cast(stride_pm_t, tl.int64)
|
||||
|
||||
offs_h_2d = offs_h[:, None]
|
||||
mask_h_2d = mask_h[:, None]
|
||||
offs_v_0 = tl.arange(0, BLOCK_D)
|
||||
offs_v_1 = BLOCK_D + tl.arange(0, BLOCK_D)
|
||||
offs_v_2 = 2 * BLOCK_D + tl.arange(0, BLOCK_D)
|
||||
offs_v_3 = 3 * BLOCK_D + tl.arange(0, BLOCK_D)
|
||||
|
||||
for k in range(split_k):
|
||||
k_64 = tl.cast(k, tl.int64)
|
||||
po_base = PartialOutput + k_64 * stride_po_s_64 + pid_t_64 * stride_po_t_64
|
||||
|
||||
p_acc_0 = tl.load(
|
||||
po_base + offs_h_2d * stride_po_h + offs_v_0[None, :] * stride_po_d,
|
||||
mask=mask_h_2d,
|
||||
other=0.0,
|
||||
)
|
||||
p_acc_1 = tl.load(
|
||||
po_base + offs_h_2d * stride_po_h + offs_v_1[None, :] * stride_po_d,
|
||||
mask=mask_h_2d & (offs_v_1[None, :] < d_v),
|
||||
other=0.0,
|
||||
)
|
||||
p_acc_2 = tl.load(
|
||||
po_base + offs_h_2d * stride_po_h + offs_v_2[None, :] * stride_po_d,
|
||||
mask=mask_h_2d & (offs_v_2[None, :] < d_v),
|
||||
other=0.0,
|
||||
)
|
||||
p_acc_3 = tl.load(
|
||||
po_base + offs_h_2d * stride_po_h + offs_v_3[None, :] * stride_po_d,
|
||||
mask=mask_h_2d & (offs_v_3[None, :] < d_v),
|
||||
other=0.0,
|
||||
)
|
||||
|
||||
plse_ptrs = (
|
||||
PartialLSE
|
||||
+ k_64 * stride_plse_s_64
|
||||
+ pid_t_64 * stride_plse_t_64
|
||||
+ offs_h * stride_plse_h
|
||||
)
|
||||
p_l = tl.load(plse_ptrs, mask=mask_h, other=0.0)
|
||||
|
||||
pm_ptrs = (
|
||||
PartialM
|
||||
+ k_64 * stride_pm_s_64
|
||||
+ pid_t_64 * stride_pm_t_64
|
||||
+ offs_h * stride_pm_h
|
||||
)
|
||||
p_m = tl.load(pm_ptrs, mask=mask_h, other=NEG_INF)
|
||||
|
||||
m_new = tl.maximum(m_acc, p_m)
|
||||
alpha_acc = tl.where(
|
||||
m_acc == NEG_INF, 0.0, tl.math.exp2((m_acc - m_new) * LOG2E)
|
||||
)
|
||||
alpha_p = tl.where(p_m == NEG_INF, 0.0, tl.math.exp2((p_m - m_new) * LOG2E))
|
||||
l_new = alpha_acc * l_acc + alpha_p * p_l
|
||||
|
||||
acc_0 = acc_0 * alpha_acc[:, None] + p_acc_0 * alpha_p[:, None]
|
||||
acc_1 = acc_1 * alpha_acc[:, None] + p_acc_1 * alpha_p[:, None]
|
||||
acc_2 = acc_2 * alpha_acc[:, None] + p_acc_2 * alpha_p[:, None]
|
||||
acc_3 = acc_3 * alpha_acc[:, None] + p_acc_3 * alpha_p[:, None]
|
||||
|
||||
m_acc = m_new
|
||||
l_acc = l_new
|
||||
|
||||
lse = m_acc + tl.math.log2(tl.where(l_acc == 0.0, 1.0, l_acc)) / LOG2E
|
||||
is_lonely_q = l_acc == 0.0
|
||||
|
||||
if HAS_ATTN_SINK:
|
||||
attn_sink_vals = tl.load(AttnSink + offs_h, mask=mask_h, other=0.0)
|
||||
exp_attn_sink_minus_m = tl.math.exp2((attn_sink_vals - m_acc) * LOG2E)
|
||||
denominator = l_acc + exp_attn_sink_minus_m
|
||||
denominator = tl.where(denominator == 0.0, 1.0, denominator)
|
||||
output_scale = 1.0 / denominator
|
||||
else:
|
||||
output_scale = tl.where(l_acc == 0.0, 0.0, 1.0 / l_acc)
|
||||
|
||||
is_lonely_q_2d = is_lonely_q[:, None]
|
||||
output_scale_2d = output_scale[:, None]
|
||||
acc_0 = tl.where(is_lonely_q_2d, 0.0, acc_0 * output_scale_2d)
|
||||
acc_1 = tl.where(is_lonely_q_2d, 0.0, acc_1 * output_scale_2d)
|
||||
acc_2 = tl.where(is_lonely_q_2d, 0.0, acc_2 * output_scale_2d)
|
||||
acc_3 = tl.where(is_lonely_q_2d, 0.0, acc_3 * output_scale_2d)
|
||||
lse = tl.where(is_lonely_q, POS_INF, lse)
|
||||
|
||||
stride_o_t_64 = tl.cast(stride_o_t, tl.int64)
|
||||
o_base = Output + pid_t_64 * stride_o_t_64
|
||||
|
||||
tl.store(
|
||||
o_base + offs_h_2d * stride_o_h + offs_v_0[None, :] * stride_o_d,
|
||||
acc_0.to(tl.bfloat16),
|
||||
mask=mask_h_2d,
|
||||
)
|
||||
tl.store(
|
||||
o_base + offs_h_2d * stride_o_h + offs_v_1[None, :] * stride_o_d,
|
||||
acc_1.to(tl.bfloat16),
|
||||
mask=mask_h_2d & (offs_v_1[None, :] < d_v),
|
||||
)
|
||||
tl.store(
|
||||
o_base + offs_h_2d * stride_o_h + offs_v_2[None, :] * stride_o_d,
|
||||
acc_2.to(tl.bfloat16),
|
||||
mask=mask_h_2d & (offs_v_2[None, :] < d_v),
|
||||
)
|
||||
tl.store(
|
||||
o_base + offs_h_2d * stride_o_h + offs_v_3[None, :] * stride_o_d,
|
||||
acc_3.to(tl.bfloat16),
|
||||
mask=mask_h_2d & (offs_v_3[None, :] < d_v),
|
||||
)
|
||||
|
||||
stride_lse_t_64 = tl.cast(stride_lse_t, tl.int64)
|
||||
tl.store(LSE + pid_t_64 * stride_lse_t_64 + offs_h * stride_lse_h, lse, mask=mask_h)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Runner Function
|
||||
# ============================================================================
|
||||
def run_splitk_attention(
|
||||
q_reshaped: torch.Tensor,
|
||||
gathered_kv: torch.Tensor,
|
||||
invalid_mask: torch.Tensor,
|
||||
d_v: int,
|
||||
sm_scale: float,
|
||||
total_tokens: int,
|
||||
h_q: int,
|
||||
total_topk: int,
|
||||
d_qk: int,
|
||||
attn_sink: Optional[torch.Tensor] = None,
|
||||
split_k: int = 4,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Run split-K attention kernel."""
|
||||
device = q_reshaped.device
|
||||
|
||||
topk_per_split = (total_topk + split_k - 1) // split_k
|
||||
|
||||
partial_output = torch.empty(
|
||||
split_k, total_tokens, h_q, d_v, dtype=torch.float32, device=device
|
||||
)
|
||||
partial_lse = torch.empty(
|
||||
split_k, total_tokens, h_q, dtype=torch.float32, device=device
|
||||
)
|
||||
partial_m = torch.empty(
|
||||
split_k, total_tokens, h_q, dtype=torch.float32, device=device
|
||||
)
|
||||
|
||||
output = torch.empty(total_tokens, h_q, d_v, dtype=torch.bfloat16, device=device)
|
||||
lse = torch.empty(total_tokens, h_q, dtype=torch.float32, device=device)
|
||||
|
||||
grid_splitk = lambda meta: (
|
||||
total_tokens,
|
||||
triton.cdiv(h_q, meta["BLOCK_H"]),
|
||||
split_k,
|
||||
)
|
||||
_splitk_attention_kernel[grid_splitk](
|
||||
q_reshaped,
|
||||
gathered_kv,
|
||||
invalid_mask,
|
||||
partial_output,
|
||||
partial_lse,
|
||||
partial_m,
|
||||
sm_scale,
|
||||
total_tokens,
|
||||
_bucket_total_tokens(total_tokens),
|
||||
h_q,
|
||||
total_topk,
|
||||
d_qk,
|
||||
d_v,
|
||||
topk_per_split,
|
||||
q_reshaped.stride(0),
|
||||
q_reshaped.stride(1),
|
||||
q_reshaped.stride(2),
|
||||
gathered_kv.stride(0),
|
||||
gathered_kv.stride(1),
|
||||
gathered_kv.stride(2),
|
||||
invalid_mask.stride(0),
|
||||
invalid_mask.stride(1),
|
||||
partial_output.stride(0),
|
||||
partial_output.stride(1),
|
||||
partial_output.stride(2),
|
||||
partial_output.stride(3),
|
||||
partial_lse.stride(0),
|
||||
partial_lse.stride(1),
|
||||
partial_lse.stride(2),
|
||||
partial_m.stride(0),
|
||||
partial_m.stride(1),
|
||||
partial_m.stride(2),
|
||||
)
|
||||
|
||||
HAS_ATTN_SINK = attn_sink is not None
|
||||
attn_sink_tensor = attn_sink if HAS_ATTN_SINK else lse[:1]
|
||||
|
||||
grid_combine = lambda meta: (total_tokens, triton.cdiv(h_q, meta["BLOCK_H"]))
|
||||
_combine_splitk_attention_kernel[grid_combine](
|
||||
partial_output,
|
||||
partial_lse,
|
||||
partial_m,
|
||||
attn_sink_tensor,
|
||||
output,
|
||||
lse,
|
||||
total_tokens,
|
||||
_bucket_total_tokens(total_tokens),
|
||||
h_q,
|
||||
d_v,
|
||||
split_k,
|
||||
partial_output.stride(0),
|
||||
partial_output.stride(1),
|
||||
partial_output.stride(2),
|
||||
partial_output.stride(3),
|
||||
partial_lse.stride(0),
|
||||
partial_lse.stride(1),
|
||||
partial_lse.stride(2),
|
||||
partial_m.stride(0),
|
||||
partial_m.stride(1),
|
||||
partial_m.stride(2),
|
||||
output.stride(0),
|
||||
output.stride(1),
|
||||
output.stride(2),
|
||||
lse.stride(0),
|
||||
lse.stride(1),
|
||||
HAS_ATTN_SINK=HAS_ATTN_SINK,
|
||||
)
|
||||
|
||||
return output, lse
|
||||
@@ -288,6 +288,92 @@ def _fused_norm_rope_kernel(
|
||||
)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _fused_softmax_pool_kernel(
|
||||
kv_score_ptr,
|
||||
out_ptr,
|
||||
stride_bs: tl.constexpr,
|
||||
stride_k: tl.constexpr,
|
||||
K: tl.constexpr,
|
||||
HEAD_DIM: tl.constexpr,
|
||||
HEAD_BLOCK: tl.constexpr,
|
||||
):
|
||||
pid = tl.program_id(0)
|
||||
base = pid * stride_bs
|
||||
|
||||
offs = tl.arange(0, HEAD_BLOCK)
|
||||
mask = offs < HEAD_DIM
|
||||
|
||||
max_val = tl.full([HEAD_BLOCK], float("-inf"), dtype=tl.float32)
|
||||
for k in range(K):
|
||||
s = tl.load(
|
||||
kv_score_ptr + base + k * stride_k + HEAD_DIM + offs,
|
||||
mask=mask,
|
||||
other=float("-inf"),
|
||||
).to(tl.float32)
|
||||
max_val = tl.maximum(max_val, s)
|
||||
|
||||
sum_exp = tl.zeros([HEAD_BLOCK], dtype=tl.float32)
|
||||
weighted = tl.zeros([HEAD_BLOCK], dtype=tl.float32)
|
||||
for k in range(K):
|
||||
s = tl.load(
|
||||
kv_score_ptr + base + k * stride_k + HEAD_DIM + offs,
|
||||
mask=mask,
|
||||
other=float("-inf"),
|
||||
).to(tl.float32)
|
||||
v = tl.load(
|
||||
kv_score_ptr + base + k * stride_k + offs,
|
||||
mask=mask,
|
||||
other=0.0,
|
||||
).to(tl.float32)
|
||||
w = tl.exp(s - max_val)
|
||||
sum_exp += w
|
||||
weighted += v * w
|
||||
|
||||
result = weighted / sum_exp
|
||||
tl.store(
|
||||
out_ptr + pid * HEAD_DIM + offs, result.to(out_ptr.dtype.element_ty), mask=mask
|
||||
)
|
||||
|
||||
|
||||
def fused_softmax_pool_triton(
|
||||
kv_score: torch.Tensor,
|
||||
head_dim: int,
|
||||
) -> torch.Tensor:
|
||||
"""Fused softmax-weighted-sum: out = (kv * softmax(score, dim=1)).sum(dim=1).
|
||||
|
||||
Replaces the generic cunn_SpatialSoftMaxForward + elementwise multiply + sum
|
||||
with a single Triton kernel.
|
||||
|
||||
Args:
|
||||
kv_score: [bs, K, 2 * head_dim] where first head_dim is kv, second is score.
|
||||
head_dim: dimension of each of kv and score.
|
||||
Returns:
|
||||
output: [bs, head_dim]
|
||||
"""
|
||||
assert kv_score.dim() == 3
|
||||
bs, K, last = kv_score.shape
|
||||
assert last == 2 * head_dim
|
||||
assert kv_score.is_contiguous()
|
||||
|
||||
out = torch.empty(bs, head_dim, dtype=kv_score.dtype, device=kv_score.device)
|
||||
if bs == 0:
|
||||
return out
|
||||
|
||||
HEAD_BLOCK = triton.next_power_of_2(head_dim)
|
||||
grid = (bs,)
|
||||
_fused_softmax_pool_kernel[grid](
|
||||
kv_score,
|
||||
out,
|
||||
stride_bs=kv_score.stride(0),
|
||||
stride_k=kv_score.stride(1),
|
||||
K=K,
|
||||
HEAD_DIM=head_dim,
|
||||
HEAD_BLOCK=HEAD_BLOCK,
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
def fused_norm_rope_inplace_triton(
|
||||
kv: torch.Tensor,
|
||||
weight: Optional[torch.Tensor],
|
||||
|
||||
@@ -0,0 +1,157 @@
|
||||
"""Fused Q/K RMSNorm in a single Triton kernel launch.
|
||||
|
||||
Ported from ATOM (atom/model_ops/layernorm.py). Fuses per-head Q RMSNorm
|
||||
(optionally weightless) and KV RMSNorm into one kernel, halving the number
|
||||
of norm kernel launches per attention layer.
|
||||
"""
|
||||
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _fused_qk_norm_kernel(
|
||||
q_ptr,
|
||||
k_ptr,
|
||||
q_out_ptr,
|
||||
k_out_ptr,
|
||||
q_weight_ptr,
|
||||
k_weight_ptr,
|
||||
eps,
|
||||
num_tokens,
|
||||
head_dim,
|
||||
q_in_stride0,
|
||||
k_in_stride0,
|
||||
q_out_stride0,
|
||||
k_out_stride0,
|
||||
num_q_heads,
|
||||
num_k_heads,
|
||||
Q_HAS_WEIGHT: tl.constexpr,
|
||||
RBLOCK: tl.constexpr,
|
||||
XBLOCK: tl.constexpr,
|
||||
):
|
||||
num_q_rows = num_tokens * num_q_heads
|
||||
total_rows = num_tokens * (num_q_heads + num_k_heads)
|
||||
|
||||
xoffset = tl.program_id(0) * XBLOCK
|
||||
xindex = xoffset + tl.arange(0, XBLOCK)[:, None]
|
||||
xmask = xindex < total_rows
|
||||
cols = tl.arange(0, RBLOCK)[None, :]
|
||||
col_mask = cols < head_dim
|
||||
|
||||
is_q = xindex < num_q_rows
|
||||
row_in_section = tl.where(is_q, xindex, xindex - num_q_rows)
|
||||
cur_num_heads = tl.where(is_q, num_q_heads, num_k_heads)
|
||||
|
||||
tokens = row_in_section // cur_num_heads
|
||||
heads = row_in_section % cur_num_heads
|
||||
|
||||
in_stride = tl.where(is_q, q_in_stride0, k_in_stride0)
|
||||
in_bases = tokens * in_stride + heads * head_dim
|
||||
|
||||
out_stride0 = tl.where(is_q, q_out_stride0, k_out_stride0)
|
||||
out_bases = tokens * out_stride0 + heads * head_dim
|
||||
|
||||
mask = xmask & col_mask
|
||||
|
||||
if Q_HAS_WEIGHT:
|
||||
qw = tl.load(
|
||||
q_weight_ptr + cols, mask=col_mask, other=0.0, eviction_policy="evict_last"
|
||||
).to(tl.float32)
|
||||
else:
|
||||
qw = tl.full((RBLOCK,), 1.0, tl.float32)
|
||||
kw = tl.load(
|
||||
k_weight_ptr + cols, mask=col_mask, other=0.0, eviction_policy="evict_last"
|
||||
).to(tl.float32)
|
||||
w = tl.where(is_q, qw, kw)
|
||||
|
||||
x = tl.load(
|
||||
q_ptr + in_bases + cols,
|
||||
mask=mask & is_q,
|
||||
other=0.0,
|
||||
eviction_policy="evict_first",
|
||||
).to(tl.float32)
|
||||
x = x + tl.load(
|
||||
k_ptr + in_bases + cols,
|
||||
mask=mask & ~is_q,
|
||||
other=0.0,
|
||||
eviction_policy="evict_first",
|
||||
).to(tl.float32)
|
||||
|
||||
var = tl.sum(x * x, 1)[:, None]
|
||||
rstd = tl.rsqrt(var / head_dim + eps)
|
||||
|
||||
out = (x * rstd * w).to(q_out_ptr.dtype.element_ty)
|
||||
tl.store(
|
||||
q_out_ptr + out_bases + cols,
|
||||
out,
|
||||
mask=mask & is_q,
|
||||
eviction_policy="evict_first",
|
||||
)
|
||||
tl.store(
|
||||
k_out_ptr + out_bases + cols,
|
||||
out,
|
||||
mask=mask & ~is_q,
|
||||
eviction_policy="evict_first",
|
||||
)
|
||||
|
||||
|
||||
def fused_qk_norm(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
q_weight: Optional[torch.Tensor],
|
||||
k_weight: torch.Tensor,
|
||||
eps: float,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Fused Q/K RMSNorm in a single Triton kernel launch.
|
||||
|
||||
Args:
|
||||
q: [num_tokens, num_heads, head_dim]
|
||||
k: [num_tokens, num_kv_heads, head_dim]
|
||||
q_weight: [head_dim] norm weight, or None for weightless Q norm
|
||||
k_weight: [head_dim] norm weight (always required)
|
||||
eps: epsilon for numerical stability
|
||||
|
||||
Returns:
|
||||
(q_normed, k_normed) same shapes as inputs
|
||||
"""
|
||||
head_dim = k_weight.shape[0]
|
||||
if q_weight is not None:
|
||||
assert q_weight.shape[0] == head_dim
|
||||
num_tokens = q.shape[0]
|
||||
num_q_heads = q.shape[1]
|
||||
num_k_heads = k.shape[1]
|
||||
total_rows = num_tokens * (num_q_heads + num_k_heads)
|
||||
RBLOCK = triton.next_power_of_2(head_dim)
|
||||
|
||||
q_out = torch.empty_like(q)
|
||||
k_out = torch.empty_like(k)
|
||||
|
||||
XBLOCK = 2 if total_rows > 8192 else 1
|
||||
NUM_WARPS = 1
|
||||
q_weight_arg = q_weight if q_weight is not None else k_weight
|
||||
_fused_qk_norm_kernel[((total_rows + XBLOCK - 1) // XBLOCK,)](
|
||||
q,
|
||||
k,
|
||||
q_out,
|
||||
k_out,
|
||||
q_weight_arg,
|
||||
k_weight,
|
||||
eps,
|
||||
num_tokens,
|
||||
head_dim,
|
||||
q.stride(0),
|
||||
k.stride(0),
|
||||
q_out.stride(0),
|
||||
k_out.stride(0),
|
||||
num_q_heads,
|
||||
num_k_heads,
|
||||
Q_HAS_WEIGHT=q_weight is not None,
|
||||
RBLOCK=RBLOCK,
|
||||
XBLOCK=XBLOCK,
|
||||
num_warps=NUM_WARPS,
|
||||
)
|
||||
return q_out, k_out
|
||||
@@ -56,6 +56,7 @@ class AiterMoeQuantInfo(MoeQuantInfo):
|
||||
doweight_stage1: bool = False
|
||||
hidden_pad: int = 0
|
||||
intermediate_pad: int = 0
|
||||
swiglu_limit: float = 0.0
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -116,6 +117,7 @@ class AiterRunnerCore(MoeRunnerCore):
|
||||
return AiterRunnerOutput(hidden_states=runner_input.hidden_states)
|
||||
|
||||
from aiter.fused_moe import fused_moe
|
||||
from aiter.ops.flydsl.moe_common import GateMode
|
||||
|
||||
a1_scale = (
|
||||
runner_input.a1_scale
|
||||
@@ -128,6 +130,9 @@ class AiterRunnerCore(MoeRunnerCore):
|
||||
extra["num_local_tokens"] = runner_input.num_local_tokens
|
||||
if runner_input.output_dtype is not None:
|
||||
extra["dtype"] = runner_input.output_dtype
|
||||
if quant_info.swiglu_limit > 0:
|
||||
extra["gate_mode"] = GateMode.INTERLEAVE.value
|
||||
extra["swiglu_limit"] = quant_info.swiglu_limit
|
||||
|
||||
output = fused_moe(
|
||||
hidden_states=runner_input.hidden_states,
|
||||
|
||||
@@ -898,20 +898,46 @@ def biased_topk_jit_kernel_impl(
|
||||
):
|
||||
assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch"
|
||||
|
||||
from sglang.jit_kernel.moe_fused_gate import moe_fused_gate
|
||||
if _use_aiter and scoring_func == "sqrtsoftplus" and num_fused_shared_experts == 0:
|
||||
from aiter import topk_gating
|
||||
|
||||
topk_weights, topk_ids = moe_fused_gate(
|
||||
gating_output,
|
||||
correction_bias,
|
||||
topk=topk,
|
||||
scoring_func=scoring_func,
|
||||
num_fused_shared_experts=num_fused_shared_experts,
|
||||
renormalize=renormalize,
|
||||
routed_scaling_factor=routed_scaling_factor,
|
||||
apply_routed_scaling_factor_on_output=apply_routed_scaling_factor_on_output,
|
||||
)
|
||||
topk_weights, topk_ids = topk_weights.to(torch.float32), topk_ids.to(torch.int32)
|
||||
return topk_weights, topk_ids
|
||||
num_tokens = gating_output.shape[0]
|
||||
topk_weights = torch.empty(
|
||||
(num_tokens, topk), dtype=torch.float32, device=gating_output.device
|
||||
)
|
||||
topk_ids = torch.empty(
|
||||
(num_tokens, topk), dtype=torch.int32, device=gating_output.device
|
||||
)
|
||||
|
||||
topk_gating(
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
gating_output,
|
||||
correction_bias,
|
||||
renormalize,
|
||||
routed_scaling_factor,
|
||||
score_func="sqrtsoftplus",
|
||||
)
|
||||
|
||||
return topk_weights, topk_ids
|
||||
|
||||
else:
|
||||
from sglang.jit_kernel.moe_fused_gate import moe_fused_gate
|
||||
|
||||
topk_weights, topk_ids = moe_fused_gate(
|
||||
gating_output,
|
||||
correction_bias,
|
||||
topk=topk,
|
||||
scoring_func=scoring_func,
|
||||
num_fused_shared_experts=num_fused_shared_experts,
|
||||
renormalize=renormalize,
|
||||
routed_scaling_factor=routed_scaling_factor,
|
||||
apply_routed_scaling_factor_on_output=apply_routed_scaling_factor_on_output,
|
||||
)
|
||||
topk_weights, topk_ids = topk_weights.to(torch.float32), topk_ids.to(
|
||||
torch.int32
|
||||
)
|
||||
return topk_weights, topk_ids
|
||||
|
||||
|
||||
@torch.compile(dynamic=True, backend=get_compiler_backend(), disable=_is_npu)
|
||||
|
||||
@@ -125,8 +125,11 @@ def _require_fp4_dtype():
|
||||
|
||||
|
||||
if _use_aiter or _use_hip_int4:
|
||||
from aiter.ops.shuffle import shuffle_weight
|
||||
from aiter.utility.fp4_utils import e8m0_shuffle
|
||||
from aiter.ops.shuffle import (
|
||||
shuffle_scale_a16w4,
|
||||
shuffle_weight,
|
||||
shuffle_weight_a16w4,
|
||||
)
|
||||
|
||||
if _use_aiter:
|
||||
from sglang.srt.layers.quantization.fp8_utils import (
|
||||
@@ -1217,8 +1220,10 @@ class Fp8MoEMethod(FusedMoEMethodBase):
|
||||
for scale_name in ("w13_weight_scale_inv", "w2_weight_scale_inv"):
|
||||
scale = getattr(layer, scale_name)
|
||||
num_experts, num_rows, _ = scale.shape
|
||||
scale.data = e8m0_shuffle(scale.view(num_experts * num_rows, -1)).view(
|
||||
num_experts, num_rows, -1
|
||||
# a8w4: aiter flydsl scale layout
|
||||
is_w13_scale = scale_name == "w13_weight_scale_inv"
|
||||
scale.data = shuffle_scale_a16w4(
|
||||
scale.view(num_experts * num_rows, -1), num_experts, is_w13_scale
|
||||
)
|
||||
|
||||
layer.w13_weight.data = layer.w13_weight.data.view(fp4_weight_dtype)
|
||||
@@ -1226,11 +1231,12 @@ class Fp8MoEMethod(FusedMoEMethodBase):
|
||||
|
||||
is_shuffled = _is_shuffle_moe_mxfp4
|
||||
if is_shuffled:
|
||||
layer.w13_weight.data = shuffle_weight(
|
||||
layer.w13_weight.contiguous(), (16, 16)
|
||||
# a8w4: aiter flydsl weight layout
|
||||
layer.w13_weight.data = shuffle_weight_a16w4(
|
||||
layer.w13_weight.contiguous(), 16, True
|
||||
)
|
||||
layer.w2_weight.data = shuffle_weight(
|
||||
layer.w2_weight.contiguous(), (16, 16)
|
||||
layer.w2_weight.data = shuffle_weight_a16w4(
|
||||
layer.w2_weight.contiguous(), 16, False
|
||||
)
|
||||
layer.w13_weight.is_shuffled = is_shuffled
|
||||
layer.w2_weight.is_shuffled = is_shuffled
|
||||
@@ -2075,6 +2081,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
|
||||
w13_scale=w13_scale,
|
||||
w2_scale=w2_scale,
|
||||
expert_mask=layer.dispatcher.expert_mask_gpu if _use_aiter else None,
|
||||
swiglu_limit=self.moe_runner_config.swiglu_limit or 0.0,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -470,8 +470,13 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
|
||||
enable_memory_saver,
|
||||
)
|
||||
|
||||
indexer_size = (
|
||||
self.c4_logical_size
|
||||
if (not _is_hip or envs.SGLANG_OPT_USE_COMPRESSOR_V2.get())
|
||||
else c4_size
|
||||
)
|
||||
self.c4_indexer_kv_pool = DeepSeekV4IndexerPool(
|
||||
self.c4_logical_size if not _is_hip else c4_size,
|
||||
indexer_size,
|
||||
c4_page_size,
|
||||
dtype,
|
||||
indexer_head_dim,
|
||||
|
||||
@@ -840,8 +840,11 @@ class DeepseekV2MoE(nn.Module):
|
||||
**topk_kwargs,
|
||||
)
|
||||
final_hidden_states = self.experts(hidden_states, topk_output)
|
||||
if not (_is_cuda or _is_musa) or isinstance(
|
||||
self.experts.quant_method, KTEPWrapperMethod
|
||||
if (
|
||||
not _is_cuda
|
||||
and not _is_musa
|
||||
and not _use_aiter
|
||||
or isinstance(self.experts.quant_method, KTEPWrapperMethod)
|
||||
):
|
||||
final_hidden_states *= self.routed_scaling_factor
|
||||
|
||||
|
||||
@@ -536,6 +536,118 @@ class MQALayer(nn.Module):
|
||||
|
||||
return q
|
||||
|
||||
def _forward_prepare_multi_stream_hip(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
forward_batch: ForwardBatch,
|
||||
attn_backend,
|
||||
q_out: Optional[torch.Tensor] = None,
|
||||
x_quant=None,
|
||||
) -> torch.Tensor:
|
||||
"""ATOM-style ROCm path: overlap compressors, keep Q/KV on main stream."""
|
||||
assert self.alt_streams is not None
|
||||
assert len(self.alt_streams) >= 1
|
||||
|
||||
current_stream = torch.cuda.current_stream()
|
||||
stream_compressor = self.alt_streams[0]
|
||||
stream_indexer_compressor = (
|
||||
self.alt_streams[1] if len(self.alt_streams) > 1 else None
|
||||
)
|
||||
|
||||
if self.compressor is not None:
|
||||
stream_compressor.wait_stream(current_stream)
|
||||
with torch.cuda.stream(stream_compressor):
|
||||
attn_backend.forward_core_compressor(
|
||||
x, forward_batch, self.layer_id, self.compressor
|
||||
)
|
||||
|
||||
if self.indexer is not None and stream_indexer_compressor is not None:
|
||||
stream_indexer_compressor.wait_stream(current_stream)
|
||||
with torch.cuda.stream(stream_indexer_compressor):
|
||||
attn_backend.forward_indexer_compressor(
|
||||
x=x,
|
||||
forward_batch=forward_batch,
|
||||
layer_id=self.indexer.layer_id,
|
||||
compressor=self.indexer.compressor,
|
||||
)
|
||||
|
||||
x_linear = x_quant if x_quant is not None else x
|
||||
if self.fuse_wqa_wkv:
|
||||
qkv_a, _ = self.wqkv_a(x_linear)
|
||||
q_lora = qkv_a[..., : self.q_lora_rank]
|
||||
else:
|
||||
q_lora, _ = self.wq_a(x_linear)
|
||||
qkv_a = None
|
||||
|
||||
if self.use_fused_qk_norm_rope:
|
||||
if _is_gfx95_supported:
|
||||
q_for_wqb, q_lora = _fused_rmsnorm_fp8_quant(
|
||||
q_lora,
|
||||
self.q_norm.weight,
|
||||
self.q_norm.variance_epsilon,
|
||||
)
|
||||
q, _ = self.wq_b(q_for_wqb)
|
||||
else:
|
||||
q_lora = self.q_norm(q_lora)
|
||||
q, _ = self.wq_b(q_lora)
|
||||
|
||||
kv = (
|
||||
qkv_a[..., self.q_lora_rank :]
|
||||
if qkv_a is not None
|
||||
else self.wkv(x_linear)[0]
|
||||
)
|
||||
|
||||
from sglang.srt.layers.fused_qk_norm_rope_store import (
|
||||
fused_qk_norm_rope_swa_store,
|
||||
)
|
||||
|
||||
token_to_kv_pool = get_token_to_kv_pool()
|
||||
swa_loc = token_to_kv_pool.translate_loc_from_full_to_swa(
|
||||
forward_batch.out_cache_loc
|
||||
)
|
||||
swa_cache = token_to_kv_pool.swa_kv_pool.kv_buffer[self.layer_id]
|
||||
swa_page_size = token_to_kv_pool.swa_kv_pool.page_size
|
||||
|
||||
q = fused_qk_norm_rope_swa_store(
|
||||
q=q,
|
||||
kv=kv,
|
||||
q_norm_weight=None,
|
||||
kv_norm_weight=self.kv_norm.weight,
|
||||
q_rms_eps=self.eps,
|
||||
kv_rms_eps=self.eps,
|
||||
rope_head_dim=self.qk_rope_head_dim,
|
||||
cos_cache=self.cos_cache,
|
||||
sin_cache=self.sin_cache,
|
||||
positions=positions,
|
||||
swa_cache=swa_cache,
|
||||
swa_loc=swa_loc,
|
||||
swa_page_size=swa_page_size,
|
||||
q_out=q_out,
|
||||
dtype=x.dtype,
|
||||
)
|
||||
else:
|
||||
q_lora = self.q_norm(q_lora)
|
||||
q = self._compute_q_b(q_lora, positions, q_out)
|
||||
self._compute_kv_to_cache(x_linear, positions, forward_batch, qkv_a=qkv_a)
|
||||
|
||||
del qkv_a
|
||||
|
||||
if self.indexer is not None:
|
||||
current_stream.wait_stream(stream_compressor)
|
||||
if stream_indexer_compressor is not None:
|
||||
current_stream.wait_stream(stream_indexer_compressor)
|
||||
self.indexer(
|
||||
x=x,
|
||||
q_lora=q_lora,
|
||||
forward_batch=forward_batch,
|
||||
skip_compressor=True,
|
||||
)
|
||||
elif self.compressor is not None:
|
||||
current_stream.wait_stream(stream_compressor)
|
||||
|
||||
return q
|
||||
|
||||
def _forward_prepare(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
@@ -695,14 +807,24 @@ class MQALayer(nn.Module):
|
||||
if enable_multi_stream:
|
||||
# Multi-stream path always fuses cache write into the K kernel,
|
||||
# so the bf16 KV intermediate is gone.
|
||||
q = self._forward_prepare_multi_stream(
|
||||
x,
|
||||
positions,
|
||||
forward_batch,
|
||||
attn_backend,
|
||||
q_out,
|
||||
x_quant=x_quant,
|
||||
)
|
||||
if _is_hip:
|
||||
q = self._forward_prepare_multi_stream_hip(
|
||||
x,
|
||||
positions,
|
||||
forward_batch,
|
||||
attn_backend,
|
||||
q_out,
|
||||
x_quant=x_quant,
|
||||
)
|
||||
else:
|
||||
q = self._forward_prepare_multi_stream(
|
||||
x,
|
||||
positions,
|
||||
forward_batch,
|
||||
attn_backend,
|
||||
q_out,
|
||||
x_quant=x_quant,
|
||||
)
|
||||
kv = None
|
||||
else:
|
||||
q, kv = self._forward_prepare(
|
||||
@@ -792,12 +914,20 @@ class DeepseekV4DecoderLayer(nn.Module):
|
||||
alt_streams=alt_streams,
|
||||
compress_ratio_override=compress_ratio_override,
|
||||
)
|
||||
moe_alt_stream = (
|
||||
alt_streams[0]
|
||||
if (
|
||||
alt_streams is not None
|
||||
and (_is_cuda or envs.SGLANG_ROCM_USE_MULTI_STREAM.get())
|
||||
)
|
||||
else None
|
||||
)
|
||||
self.mlp = deepseek_v2.DeepseekV2MoE(
|
||||
config=config,
|
||||
quant_config=moe_quant_config_override or quant_config,
|
||||
prefix=add_prefix("mlp", prefix),
|
||||
layer_id=self.layer_id,
|
||||
alt_stream=alt_streams[0] if alt_streams is not None else None,
|
||||
alt_stream=moe_alt_stream,
|
||||
is_nextn=is_nextn,
|
||||
is_deepseek_v4=True,
|
||||
)
|
||||
@@ -1147,7 +1277,19 @@ class DeepseekV4Model(nn.Module):
|
||||
else:
|
||||
self.embed_tokens = PPMissingLayer()
|
||||
self.rms_norm_eps = config.rms_norm_eps
|
||||
self.alt_streams = [torch.cuda.Stream() for _ in range(5)] if _is_cuda else None
|
||||
use_stream_pool = _is_cuda or (
|
||||
_is_hip
|
||||
and (
|
||||
envs.SGLANG_ROCM_USE_MULTI_STREAM.get()
|
||||
or envs.SGLANG_OPT_USE_MULTI_STREAM_OVERLAP.get()
|
||||
)
|
||||
)
|
||||
num_alt_streams = 5 if _is_cuda else 2
|
||||
self.alt_streams = (
|
||||
[torch.cuda.Stream() for _ in range(num_alt_streams)]
|
||||
if use_stream_pool
|
||||
else None
|
||||
)
|
||||
self.layers, self.start_layer, self.end_layer = make_layers(
|
||||
config.num_hidden_layers,
|
||||
lambda idx, prefix: DeepseekV4DecoderLayer(
|
||||
|
||||
Reference in New Issue
Block a user