Fix Qwen3.5 GDN multi-item scoring (#33922)
Co-authored-by: Po-Han Huang (NVIDIA) <53919306+nvpohanh@users.noreply.github.com>
This commit is contained in:
@@ -44,6 +44,7 @@ def chunk_gated_delta_rule_fwd(
|
|||||||
initial_state_indices: torch.Tensor,
|
initial_state_indices: torch.Tensor,
|
||||||
cu_seqlens: Optional[torch.LongTensor] = None,
|
cu_seqlens: Optional[torch.LongTensor] = None,
|
||||||
chunk_indices: torch.LongTensor | None = None,
|
chunk_indices: torch.LongTensor | None = None,
|
||||||
|
inplace_update: bool = True,
|
||||||
):
|
):
|
||||||
g = chunk_local_cumsum(
|
g = chunk_local_cumsum(
|
||||||
g, chunk_size=CHUNK_SIZE, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices
|
g, chunk_size=CHUNK_SIZE, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices
|
||||||
@@ -68,6 +69,7 @@ def chunk_gated_delta_rule_fwd(
|
|||||||
initial_state_indices=initial_state_indices,
|
initial_state_indices=initial_state_indices,
|
||||||
cu_seqlens=cu_seqlens,
|
cu_seqlens=cu_seqlens,
|
||||||
chunk_indices=chunk_indices,
|
chunk_indices=chunk_indices,
|
||||||
|
inplace_update=inplace_update,
|
||||||
)
|
)
|
||||||
o = chunk_fwd_o(
|
o = chunk_fwd_o(
|
||||||
q=q,
|
q=q,
|
||||||
@@ -100,6 +102,7 @@ class ChunkGatedDeltaRuleFunction(torch.autograd.Function):
|
|||||||
initial_state_indices: torch.Tensor,
|
initial_state_indices: torch.Tensor,
|
||||||
cu_seqlens: Optional[torch.LongTensor] = None,
|
cu_seqlens: Optional[torch.LongTensor] = None,
|
||||||
use_qk_l2norm_in_kernel: bool = False,
|
use_qk_l2norm_in_kernel: bool = False,
|
||||||
|
inplace_update: bool = True,
|
||||||
):
|
):
|
||||||
q_orig = q
|
q_orig = q
|
||||||
k_orig = k
|
k_orig = k
|
||||||
@@ -124,6 +127,7 @@ class ChunkGatedDeltaRuleFunction(torch.autograd.Function):
|
|||||||
initial_state_indices=initial_state_indices,
|
initial_state_indices=initial_state_indices,
|
||||||
cu_seqlens=cu_seqlens,
|
cu_seqlens=cu_seqlens,
|
||||||
chunk_indices=chunk_indices,
|
chunk_indices=chunk_indices,
|
||||||
|
inplace_update=inplace_update,
|
||||||
)
|
)
|
||||||
return o.to(q.dtype), h
|
return o.to(q.dtype), h
|
||||||
|
|
||||||
@@ -141,6 +145,7 @@ def chunk_gated_delta_rule(
|
|||||||
cu_seqlens: Optional[torch.LongTensor] = None,
|
cu_seqlens: Optional[torch.LongTensor] = None,
|
||||||
head_first: bool = False,
|
head_first: bool = False,
|
||||||
use_qk_l2norm_in_kernel: bool = False,
|
use_qk_l2norm_in_kernel: bool = False,
|
||||||
|
inplace_update: bool = True,
|
||||||
):
|
):
|
||||||
r"""
|
r"""
|
||||||
Args:
|
Args:
|
||||||
@@ -169,6 +174,8 @@ def chunk_gated_delta_rule(
|
|||||||
head_first (Optional[bool]):
|
head_first (Optional[bool]):
|
||||||
Whether the inputs are in the head-first format, which is not supported for variable-length inputs.
|
Whether the inputs are in the head-first format, which is not supported for variable-length inputs.
|
||||||
Default: `False`.
|
Default: `False`.
|
||||||
|
inplace_update (Optional[bool]):
|
||||||
|
Whether to write final states back to `initial_state`. Default: `True`.
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
o (torch.Tensor):
|
o (torch.Tensor):
|
||||||
@@ -255,6 +262,7 @@ def chunk_gated_delta_rule(
|
|||||||
initial_state_indices,
|
initial_state_indices,
|
||||||
cu_seqlens,
|
cu_seqlens,
|
||||||
use_qk_l2norm_in_kernel,
|
use_qk_l2norm_in_kernel,
|
||||||
|
inplace_update,
|
||||||
)
|
)
|
||||||
if head_first:
|
if head_first:
|
||||||
o = rearrange(o, "b t h ... -> b h t ...")
|
o = rearrange(o, "b t h ... -> b h t ...")
|
||||||
|
|||||||
@@ -360,6 +360,7 @@ def chunk_gated_delta_rule_fwd_h(
|
|||||||
cu_seqlens: Optional[torch.LongTensor] = None,
|
cu_seqlens: Optional[torch.LongTensor] = None,
|
||||||
chunk_indices: Optional[torch.LongTensor] = None,
|
chunk_indices: Optional[torch.LongTensor] = None,
|
||||||
use_exp2: bool = False,
|
use_exp2: bool = False,
|
||||||
|
inplace_update: bool = True,
|
||||||
track_state: Optional[torch.Tensor] = None,
|
track_state: Optional[torch.Tensor] = None,
|
||||||
track_chunk_idx: Optional[torch.Tensor] = None,
|
track_chunk_idx: Optional[torch.Tensor] = None,
|
||||||
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||||
@@ -426,7 +427,7 @@ def chunk_gated_delta_rule_fwd_h(
|
|||||||
USE_G=g is not None,
|
USE_G=g is not None,
|
||||||
USE_GK=gk is not None,
|
USE_GK=gk is not None,
|
||||||
USE_INITIAL_STATE=initial_state is not None,
|
USE_INITIAL_STATE=initial_state is not None,
|
||||||
INPLACE_UPDATE=True,
|
INPLACE_UPDATE=inplace_update,
|
||||||
SAVE_NEW_VALUE=v_new is not None,
|
SAVE_NEW_VALUE=v_new is not None,
|
||||||
IS_VARLEN=cu_seqlens is not None,
|
IS_VARLEN=cu_seqlens is not None,
|
||||||
NT_BUCKET=(0 if NT <= 32 else (1 if NT <= 128 else 2)),
|
NT_BUCKET=(0 if NT <= 32 else (1 if NT <= 128 else 2)),
|
||||||
|
|||||||
@@ -242,7 +242,12 @@ def chunk_gated_delta_rule_fwd_h(
|
|||||||
cu_seqlens: Optional[torch.LongTensor] = None,
|
cu_seqlens: Optional[torch.LongTensor] = None,
|
||||||
chunk_indices: Optional[torch.LongTensor] = None,
|
chunk_indices: Optional[torch.LongTensor] = None,
|
||||||
use_exp2: bool = False,
|
use_exp2: bool = False,
|
||||||
|
inplace_update: bool = True,
|
||||||
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||||
|
if not inplace_update:
|
||||||
|
raise NotImplementedError(
|
||||||
|
"GDN multi-item scoring is not supported by the XPU chunk kernel"
|
||||||
|
)
|
||||||
assert not (use_exp2 and g is not None), (
|
assert not (use_exp2 and g is not None), (
|
||||||
"use_exp2 covers only the per-channel gk path; scalar g stays natural-exp"
|
"use_exp2 covers only the per-channel gk path; scalar g stays natural-exp"
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -47,12 +47,24 @@ _validate_mamba_replay_state_indices = (
|
|||||||
|
|
||||||
|
|
||||||
class MambaAttnBackendBase(AttentionBackend):
|
class MambaAttnBackendBase(AttentionBackend):
|
||||||
|
supports_mis: bool = False
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def validate_mis_support(cls, server_args) -> None:
|
||||||
|
if server_args.enable_mis and not cls.supports_mis:
|
||||||
|
raise ValueError(
|
||||||
|
f"{cls.__name__} does not support multi-item scoring. "
|
||||||
|
"Hybrid models require a linear-attention backend that explicitly "
|
||||||
|
"declares MIS support."
|
||||||
|
)
|
||||||
|
|
||||||
# Per-slot accept lengths for the KDA fused-accept spec path; allocated only
|
# Per-slot accept lengths for the KDA fused-accept spec path; allocated only
|
||||||
# by KDAAttnBackend where `_can_fuse_accept_state` holds. None everywhere
|
# by KDAAttnBackend where `_can_fuse_accept_state` holds. None everywhere
|
||||||
# else — update_mamba_state_after_mtp_verify keys the fused branch on it.
|
# else — update_mamba_state_after_mtp_verify keys the fused branch on it.
|
||||||
accept_lens_pool: Optional[torch.Tensor] = None
|
accept_lens_pool: Optional[torch.Tensor] = None
|
||||||
|
|
||||||
def __init__(self, model_runner: ModelRunner):
|
def __init__(self, model_runner: ModelRunner):
|
||||||
|
self.validate_mis_support(model_runner.server_args)
|
||||||
super().__init__()
|
super().__init__()
|
||||||
self.pad_slot_id = PAD_SLOT_ID
|
self.pad_slot_id = PAD_SLOT_ID
|
||||||
self.device = model_runner.device
|
self.device = model_runner.device
|
||||||
|
|||||||
@@ -1,5 +1,6 @@
|
|||||||
from typing import Optional, Tuple, Union
|
from typing import Optional, Tuple, Union
|
||||||
|
|
||||||
|
import msgspec
|
||||||
import torch
|
import torch
|
||||||
|
|
||||||
from sglang.kernels.ops.attention.fla.fused_gdn_gating import fused_gdn_gating
|
from sglang.kernels.ops.attention.fla.fused_gdn_gating import fused_gdn_gating
|
||||||
@@ -49,6 +50,117 @@ _fused_decode_verify_real_tensors = (
|
|||||||
envs.SGLANG_GDN_DECODE_FUSION_VERIFY_REAL_TENSORS.get()
|
envs.SGLANG_GDN_DECODE_FUSION_VERIFY_REAL_TENSORS.get()
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class GDNMISMetadata(msgspec.Struct, frozen=True):
|
||||||
|
query_token_indices: torch.Tensor
|
||||||
|
query_cu_seqlens: torch.Tensor
|
||||||
|
query_seq_lens_cpu: list[int]
|
||||||
|
query_request_indices: torch.Tensor
|
||||||
|
item_token_indices: torch.Tensor
|
||||||
|
item_cu_seqlens: torch.Tensor
|
||||||
|
item_seq_lens_cpu: list[int]
|
||||||
|
item_request_indices: torch.Tensor
|
||||||
|
|
||||||
|
|
||||||
|
def build_gdn_mis_metadata(forward_batch: ForwardBatch) -> GDNMISMetadata:
|
||||||
|
"""Build compact query/item segments from request-local MIS delimiters."""
|
||||||
|
if not forward_batch.is_prefill_only:
|
||||||
|
raise ValueError("GDN MIS is only supported for prefill-only requests")
|
||||||
|
|
||||||
|
prefix_lens = forward_batch.extend_prefix_lens_cpu
|
||||||
|
if isinstance(prefix_lens, torch.Tensor):
|
||||||
|
prefix_lens = prefix_lens.tolist()
|
||||||
|
if any(int(prefix_len) != 0 for prefix_len in prefix_lens):
|
||||||
|
raise ValueError("GDN MIS does not support cached prefixes")
|
||||||
|
|
||||||
|
seq_lens = forward_batch.extend_seq_lens_cpu
|
||||||
|
if isinstance(seq_lens, torch.Tensor):
|
||||||
|
seq_lens = seq_lens.tolist()
|
||||||
|
seq_lens = [int(seq_len) for seq_len in seq_lens]
|
||||||
|
delimiter_indices = forward_batch.multi_item_delimiter_indices
|
||||||
|
if delimiter_indices is None or len(delimiter_indices) != len(seq_lens):
|
||||||
|
raise ValueError("GDN MIS requires delimiter indices for every request")
|
||||||
|
if sum(seq_lens) > forward_batch.input_ids.numel():
|
||||||
|
raise ValueError("GDN MIS sequence lengths exceed the input tokens")
|
||||||
|
|
||||||
|
query_token_indices: list[int] = []
|
||||||
|
query_seq_lens_cpu: list[int] = []
|
||||||
|
query_request_indices: list[int] = []
|
||||||
|
item_token_indices: list[int] = []
|
||||||
|
item_seq_lens_cpu: list[int] = []
|
||||||
|
item_request_indices: list[int] = []
|
||||||
|
|
||||||
|
request_start = 0
|
||||||
|
for request_idx, (seq_len, request_delimiters) in enumerate(
|
||||||
|
zip(seq_lens, delimiter_indices)
|
||||||
|
):
|
||||||
|
delimiters = [int(index) for index in request_delimiters.tolist()]
|
||||||
|
if len(delimiters) < 2:
|
||||||
|
raise ValueError("GDN MIS requires at least two delimiters per request")
|
||||||
|
if any(
|
||||||
|
current >= following
|
||||||
|
for current, following in zip(delimiters, delimiters[1:])
|
||||||
|
):
|
||||||
|
raise ValueError("GDN MIS delimiter indices must be strictly increasing")
|
||||||
|
if delimiters[0] < 0 or delimiters[-1] >= seq_len:
|
||||||
|
raise ValueError("GDN MIS delimiter index is outside the request")
|
||||||
|
if delimiters[-1] != seq_len - 1:
|
||||||
|
raise ValueError("GDN MIS final delimiter must be the last request token")
|
||||||
|
|
||||||
|
query_len = delimiters[0]
|
||||||
|
if query_len > 0:
|
||||||
|
query_token_indices.extend(range(request_start, request_start + query_len))
|
||||||
|
query_seq_lens_cpu.append(query_len)
|
||||||
|
query_request_indices.append(request_idx)
|
||||||
|
|
||||||
|
branch_ends = delimiters[1:] + [seq_len]
|
||||||
|
for branch_start, branch_end in zip(delimiters, branch_ends):
|
||||||
|
branch_len = branch_end - branch_start
|
||||||
|
item_token_indices.extend(
|
||||||
|
range(request_start + branch_start, request_start + branch_end)
|
||||||
|
)
|
||||||
|
item_seq_lens_cpu.append(branch_len)
|
||||||
|
item_request_indices.append(request_idx)
|
||||||
|
|
||||||
|
request_start += seq_len
|
||||||
|
|
||||||
|
device = forward_batch.input_ids.device
|
||||||
|
|
||||||
|
def _indices(values: list[int], dtype: torch.dtype) -> torch.Tensor:
|
||||||
|
return torch.tensor(values, dtype=dtype, device=device)
|
||||||
|
|
||||||
|
def _cu_seqlens(lengths: list[int]) -> torch.Tensor:
|
||||||
|
result = torch.zeros(len(lengths) + 1, dtype=torch.int32, device=device)
|
||||||
|
if lengths:
|
||||||
|
result[1:] = torch.tensor(lengths, dtype=torch.int32, device=device).cumsum(
|
||||||
|
dim=0
|
||||||
|
)
|
||||||
|
return result
|
||||||
|
|
||||||
|
return GDNMISMetadata(
|
||||||
|
query_token_indices=_indices(query_token_indices, torch.int64),
|
||||||
|
query_cu_seqlens=_cu_seqlens(query_seq_lens_cpu),
|
||||||
|
query_seq_lens_cpu=query_seq_lens_cpu,
|
||||||
|
query_request_indices=_indices(query_request_indices, torch.int64),
|
||||||
|
item_token_indices=_indices(item_token_indices, torch.int64),
|
||||||
|
item_cu_seqlens=_cu_seqlens(item_seq_lens_cpu),
|
||||||
|
item_seq_lens_cpu=item_seq_lens_cpu,
|
||||||
|
item_request_indices=_indices(item_request_indices, torch.int64),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def validate_gdn_mis_backend(prefill_backend: LinearAttnKernelBackend) -> None:
|
||||||
|
if not get_exec().features.enable_mis:
|
||||||
|
return
|
||||||
|
if not prefill_backend.is_triton():
|
||||||
|
raise ValueError(
|
||||||
|
"GDN multi-item scoring requires the Triton linear-attention prefill "
|
||||||
|
"backend. Set --linear-attn-prefill-backend triton."
|
||||||
|
)
|
||||||
|
if get_memory().enable_page_major_kv_layout:
|
||||||
|
raise ValueError("GDN multi-item scoring does not support page-major layout")
|
||||||
|
|
||||||
|
|
||||||
if is_cuda():
|
if is_cuda():
|
||||||
from sglang.srt.layers.attention.mamba.causal_conv1d import (
|
from sglang.srt.layers.attention.mamba.causal_conv1d import (
|
||||||
causal_conv1d_fn as causal_conv1d_fn_cuda,
|
causal_conv1d_fn as causal_conv1d_fn_cuda,
|
||||||
@@ -391,10 +503,13 @@ class GDNAttnBackend(MambaAttnBackendBase):
|
|||||||
"""Attention backend for GDN (Gated Delta Network) linear attention."""
|
"""Attention backend for GDN (Gated Delta Network) linear attention."""
|
||||||
|
|
||||||
needs_cpu_seq_lens: bool = False
|
needs_cpu_seq_lens: bool = False
|
||||||
|
supports_mis: bool = True
|
||||||
|
|
||||||
def __init__(self, model_runner: ModelRunner):
|
def __init__(self, model_runner: ModelRunner):
|
||||||
_validate_gdn_linear_attn_backends(model_runner.linear_attn_backends)
|
_validate_gdn_linear_attn_backends(model_runner.linear_attn_backends)
|
||||||
super().__init__(model_runner)
|
super().__init__(model_runner)
|
||||||
|
self.enable_mis = get_exec().features.enable_mis
|
||||||
|
self.mis_metadata: Optional[GDNMISMetadata] = None
|
||||||
self.conv_states_shape = (
|
self.conv_states_shape = (
|
||||||
model_runner.req_to_token_pool.mamba_pool.mamba_cache.conv[0].shape
|
model_runner.req_to_token_pool.mamba_pool.mamba_cache.conv[0].shape
|
||||||
)
|
)
|
||||||
@@ -404,6 +519,7 @@ class GDNAttnBackend(MambaAttnBackendBase):
|
|||||||
)
|
)
|
||||||
|
|
||||||
backends = model_runner.linear_attn_backends
|
backends = model_runner.linear_attn_backends
|
||||||
|
validate_gdn_mis_backend(backends.prefill)
|
||||||
self.linear_attn_backends = backends
|
self.linear_attn_backends = backends
|
||||||
self.kernel_dispatcher = GDNKernelDispatcher(
|
self.kernel_dispatcher = GDNKernelDispatcher(
|
||||||
backends.decode, backends.prefill, backends.verify
|
backends.decode, backends.prefill, backends.verify
|
||||||
@@ -418,6 +534,11 @@ class GDNAttnBackend(MambaAttnBackendBase):
|
|||||||
|
|
||||||
def init_forward_metadata(self, forward_batch: ForwardBatch):
|
def init_forward_metadata(self, forward_batch: ForwardBatch):
|
||||||
super().init_forward_metadata(forward_batch)
|
super().init_forward_metadata(forward_batch)
|
||||||
|
self.mis_metadata = None
|
||||||
|
if forward_batch.multi_item_delimiter_indices is not None:
|
||||||
|
if not self.enable_mis:
|
||||||
|
raise ValueError("GDN MIS metadata requires --enable-mis")
|
||||||
|
self.mis_metadata = build_gdn_mis_metadata(forward_batch)
|
||||||
if self.forward_metadata.has_mamba_track_mask:
|
if self.forward_metadata.has_mamba_track_mask:
|
||||||
self.forward_metadata.mamba_track_mask_indices = (
|
self.forward_metadata.mamba_track_mask_indices = (
|
||||||
forward_batch.mamba_track_mask.nonzero(as_tuple=True)[0]
|
forward_batch.mamba_track_mask.nonzero(as_tuple=True)[0]
|
||||||
@@ -710,6 +831,18 @@ class GDNAttnBackend(MambaAttnBackendBase):
|
|||||||
mamba_cache_params = self.req_to_token_pool.mamba2_layer_cache(layer.layer_id)
|
mamba_cache_params = self.req_to_token_pool.mamba2_layer_cache(layer.layer_id)
|
||||||
conv_states = mamba_cache_params.conv[0]
|
conv_states = mamba_cache_params.conv[0]
|
||||||
ssm_states = mamba_cache_params.temporal
|
ssm_states = mamba_cache_params.temporal
|
||||||
|
if self.mis_metadata is not None:
|
||||||
|
if is_target_verify:
|
||||||
|
raise ValueError("GDN MIS does not support target verify")
|
||||||
|
return self._forward_extend_mis(
|
||||||
|
layer=layer,
|
||||||
|
mixed_qkv=mixed_qkv,
|
||||||
|
a=a,
|
||||||
|
b=b,
|
||||||
|
conv_states=conv_states,
|
||||||
|
ssm_states=ssm_states,
|
||||||
|
cache_indices=cache_indices,
|
||||||
|
)
|
||||||
if is_target_verify:
|
if is_target_verify:
|
||||||
assert isinstance(mamba_cache_params, MambaPool.SpeculativeState)
|
assert isinstance(mamba_cache_params, MambaPool.SpeculativeState)
|
||||||
intermediate_state_cache = mamba_cache_params.intermediate_ssm
|
intermediate_state_cache = mamba_cache_params.intermediate_ssm
|
||||||
@@ -923,6 +1056,138 @@ class GDNAttnBackend(MambaAttnBackendBase):
|
|||||||
|
|
||||||
return core_attn_out
|
return core_attn_out
|
||||||
|
|
||||||
|
def _forward_extend_mis(
|
||||||
|
self,
|
||||||
|
*,
|
||||||
|
layer: RadixLinearAttention,
|
||||||
|
mixed_qkv: torch.Tensor,
|
||||||
|
a: torch.Tensor,
|
||||||
|
b: torch.Tensor,
|
||||||
|
conv_states: torch.Tensor,
|
||||||
|
ssm_states: torch.Tensor,
|
||||||
|
cache_indices: torch.Tensor,
|
||||||
|
) -> torch.Tensor:
|
||||||
|
metadata = self.mis_metadata
|
||||||
|
assert metadata is not None
|
||||||
|
|
||||||
|
conv_states[cache_indices] = 0
|
||||||
|
ssm_states[cache_indices] = 0
|
||||||
|
output = mixed_qkv.new_zeros(
|
||||||
|
1, mixed_qkv.shape[0], layer.num_v_heads, layer.head_v_dim
|
||||||
|
)
|
||||||
|
|
||||||
|
if metadata.query_token_indices.numel() > 0:
|
||||||
|
query_cache_indices = cache_indices[metadata.query_request_indices]
|
||||||
|
query_output = self._forward_mis_segments(
|
||||||
|
layer=layer,
|
||||||
|
mixed_qkv=mixed_qkv[metadata.query_token_indices],
|
||||||
|
a=a[metadata.query_token_indices],
|
||||||
|
b=b[metadata.query_token_indices],
|
||||||
|
conv_states=conv_states,
|
||||||
|
conv_cache_indices=query_cache_indices,
|
||||||
|
has_initial_states=torch.zeros(
|
||||||
|
len(metadata.query_seq_lens_cpu),
|
||||||
|
dtype=torch.bool,
|
||||||
|
device=mixed_qkv.device,
|
||||||
|
),
|
||||||
|
query_start_loc=metadata.query_cu_seqlens,
|
||||||
|
seq_lens_cpu=metadata.query_seq_lens_cpu,
|
||||||
|
ssm_states=ssm_states,
|
||||||
|
ssm_cache_indices=query_cache_indices,
|
||||||
|
inplace_update=True,
|
||||||
|
)
|
||||||
|
output[:, metadata.query_token_indices] = query_output
|
||||||
|
|
||||||
|
item_ssm_indices = cache_indices[metadata.item_request_indices]
|
||||||
|
item_conv_states = conv_states[item_ssm_indices]
|
||||||
|
item_conv_indices = torch.arange(
|
||||||
|
item_ssm_indices.shape[0],
|
||||||
|
dtype=cache_indices.dtype,
|
||||||
|
device=cache_indices.device,
|
||||||
|
)
|
||||||
|
item_output = self._forward_mis_segments(
|
||||||
|
layer=layer,
|
||||||
|
mixed_qkv=mixed_qkv[metadata.item_token_indices],
|
||||||
|
a=a[metadata.item_token_indices],
|
||||||
|
b=b[metadata.item_token_indices],
|
||||||
|
conv_states=item_conv_states,
|
||||||
|
conv_cache_indices=item_conv_indices,
|
||||||
|
has_initial_states=torch.ones(
|
||||||
|
len(metadata.item_seq_lens_cpu),
|
||||||
|
dtype=torch.bool,
|
||||||
|
device=mixed_qkv.device,
|
||||||
|
),
|
||||||
|
query_start_loc=metadata.item_cu_seqlens,
|
||||||
|
seq_lens_cpu=metadata.item_seq_lens_cpu,
|
||||||
|
ssm_states=ssm_states,
|
||||||
|
ssm_cache_indices=item_ssm_indices,
|
||||||
|
inplace_update=False,
|
||||||
|
)
|
||||||
|
output[:, metadata.item_token_indices] = item_output
|
||||||
|
return output
|
||||||
|
|
||||||
|
def _forward_mis_segments(
|
||||||
|
self,
|
||||||
|
*,
|
||||||
|
layer: RadixLinearAttention,
|
||||||
|
mixed_qkv: torch.Tensor,
|
||||||
|
a: torch.Tensor,
|
||||||
|
b: torch.Tensor,
|
||||||
|
conv_states: torch.Tensor,
|
||||||
|
conv_cache_indices: torch.Tensor,
|
||||||
|
has_initial_states: torch.Tensor,
|
||||||
|
query_start_loc: torch.Tensor,
|
||||||
|
seq_lens_cpu: list[int],
|
||||||
|
ssm_states: torch.Tensor,
|
||||||
|
ssm_cache_indices: torch.Tensor,
|
||||||
|
inplace_update: bool,
|
||||||
|
) -> torch.Tensor:
|
||||||
|
mixed_qkv = causal_conv1d_fn(
|
||||||
|
mixed_qkv.transpose(0, 1),
|
||||||
|
layer.conv_weights,
|
||||||
|
layer.bias,
|
||||||
|
activation=layer.activation,
|
||||||
|
conv_states=conv_states,
|
||||||
|
has_initial_state=has_initial_states,
|
||||||
|
cache_indices=conv_cache_indices,
|
||||||
|
query_start_loc=query_start_loc,
|
||||||
|
seq_lens_cpu=seq_lens_cpu,
|
||||||
|
).transpose(0, 1)
|
||||||
|
|
||||||
|
qkv_dim = layer.q_dim + layer.k_dim + layer.v_dim
|
||||||
|
if (is_cuda() or is_hip()) and qkv_dim <= MAX_FUSED_QKV_SPLIT_DIM:
|
||||||
|
query, key, value = fused_qkv_split_gdn_prefill(
|
||||||
|
mixed_qkv,
|
||||||
|
layer.num_q_heads,
|
||||||
|
layer.num_k_heads,
|
||||||
|
layer.num_v_heads,
|
||||||
|
layer.head_q_dim,
|
||||||
|
layer.head_k_dim,
|
||||||
|
layer.head_v_dim,
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
query, key, value = torch.split(
|
||||||
|
mixed_qkv, [layer.q_dim, layer.k_dim, layer.v_dim], dim=-1
|
||||||
|
)
|
||||||
|
num_tokens = mixed_qkv.shape[0]
|
||||||
|
query = query.view(1, num_tokens, layer.num_q_heads, layer.head_q_dim)
|
||||||
|
key = key.view(1, num_tokens, layer.num_k_heads, layer.head_k_dim)
|
||||||
|
value = value.view(1, num_tokens, layer.num_v_heads, layer.head_v_dim)
|
||||||
|
|
||||||
|
g, beta = fused_gdn_gating(layer.A_log, a, b, layer.dt_bias)
|
||||||
|
output, _, _ = self.kernel_dispatcher.extend(
|
||||||
|
q=query,
|
||||||
|
k=key,
|
||||||
|
v=value,
|
||||||
|
g=g,
|
||||||
|
beta=beta,
|
||||||
|
ssm_states=ssm_states,
|
||||||
|
cache_indices=ssm_cache_indices,
|
||||||
|
query_start_loc=query_start_loc,
|
||||||
|
inplace_update=inplace_update,
|
||||||
|
)
|
||||||
|
return output
|
||||||
|
|
||||||
def _replayssm_fold_target_verify(
|
def _replayssm_fold_target_verify(
|
||||||
self,
|
self,
|
||||||
*,
|
*,
|
||||||
|
|||||||
@@ -177,13 +177,28 @@ class TritonGDNKernel(LinearAttnKernelBase):
|
|||||||
ssm_states: torch.Tensor,
|
ssm_states: torch.Tensor,
|
||||||
cache_indices: torch.Tensor,
|
cache_indices: torch.Tensor,
|
||||||
query_start_loc: torch.Tensor,
|
query_start_loc: torch.Tensor,
|
||||||
|
inplace_update: bool = True,
|
||||||
**kwargs,
|
**kwargs,
|
||||||
) -> tuple:
|
) -> tuple:
|
||||||
recurrent_state = ssm_states
|
recurrent_state = ssm_states
|
||||||
recurrent_state_indices_args = {"initial_state_indices": cache_indices}
|
recurrent_state_indices_args = {"initial_state_indices": cache_indices}
|
||||||
if is_npu():
|
inplace_update_args = {"inplace_update": inplace_update}
|
||||||
|
if is_cpu():
|
||||||
|
if not inplace_update:
|
||||||
|
raise NotImplementedError(
|
||||||
|
"GDN multi-item scoring is not supported by the CPU chunk kernel"
|
||||||
|
)
|
||||||
|
inplace_update_args = {}
|
||||||
|
elif is_npu():
|
||||||
|
if not inplace_update:
|
||||||
|
raise NotImplementedError(
|
||||||
|
"GDN multi-item scoring is not supported by the NPU chunk kernel"
|
||||||
|
)
|
||||||
recurrent_state = ssm_states[cache_indices]
|
recurrent_state = ssm_states[cache_indices]
|
||||||
recurrent_state_indices_args = {}
|
recurrent_state_indices_args = {}
|
||||||
|
# The external NPU kernel does not expose the optional write-back
|
||||||
|
# control. Its existing behavior is equivalent to True.
|
||||||
|
inplace_update_args = {}
|
||||||
|
|
||||||
return chunk_gated_delta_rule(
|
return chunk_gated_delta_rule(
|
||||||
q=q,
|
q=q,
|
||||||
@@ -196,6 +211,7 @@ class TritonGDNKernel(LinearAttnKernelBase):
|
|||||||
head_first=False,
|
head_first=False,
|
||||||
use_qk_l2norm_in_kernel=True,
|
use_qk_l2norm_in_kernel=True,
|
||||||
**recurrent_state_indices_args,
|
**recurrent_state_indices_args,
|
||||||
|
**inplace_update_args,
|
||||||
)
|
)
|
||||||
|
|
||||||
def target_verify(
|
def target_verify(
|
||||||
|
|||||||
@@ -30,7 +30,7 @@ from sglang.srt.model_executor.cuda_graph_config import (
|
|||||||
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
|
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
|
||||||
from sglang.srt.model_executor.forward_context import ForwardContext, forward_context
|
from sglang.srt.model_executor.forward_context import ForwardContext, forward_context
|
||||||
from sglang.srt.model_executor.model_runner import ModelRunner
|
from sglang.srt.model_executor.model_runner import ModelRunner
|
||||||
from sglang.srt.runtime_context import get_context, get_parallel
|
from sglang.srt.runtime_context import get_context, get_parallel, get_server_args
|
||||||
|
|
||||||
_parallel_override = get_parallel().override(attn_tp_size=1)
|
_parallel_override = get_parallel().override(attn_tp_size=1)
|
||||||
_parallel_override.__enter__()
|
_parallel_override.__enter__()
|
||||||
@@ -56,6 +56,8 @@ class GDNAttentionCase:
|
|||||||
prefix_lens: tuple[int, ...]
|
prefix_lens: tuple[int, ...]
|
||||||
extend_lens: tuple[int, ...] = ()
|
extend_lens: tuple[int, ...] = ()
|
||||||
linear_attn_prefill_backend: str | None = None
|
linear_attn_prefill_backend: str | None = None
|
||||||
|
mis_delimiter_indices: tuple[tuple[int, ...], ...] = ()
|
||||||
|
conv_history_weight: float = 0.0
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def batch_size(self) -> int:
|
def batch_size(self) -> int:
|
||||||
@@ -252,7 +254,7 @@ class MockGDNModelRunner(ModelRunner):
|
|||||||
dllm_algorithm=None,
|
dllm_algorithm=None,
|
||||||
dllm_algorithm_config=None,
|
dllm_algorithm_config=None,
|
||||||
enable_deterministic_inference=False,
|
enable_deterministic_inference=False,
|
||||||
enable_mis=False,
|
enable_mis=bool(case.mis_delimiter_indices),
|
||||||
linear_attn_backend="triton",
|
linear_attn_backend="triton",
|
||||||
linear_attn_decode_backend=None,
|
linear_attn_decode_backend=None,
|
||||||
linear_attn_prefill_backend=case.linear_attn_prefill_backend,
|
linear_attn_prefill_backend=case.linear_attn_prefill_backend,
|
||||||
@@ -268,7 +270,8 @@ class MockGDNModelRunner(ModelRunner):
|
|||||||
# derives it from hf_config + page_size, which needs a real model.
|
# derives it from hf_config + page_size, which needs a real model.
|
||||||
_mamba_cache_chunk_size=64,
|
_mamba_cache_chunk_size=64,
|
||||||
)
|
)
|
||||||
self.server_args = self._server_args_override.install()
|
self._server_args_override.install()
|
||||||
|
self.server_args = get_server_args()
|
||||||
cache_shape = Mamba2StateShape.create(
|
cache_shape = Mamba2StateShape.create(
|
||||||
tp_world_size=1,
|
tp_world_size=1,
|
||||||
intermediate_size=case.num_v_heads * head_v_dim,
|
intermediate_size=case.num_v_heads * head_v_dim,
|
||||||
@@ -364,6 +367,7 @@ class ProjectedGDNAttention(nn.Module):
|
|||||||
head_v_dim: int,
|
head_v_dim: int,
|
||||||
dtype: torch.dtype,
|
dtype: torch.dtype,
|
||||||
device: str,
|
device: str,
|
||||||
|
conv_history_weight: float = 0.0,
|
||||||
):
|
):
|
||||||
super().__init__()
|
super().__init__()
|
||||||
self.num_k_heads = num_k_heads
|
self.num_k_heads = num_k_heads
|
||||||
@@ -372,6 +376,7 @@ class ProjectedGDNAttention(nn.Module):
|
|||||||
self.head_v_dim = head_v_dim
|
self.head_v_dim = head_v_dim
|
||||||
mixed_qkv_dim = 2 * num_k_heads * head_k_dim + num_v_heads * head_v_dim
|
mixed_qkv_dim = 2 * num_k_heads * head_k_dim + num_v_heads * head_v_dim
|
||||||
conv_weights = torch.zeros(mixed_qkv_dim, 2, dtype=dtype, device=device)
|
conv_weights = torch.zeros(mixed_qkv_dim, 2, dtype=dtype, device=device)
|
||||||
|
conv_weights[:, 0] = conv_history_weight
|
||||||
conv_weights[:, 1] = 1
|
conv_weights[:, 1] = 1
|
||||||
self.A_log = nn.Parameter(
|
self.A_log = nn.Parameter(
|
||||||
torch.randn(num_v_heads, dtype=torch.float32, device=device) * 0.1
|
torch.randn(num_v_heads, dtype=torch.float32, device=device) * 0.1
|
||||||
@@ -547,6 +552,15 @@ def _make_forward_batch(
|
|||||||
out_cache_loc=torch.tensor(out_cache_locs, dtype=torch.int64, device=device),
|
out_cache_loc=torch.tensor(out_cache_locs, dtype=torch.int64, device=device),
|
||||||
seq_lens_sum=sum(seq_lens),
|
seq_lens_sum=sum(seq_lens),
|
||||||
positions=torch.tensor(positions, dtype=torch.int64, device=device),
|
positions=torch.tensor(positions, dtype=torch.int64, device=device),
|
||||||
|
is_prefill_only=bool(case.mis_delimiter_indices),
|
||||||
|
multi_item_delimiter_indices=(
|
||||||
|
[
|
||||||
|
torch.tensor(indices, dtype=torch.int64)
|
||||||
|
for indices in case.mis_delimiter_indices
|
||||||
|
]
|
||||||
|
if case.mis_delimiter_indices
|
||||||
|
else None
|
||||||
|
),
|
||||||
)
|
)
|
||||||
|
|
||||||
if case.forward_mode.is_extend(include_draft_extend_v2=True):
|
if case.forward_mode.is_extend(include_draft_extend_v2=True):
|
||||||
@@ -626,6 +640,7 @@ def build_gdn_attention_fixture(
|
|||||||
head_v_dim=head_v_dim,
|
head_v_dim=head_v_dim,
|
||||||
dtype=dtype,
|
dtype=dtype,
|
||||||
device=device,
|
device=device,
|
||||||
|
conv_history_weight=case.conv_history_weight,
|
||||||
)
|
)
|
||||||
reference_module = ReferenceGDNAttention(
|
reference_module = ReferenceGDNAttention(
|
||||||
num_k_heads=case.num_k_heads,
|
num_k_heads=case.num_k_heads,
|
||||||
@@ -636,6 +651,13 @@ def build_gdn_attention_fixture(
|
|||||||
device=device,
|
device=device,
|
||||||
)
|
)
|
||||||
_copy_gdn_parameters(actual_module, reference_module)
|
_copy_gdn_parameters(actual_module, reference_module)
|
||||||
|
if case.mis_delimiter_indices:
|
||||||
|
# Keep recurrent history strong so cross-item state leakage cannot hide
|
||||||
|
# behind the normal random gate decay in short focused test sequences.
|
||||||
|
with torch.no_grad():
|
||||||
|
actual_module.A_log.fill_(-4.0)
|
||||||
|
actual_module.dt_bias.fill_(-4.0)
|
||||||
|
_copy_gdn_parameters(actual_module, reference_module)
|
||||||
from .dense_attention import make_loc_fn as _dense_make_loc_fn
|
from .dense_attention import make_loc_fn as _dense_make_loc_fn
|
||||||
|
|
||||||
loc_fn = _dense_make_loc_fn(
|
loc_fn = _dense_make_loc_fn(
|
||||||
@@ -785,7 +807,8 @@ def _pure_torch_gdn_reference(
|
|||||||
initial_ssm_states: torch.Tensor,
|
initial_ssm_states: torch.Tensor,
|
||||||
) -> GDNReferenceOutput:
|
) -> GDNReferenceOutput:
|
||||||
module = fixture.reference_module
|
module = fixture.reference_module
|
||||||
q, k, v = module.split_qkv(fixture.mixed_qkv)
|
mixed_qkv = _pure_torch_mis_conv_reference(fixture)
|
||||||
|
q, k, v = module.split_qkv(mixed_qkv)
|
||||||
cache_indices = _cache_indices(fixture)
|
cache_indices = _cache_indices(fixture)
|
||||||
g, beta = _pure_torch_gdn_gating(module, fixture.a, fixture.b)
|
g, beta = _pure_torch_gdn_gating(module, fixture.a, fixture.b)
|
||||||
q = q.float()
|
q = q.float()
|
||||||
@@ -808,7 +831,21 @@ def _pure_torch_gdn_reference(
|
|||||||
state_idx = cache_indices[req_idx]
|
state_idx = cache_indices[req_idx]
|
||||||
state = initial_ssm_states[state_idx].float().clone()
|
state = initial_ssm_states[state_idx].float().clone()
|
||||||
|
|
||||||
for offset in range(input_len):
|
if fixture.case.mis_delimiter_indices:
|
||||||
|
delimiters = fixture.case.mis_delimiter_indices[req_idx]
|
||||||
|
segments = [(0, delimiters[0])]
|
||||||
|
segments.extend(zip(delimiters, tuple(delimiters[1:]) + (input_len,)))
|
||||||
|
else:
|
||||||
|
segments = [(0, input_len)]
|
||||||
|
|
||||||
|
query_final_state = state.clone()
|
||||||
|
for segment_idx, (segment_start, segment_end) in enumerate(segments):
|
||||||
|
state = (
|
||||||
|
query_final_state.clone()
|
||||||
|
if segment_idx > 0
|
||||||
|
else initial_ssm_states[state_idx].float().clone()
|
||||||
|
)
|
||||||
|
for offset in range(segment_start, segment_end):
|
||||||
token_idx = start + offset
|
token_idx = start + offset
|
||||||
for v_head in range(fixture.case.num_v_heads):
|
for v_head in range(fixture.case.num_v_heads):
|
||||||
k_head = v_head // q_head_ratio
|
k_head = v_head // q_head_ratio
|
||||||
@@ -822,15 +859,24 @@ def _pure_torch_gdn_reference(
|
|||||||
|
|
||||||
head_state = state[v_head]
|
head_state = state[v_head]
|
||||||
head_state = head_state * torch.exp(g[token_idx, v_head])
|
head_state = head_state * torch.exp(g[token_idx, v_head])
|
||||||
residual_v = v_vec - torch.sum(head_state * k_norm.unsqueeze(0), dim=1)
|
residual_v = v_vec - torch.sum(
|
||||||
|
head_state * k_norm.unsqueeze(0), dim=1
|
||||||
|
)
|
||||||
residual_v = residual_v * beta[token_idx, v_head]
|
residual_v = residual_v * beta[token_idx, v_head]
|
||||||
head_state = head_state + residual_v.unsqueeze(1) * k_norm.unsqueeze(0)
|
head_state = head_state + residual_v.unsqueeze(
|
||||||
|
1
|
||||||
|
) * k_norm.unsqueeze(0)
|
||||||
state[v_head] = head_state
|
state[v_head] = head_state
|
||||||
outputs[0, token_idx, v_head] = torch.sum(
|
outputs[0, token_idx, v_head] = torch.sum(
|
||||||
head_state * q_norm.unsqueeze(0), dim=1
|
head_state * q_norm.unsqueeze(0), dim=1
|
||||||
)
|
)
|
||||||
|
|
||||||
final_states[state_idx] = state.to(final_states.dtype)
|
if segment_idx == 0:
|
||||||
|
query_final_state = state.clone()
|
||||||
|
|
||||||
|
final_states[state_idx] = (
|
||||||
|
query_final_state if fixture.case.mis_delimiter_indices else state
|
||||||
|
).to(final_states.dtype)
|
||||||
start += input_len
|
start += input_len
|
||||||
|
|
||||||
return GDNReferenceOutput(
|
return GDNReferenceOutput(
|
||||||
@@ -839,6 +885,45 @@ def _pure_torch_gdn_reference(
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _pure_torch_mis_conv_reference(fixture: GDNAttentionFixture) -> torch.Tensor:
|
||||||
|
"""Token-by-token causal-conv reference with MIS query-state branching."""
|
||||||
|
if fixture.case.conv_history_weight == 0:
|
||||||
|
return fixture.mixed_qkv
|
||||||
|
|
||||||
|
weights = fixture.actual_module.attn.conv_weights.float()
|
||||||
|
width = weights.shape[1]
|
||||||
|
result = torch.empty_like(fixture.mixed_qkv)
|
||||||
|
request_start = 0
|
||||||
|
for request_idx, input_len in enumerate(fixture.case.input_lens):
|
||||||
|
request_input = fixture.mixed_qkv[
|
||||||
|
request_start : request_start + input_len
|
||||||
|
].float()
|
||||||
|
delimiters = fixture.case.mis_delimiter_indices[request_idx]
|
||||||
|
segments = [(0, delimiters[0])]
|
||||||
|
segments.extend(zip(delimiters, tuple(delimiters[1:]) + (input_len,)))
|
||||||
|
|
||||||
|
query_state = torch.zeros(
|
||||||
|
width - 1,
|
||||||
|
request_input.shape[1],
|
||||||
|
dtype=torch.float32,
|
||||||
|
device=request_input.device,
|
||||||
|
)
|
||||||
|
for segment_idx, (segment_start, segment_end) in enumerate(segments):
|
||||||
|
state = query_state.clone()
|
||||||
|
for token_offset in range(segment_start, segment_end):
|
||||||
|
window = torch.cat(
|
||||||
|
[state, request_input[token_offset : token_offset + 1]]
|
||||||
|
)
|
||||||
|
result[request_start + token_offset] = torch.sum(
|
||||||
|
window.transpose(0, 1) * weights, dim=1
|
||||||
|
).to(result.dtype)
|
||||||
|
state = window[1:]
|
||||||
|
if segment_idx == 0:
|
||||||
|
query_state = state
|
||||||
|
request_start += input_len
|
||||||
|
return result
|
||||||
|
|
||||||
|
|
||||||
def make_gdn_case_with_prefix_lens(
|
def make_gdn_case_with_prefix_lens(
|
||||||
case: GDNAttentionCase,
|
case: GDNAttentionCase,
|
||||||
name: str,
|
name: str,
|
||||||
@@ -864,6 +949,8 @@ def make_gdn_case_with_prefix_lens(
|
|||||||
page_size=case.page_size,
|
page_size=case.page_size,
|
||||||
prefix_lens=prefix_lens,
|
prefix_lens=prefix_lens,
|
||||||
extend_lens=extend_lens,
|
extend_lens=extend_lens,
|
||||||
|
mis_delimiter_indices=case.mis_delimiter_indices,
|
||||||
|
conv_history_weight=case.conv_history_weight,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
@@ -1117,7 +1204,7 @@ def run_gdn_attention_case(
|
|||||||
expected = _pure_torch_gdn_reference(fixture, initial_ssm_states)
|
expected = _pure_torch_gdn_reference(fixture, initial_ssm_states)
|
||||||
|
|
||||||
torch.testing.assert_close(actual, expected.output, atol=GDN_ATOL, rtol=GDN_RTOL)
|
torch.testing.assert_close(actual, expected.output, atol=GDN_ATOL, rtol=GDN_RTOL)
|
||||||
if case.forward_mode.is_decode():
|
if case.forward_mode.is_decode() or case.mis_delimiter_indices:
|
||||||
torch.testing.assert_close(
|
torch.testing.assert_close(
|
||||||
_ssm_states(fixture)[_cache_indices(fixture)],
|
_ssm_states(fixture)[_cache_indices(fixture)],
|
||||||
expected.final_states[_cache_indices(fixture)],
|
expected.final_states[_cache_indices(fixture)],
|
||||||
|
|||||||
@@ -161,6 +161,89 @@ class TestChunkGatedDeltaRule(unittest.TestCase):
|
|||||||
def test_batch_32(self):
|
def test_batch_32(self):
|
||||||
self._check_shape(B=32, T_per_seq=32, H=16, K=128, V=128, pool_size=256)
|
self._check_shape(B=32, T_per_seq=32, H=16, K=128, V=128, pool_size=256)
|
||||||
|
|
||||||
|
@unittest.skipUnless(
|
||||||
|
torch.cuda.is_available(),
|
||||||
|
"The read-only initial-state path is not implemented by the XPU chunk kernel",
|
||||||
|
)
|
||||||
|
def test_read_only_initial_state_supports_duplicate_indices(self):
|
||||||
|
"""MIS item branches may share one query-end state without updating it."""
|
||||||
|
device = get_device()
|
||||||
|
dtype = torch.bfloat16
|
||||||
|
batch_size, tokens_per_item = 3, 65
|
||||||
|
num_heads, key_dim, value_dim = 4, 32, 32
|
||||||
|
total_tokens = batch_size * tokens_per_item
|
||||||
|
|
||||||
|
torch.manual_seed(1234)
|
||||||
|
pool_init = torch.randn(
|
||||||
|
4,
|
||||||
|
num_heads,
|
||||||
|
value_dim,
|
||||||
|
key_dim,
|
||||||
|
dtype=torch.float32,
|
||||||
|
device=device,
|
||||||
|
)
|
||||||
|
duplicate_indices = torch.tensor([2, 2, 2], dtype=torch.int32, device=device)
|
||||||
|
cu_seqlens = torch.arange(
|
||||||
|
0,
|
||||||
|
total_tokens + 1,
|
||||||
|
tokens_per_item,
|
||||||
|
dtype=torch.int32,
|
||||||
|
device=device,
|
||||||
|
)
|
||||||
|
q = torch.randn(1, total_tokens, num_heads, key_dim, dtype=dtype, device=device)
|
||||||
|
k = torch.randn_like(q)
|
||||||
|
v = torch.randn(
|
||||||
|
1, total_tokens, num_heads, value_dim, dtype=dtype, device=device
|
||||||
|
)
|
||||||
|
g = torch.nn.functional.logsigmoid(
|
||||||
|
torch.randn(1, total_tokens, num_heads, dtype=dtype, device=device)
|
||||||
|
)
|
||||||
|
beta = torch.sigmoid(
|
||||||
|
torch.randn(1, total_tokens, num_heads, dtype=dtype, device=device)
|
||||||
|
)
|
||||||
|
|
||||||
|
expected, _ = self._run_reference(
|
||||||
|
pool_init, duplicate_indices, q, k, v, g, beta
|
||||||
|
)
|
||||||
|
actual_pool = pool_init.clone()
|
||||||
|
actual, _, _ = chunk_gated_delta_rule(
|
||||||
|
q=q,
|
||||||
|
k=k,
|
||||||
|
v=v,
|
||||||
|
g=g,
|
||||||
|
beta=beta,
|
||||||
|
initial_state=actual_pool,
|
||||||
|
initial_state_indices=duplicate_indices,
|
||||||
|
cu_seqlens=cu_seqlens,
|
||||||
|
head_first=False,
|
||||||
|
use_qk_l2norm_in_kernel=True,
|
||||||
|
inplace_update=False,
|
||||||
|
)
|
||||||
|
|
||||||
|
torch.testing.assert_close(
|
||||||
|
actual.float(), expected.float(), atol=self.ATOL, rtol=self.RTOL
|
||||||
|
)
|
||||||
|
self.assertTrue(torch.equal(actual_pool, pool_init))
|
||||||
|
|
||||||
|
updating_pool = pool_init.clone()
|
||||||
|
chunk_gated_delta_rule(
|
||||||
|
q=q,
|
||||||
|
k=k,
|
||||||
|
v=v,
|
||||||
|
g=g,
|
||||||
|
beta=beta,
|
||||||
|
initial_state=updating_pool,
|
||||||
|
initial_state_indices=torch.arange(
|
||||||
|
batch_size, dtype=torch.int32, device=device
|
||||||
|
),
|
||||||
|
cu_seqlens=cu_seqlens,
|
||||||
|
head_first=False,
|
||||||
|
use_qk_l2norm_in_kernel=True,
|
||||||
|
)
|
||||||
|
self.assertFalse(
|
||||||
|
torch.equal(updating_pool[:batch_size], pool_init[:batch_size])
|
||||||
|
)
|
||||||
|
|
||||||
# ------------------------------------------------------------------
|
# ------------------------------------------------------------------
|
||||||
# Head count sweep
|
# Head count sweep
|
||||||
# ------------------------------------------------------------------
|
# ------------------------------------------------------------------
|
||||||
|
|||||||
@@ -216,6 +216,36 @@ class TestTritonGDNBackendCorrectness(CustomTestCase):
|
|||||||
with self.subTest(case=case.name, backend=case.backend):
|
with self.subTest(case=case.name, backend=case.backend):
|
||||||
run_gdn_attention_case(self, case)
|
run_gdn_attention_case(self, case)
|
||||||
|
|
||||||
|
def test_multi_item_scoring_mixed_batch_with_empty_query(self):
|
||||||
|
case = GDNAttentionCase(
|
||||||
|
name="gdn_mis_mixed_batch_empty_query",
|
||||||
|
backend="triton",
|
||||||
|
forward_mode=ForwardMode.EXTEND,
|
||||||
|
num_k_heads=2,
|
||||||
|
num_v_heads=2,
|
||||||
|
page_size=1,
|
||||||
|
prefix_lens=(0, 0),
|
||||||
|
extend_lens=(9, 7),
|
||||||
|
mis_delimiter_indices=((0, 3, 8), (4, 6)),
|
||||||
|
conv_history_weight=0.25,
|
||||||
|
)
|
||||||
|
run_gdn_attention_case(self, case)
|
||||||
|
|
||||||
|
def test_multi_item_scoring_crosses_chunk_boundaries(self):
|
||||||
|
case = GDNAttentionCase(
|
||||||
|
name="gdn_mis_chunk_boundaries",
|
||||||
|
backend="triton",
|
||||||
|
forward_mode=ForwardMode.EXTEND,
|
||||||
|
num_k_heads=2,
|
||||||
|
num_v_heads=2,
|
||||||
|
page_size=1,
|
||||||
|
prefix_lens=(0,),
|
||||||
|
extend_lens=(198,),
|
||||||
|
mis_delimiter_indices=((5, 68, 132, 197),),
|
||||||
|
conv_history_weight=0.25,
|
||||||
|
)
|
||||||
|
run_gdn_attention_case(self, case, max_context_len=256)
|
||||||
|
|
||||||
# Layout-robustness. See dense/test_triton.py for the rationale.
|
# Layout-robustness. See dense/test_triton.py for the rationale.
|
||||||
# shuffled_pages is the default for all tests; this method opts
|
# shuffled_pages is the default for all tests; this method opts
|
||||||
# into the more aggressive interleaved_pages + non_monotonic_extend.
|
# into the more aggressive interleaved_pages + non_monotonic_extend.
|
||||||
|
|||||||
@@ -0,0 +1,133 @@
|
|||||||
|
"""End-to-end MIS coverage for the hybrid Qwen3.5 GDN architecture."""
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import os
|
||||||
|
import unittest
|
||||||
|
|
||||||
|
import torch
|
||||||
|
from transformers import AutoTokenizer
|
||||||
|
|
||||||
|
from sglang.srt.entrypoints.engine import Engine
|
||||||
|
from sglang.test.ci.ci_register import register_cuda_ci
|
||||||
|
from sglang.test.test_utils import (
|
||||||
|
DEFAULT_HYBRID_GDN_SMALL_MODEL_NAME_FOR_TEST,
|
||||||
|
CustomTestCase,
|
||||||
|
)
|
||||||
|
|
||||||
|
register_cuda_ci(est_time=240, stage="extra-a", runner_config="1-gpu-large")
|
||||||
|
|
||||||
|
|
||||||
|
class TestQwen35GDNMultiItemScoring(CustomTestCase):
|
||||||
|
model = os.environ.get(
|
||||||
|
"QWEN35_GDN_TEST_MODEL", DEFAULT_HYBRID_GDN_SMALL_MODEL_NAME_FOR_TEST
|
||||||
|
)
|
||||||
|
atol = 2e-2
|
||||||
|
rtol = 2e-2
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def setUpClass(cls):
|
||||||
|
cls.engine = Engine(
|
||||||
|
model_path=cls.model,
|
||||||
|
trust_remote_code=True,
|
||||||
|
dtype="bfloat16",
|
||||||
|
enable_mis=True,
|
||||||
|
attention_backend="flashinfer",
|
||||||
|
linear_attn_prefill_backend="triton",
|
||||||
|
disable_radix_cache=True,
|
||||||
|
chunked_prefill_size=-1,
|
||||||
|
mem_fraction_static=0.8,
|
||||||
|
log_level="error",
|
||||||
|
)
|
||||||
|
tokenizer = AutoTokenizer.from_pretrained(cls.model, trust_remote_code=True)
|
||||||
|
cls.label_token_ids = [
|
||||||
|
tokenizer.encode(label, add_special_tokens=False)[0]
|
||||||
|
for label in (" yes", " no")
|
||||||
|
]
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def tearDownClass(cls):
|
||||||
|
if getattr(cls, "engine", None) is not None:
|
||||||
|
cls.engine.shutdown()
|
||||||
|
torch.cuda.empty_cache()
|
||||||
|
|
||||||
|
def _score(self, query, items):
|
||||||
|
return self.engine.score(
|
||||||
|
query=query,
|
||||||
|
items=items,
|
||||||
|
label_token_ids=self.label_token_ids,
|
||||||
|
apply_softmax=False,
|
||||||
|
).scores
|
||||||
|
|
||||||
|
async def _async_score(self, query, items):
|
||||||
|
result = await self.engine.async_score(
|
||||||
|
query=query,
|
||||||
|
items=items,
|
||||||
|
label_token_ids=self.label_token_ids,
|
||||||
|
apply_softmax=False,
|
||||||
|
)
|
||||||
|
return result.scores
|
||||||
|
|
||||||
|
def _pointwise(self, query, items):
|
||||||
|
return [self._score(query, [item])[0] for item in items]
|
||||||
|
|
||||||
|
def _assert_scores_close(self, actual, expected):
|
||||||
|
torch.testing.assert_close(
|
||||||
|
torch.tensor(actual),
|
||||||
|
torch.tensor(expected),
|
||||||
|
atol=self.atol,
|
||||||
|
rtol=self.rtol,
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_batched_matches_pointwise_for_varied_requests(self):
|
||||||
|
cases = [
|
||||||
|
(
|
||||||
|
"Decide whether each statement is true:",
|
||||||
|
["The sky is blue.", "Two plus two is five.", "Water freezes."],
|
||||||
|
),
|
||||||
|
("", ["empty query short", "empty query with a much longer item " * 8]),
|
||||||
|
("Classify:", ["one"]),
|
||||||
|
(
|
||||||
|
"Judge each passage:",
|
||||||
|
["tiny", "medium length passage " * 5, "long passage " * 24],
|
||||||
|
),
|
||||||
|
]
|
||||||
|
|
||||||
|
for query, items in cases:
|
||||||
|
with self.subTest(query=query, item_count=len(items)):
|
||||||
|
self._assert_scores_close(
|
||||||
|
self._score(query, items), self._pointwise(query, items)
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_sibling_changes_and_reordering_do_not_change_target(self):
|
||||||
|
query = "Rate each answer:"
|
||||||
|
target = "The target answer remains unchanged."
|
||||||
|
baseline = self._score(query, [target, "sibling A", "sibling B"])[0]
|
||||||
|
changed = self._score(
|
||||||
|
query, [target, "a completely different sibling " * 8, "sibling B"]
|
||||||
|
)[0]
|
||||||
|
reordered = self._score(query, ["sibling B", "sibling A", target])[2]
|
||||||
|
|
||||||
|
self._assert_scores_close(changed, baseline)
|
||||||
|
self._assert_scores_close(reordered, baseline)
|
||||||
|
|
||||||
|
def test_concurrent_requests_match_pointwise(self):
|
||||||
|
cases = [
|
||||||
|
("Is it an animal?", ["cat", "table", "blue whale"]),
|
||||||
|
("", ["alpha", "beta"]),
|
||||||
|
("Choose:", ["first"]),
|
||||||
|
("Is it a city?", ["Paris", "bread", "Tokyo", "chair"]),
|
||||||
|
]
|
||||||
|
expected = [self._pointwise(query, items) for query, items in cases]
|
||||||
|
|
||||||
|
async def gather_scores():
|
||||||
|
return await asyncio.gather(
|
||||||
|
*(self._async_score(query, items) for query, items in cases)
|
||||||
|
)
|
||||||
|
|
||||||
|
actual = self.engine.loop.run_until_complete(gather_scores())
|
||||||
|
for actual_scores, expected_scores in zip(actual, expected):
|
||||||
|
self._assert_scores_close(actual_scores, expected_scores)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
unittest.main()
|
||||||
@@ -119,6 +119,63 @@ def causal_conv1d_update_ref(
|
|||||||
return (out if activation is None else F.silu(out)).to(dtype=dtype_in)
|
return (out if activation is None else F.silu(out)).to(dtype=dtype_in)
|
||||||
|
|
||||||
|
|
||||||
|
def test_causal_conv1d_branches_from_cloned_query_state():
|
||||||
|
device = get_device()
|
||||||
|
dtype = torch.bfloat16
|
||||||
|
torch.manual_seed(7)
|
||||||
|
dim, width = 96, 4
|
||||||
|
item_lens = [2, 5, 7]
|
||||||
|
total_tokens = sum(item_lens)
|
||||||
|
|
||||||
|
source_state = torch.randn(1, dim, width - 1, device=device, dtype=dtype)
|
||||||
|
source_state_before = source_state.clone()
|
||||||
|
branch_states = source_state.expand(len(item_lens), -1, -1).contiguous().clone()
|
||||||
|
x = torch.randn(dim, total_tokens, device=device, dtype=dtype)
|
||||||
|
weight = torch.randn(dim, width, device=device, dtype=dtype)
|
||||||
|
bias = torch.randn(dim, device=device, dtype=dtype)
|
||||||
|
query_start_loc = torch.tensor(
|
||||||
|
[0, *torch.tensor(item_lens).cumsum(0).tolist()],
|
||||||
|
dtype=torch.int32,
|
||||||
|
device=device,
|
||||||
|
)
|
||||||
|
cache_indices = torch.arange(len(item_lens), dtype=torch.int32, device=device)
|
||||||
|
|
||||||
|
actual = causal_conv1d_fn(
|
||||||
|
x,
|
||||||
|
weight,
|
||||||
|
bias=bias,
|
||||||
|
conv_states=branch_states,
|
||||||
|
query_start_loc=query_start_loc,
|
||||||
|
seq_lens_cpu=torch.tensor(item_lens),
|
||||||
|
cache_indices=cache_indices,
|
||||||
|
has_initial_state=torch.ones(len(item_lens), dtype=torch.bool, device=device),
|
||||||
|
activation="silu",
|
||||||
|
pad_slot_id=PAD_SLOT_ID,
|
||||||
|
)
|
||||||
|
|
||||||
|
expected_outputs = []
|
||||||
|
expected_states = []
|
||||||
|
for item in torch.split(x, item_lens, dim=-1):
|
||||||
|
item_output, item_state = causal_conv1d_ref(
|
||||||
|
item.unsqueeze(0),
|
||||||
|
weight,
|
||||||
|
bias,
|
||||||
|
initial_states=source_state,
|
||||||
|
return_final_states=True,
|
||||||
|
activation="silu",
|
||||||
|
)
|
||||||
|
expected_outputs.append(item_output.squeeze(0))
|
||||||
|
expected_states.append(item_state.squeeze(0))
|
||||||
|
|
||||||
|
torch.testing.assert_close(
|
||||||
|
actual, torch.cat(expected_outputs, dim=-1), atol=5e-2, rtol=1e-2
|
||||||
|
)
|
||||||
|
torch.testing.assert_close(
|
||||||
|
branch_states, torch.stack(expected_states), atol=5e-2, rtol=1e-2
|
||||||
|
)
|
||||||
|
assert torch.equal(source_state, source_state_before)
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.parametrize("itype", [torch.bfloat16, torch.float])
|
@pytest.mark.parametrize("itype", [torch.bfloat16, torch.float])
|
||||||
@pytest.mark.parametrize("silu_activation", [True])
|
@pytest.mark.parametrize("silu_activation", [True])
|
||||||
@pytest.mark.parametrize("has_bias", [True])
|
@pytest.mark.parametrize("has_bias", [True])
|
||||||
|
|||||||
@@ -13,6 +13,7 @@ from sglang.srt.layers.attention.linear.gdn_backend import (
|
|||||||
GDNKernelDispatcher,
|
GDNKernelDispatcher,
|
||||||
_validate_gdn_linear_attn_backends,
|
_validate_gdn_linear_attn_backends,
|
||||||
flashinfer_gdn_prefill_default,
|
flashinfer_gdn_prefill_default,
|
||||||
|
validate_gdn_mis_backend,
|
||||||
)
|
)
|
||||||
from sglang.srt.layers.attention.linear.kernels.gdn_flashinfer import (
|
from sglang.srt.layers.attention.linear.kernels.gdn_flashinfer import (
|
||||||
maybe_build_flashinfer_checkpoint_plan,
|
maybe_build_flashinfer_checkpoint_plan,
|
||||||
@@ -59,6 +60,7 @@ def make_runner(
|
|||||||
mamba_radix_cache_strategy="no_buffer",
|
mamba_radix_cache_strategy="no_buffer",
|
||||||
enable_dynamic_chunking=False,
|
enable_dynamic_chunking=False,
|
||||||
chunked_prefill_size=8192,
|
chunked_prefill_size=8192,
|
||||||
|
enable_mis=False,
|
||||||
)
|
)
|
||||||
fields.update(arg_overrides)
|
fields.update(arg_overrides)
|
||||||
args = _publish(testcase, **fields)
|
args = _publish(testcase, **fields)
|
||||||
@@ -79,6 +81,22 @@ def make_runner(
|
|||||||
|
|
||||||
|
|
||||||
class TestFlashInferGDNPrefillBackendPolicy(CustomTestCase):
|
class TestFlashInferGDNPrefillBackendPolicy(CustomTestCase):
|
||||||
|
def test_mis_requires_triton_prefill_backend(self):
|
||||||
|
runner = make_runner(self, enable_mis=True)
|
||||||
|
with self.assertRaisesRegex(ValueError, "Triton linear-attention prefill"):
|
||||||
|
validate_gdn_mis_backend(LinearAttnKernelBackend.FLASHINFER)
|
||||||
|
|
||||||
|
def test_mis_rejects_page_major_layout(self):
|
||||||
|
make_runner(self, enable_mis=True, enable_page_major_kv_layout=True)
|
||||||
|
with self.assertRaisesRegex(ValueError, "page-major"):
|
||||||
|
validate_gdn_mis_backend(LinearAttnKernelBackend.TRITON)
|
||||||
|
|
||||||
|
def test_non_gdn_linear_backend_rejects_mis(self):
|
||||||
|
with self.assertRaisesRegex(ValueError, "does not support multi-item scoring"):
|
||||||
|
MambaAttnBackendBase.validate_mis_support(SimpleNamespace(enable_mis=True))
|
||||||
|
|
||||||
|
GDNAttnBackend.validate_mis_support(SimpleNamespace(enable_mis=True))
|
||||||
|
|
||||||
def apply_policy(
|
def apply_policy(
|
||||||
self,
|
self,
|
||||||
runner,
|
runner,
|
||||||
@@ -232,6 +250,7 @@ class TestFlashInferGDNPrefillBackendPolicy(CustomTestCase):
|
|||||||
backend.kernel_dispatcher = SimpleNamespace(extend_uses_state_checkpoints=True)
|
backend.kernel_dispatcher = SimpleNamespace(extend_uses_state_checkpoints=True)
|
||||||
metadata = SimpleNamespace(has_mamba_track_mask=True, track_ssm_h_src=None)
|
metadata = SimpleNamespace(has_mamba_track_mask=True, track_ssm_h_src=None)
|
||||||
forward_batch = SimpleNamespace(
|
forward_batch = SimpleNamespace(
|
||||||
|
multi_item_delimiter_indices=None,
|
||||||
mamba_track_mask=torch.tensor([True]),
|
mamba_track_mask=torch.tensor([True]),
|
||||||
mamba_track_indices=torch.tensor([7]),
|
mamba_track_indices=torch.tensor([7]),
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -0,0 +1,99 @@
|
|||||||
|
import unittest
|
||||||
|
from types import SimpleNamespace
|
||||||
|
|
||||||
|
import torch
|
||||||
|
|
||||||
|
from sglang.srt.layers.attention.linear.gdn_backend import build_gdn_mis_metadata
|
||||||
|
from sglang.test.ci.ci_register import register_cpu_ci
|
||||||
|
from sglang.test.test_utils import CustomTestCase
|
||||||
|
|
||||||
|
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
|
||||||
|
|
||||||
|
|
||||||
|
class TestGDNMISMetadata(CustomTestCase):
|
||||||
|
def test_allows_trailing_attention_padding(self):
|
||||||
|
forward_batch = SimpleNamespace(
|
||||||
|
input_ids=torch.empty(8, dtype=torch.int64),
|
||||||
|
extend_seq_lens_cpu=[5],
|
||||||
|
extend_prefix_lens_cpu=[0],
|
||||||
|
multi_item_delimiter_indices=[torch.tensor([1, 4], dtype=torch.int64)],
|
||||||
|
is_prefill_only=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
metadata = build_gdn_mis_metadata(forward_batch)
|
||||||
|
|
||||||
|
torch.testing.assert_close(
|
||||||
|
torch.cat([metadata.query_token_indices, metadata.item_token_indices])
|
||||||
|
.sort()
|
||||||
|
.values,
|
||||||
|
torch.arange(5, dtype=torch.int64),
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_mixed_batch_with_empty_query(self):
|
||||||
|
forward_batch = SimpleNamespace(
|
||||||
|
input_ids=torch.empty(16, dtype=torch.int64),
|
||||||
|
extend_seq_lens_cpu=[9, 7],
|
||||||
|
extend_prefix_lens_cpu=[0, 0],
|
||||||
|
multi_item_delimiter_indices=[
|
||||||
|
torch.tensor([0, 3, 8], dtype=torch.int64),
|
||||||
|
torch.tensor([4, 6], dtype=torch.int64),
|
||||||
|
],
|
||||||
|
is_prefill_only=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
metadata = build_gdn_mis_metadata(forward_batch)
|
||||||
|
|
||||||
|
self.assertEqual(metadata.query_seq_lens_cpu, [4])
|
||||||
|
self.assertEqual(metadata.item_seq_lens_cpu, [3, 5, 1, 2, 1])
|
||||||
|
torch.testing.assert_close(
|
||||||
|
metadata.query_token_indices,
|
||||||
|
torch.tensor([9, 10, 11, 12], dtype=torch.int64),
|
||||||
|
)
|
||||||
|
torch.testing.assert_close(
|
||||||
|
metadata.query_cu_seqlens,
|
||||||
|
torch.tensor([0, 4], dtype=torch.int32),
|
||||||
|
)
|
||||||
|
torch.testing.assert_close(
|
||||||
|
metadata.query_request_indices,
|
||||||
|
torch.tensor([1], dtype=torch.int64),
|
||||||
|
)
|
||||||
|
torch.testing.assert_close(
|
||||||
|
metadata.item_token_indices,
|
||||||
|
torch.tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 13, 14, 15], dtype=torch.int64),
|
||||||
|
)
|
||||||
|
torch.testing.assert_close(
|
||||||
|
metadata.item_cu_seqlens,
|
||||||
|
torch.tensor([0, 3, 8, 9, 11, 12], dtype=torch.int32),
|
||||||
|
)
|
||||||
|
torch.testing.assert_close(
|
||||||
|
metadata.item_request_indices,
|
||||||
|
torch.tensor([0, 0, 0, 1, 1], dtype=torch.int64),
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_rejects_sequence_lengths_beyond_input(self):
|
||||||
|
forward_batch = SimpleNamespace(
|
||||||
|
input_ids=torch.empty(4, dtype=torch.int64),
|
||||||
|
extend_seq_lens_cpu=[5],
|
||||||
|
extend_prefix_lens_cpu=[0],
|
||||||
|
multi_item_delimiter_indices=[torch.tensor([1, 4], dtype=torch.int64)],
|
||||||
|
is_prefill_only=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
with self.assertRaisesRegex(ValueError, "exceed the input tokens"):
|
||||||
|
build_gdn_mis_metadata(forward_batch)
|
||||||
|
|
||||||
|
def test_rejects_missing_final_delimiter(self):
|
||||||
|
forward_batch = SimpleNamespace(
|
||||||
|
input_ids=torch.empty(5, dtype=torch.int64),
|
||||||
|
extend_seq_lens_cpu=[5],
|
||||||
|
extend_prefix_lens_cpu=[0],
|
||||||
|
multi_item_delimiter_indices=[torch.tensor([2, 3], dtype=torch.int64)],
|
||||||
|
is_prefill_only=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
with self.assertRaisesRegex(ValueError, "final delimiter"):
|
||||||
|
build_gdn_mis_metadata(forward_batch)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
unittest.main()
|
||||||
Reference in New Issue
Block a user