[Fix] Read the granite sinks dtype from the exec bag, not the legacy global shim (#35921)
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 5
parent
7fd5454335
commit
0ee3749e07
@@ -47,8 +47,7 @@ from sglang.srt.model_loader.weight_utils import (
|
||||
default_weight_loader,
|
||||
sharded_weight_loader,
|
||||
)
|
||||
from sglang.srt.runtime_context import get_parallel
|
||||
from sglang.srt.server_args import get_global_server_args
|
||||
from sglang.srt.runtime_context import get_exec, get_parallel
|
||||
from sglang.srt.utils import add_prefix, set_weight_attrs
|
||||
from sglang.utils import get_exception_traceback
|
||||
|
||||
@@ -74,8 +73,10 @@ def granite_layer_attn_params(
|
||||
|
||||
|
||||
def build_attention_sinks(num_heads: int) -> nn.Parameter:
|
||||
# trtllm_mha requires float32 sinks, other backends use bfloat16.
|
||||
attn_backend = get_global_server_args().attention_backend
|
||||
# TODO(kpham-sgl): one parameter cannot serve a split launch -- trtllm_mha
|
||||
# wants float32 sinks, FA4 wants bfloat16. Pick the dtype at init instead,
|
||||
# once the serving backends are known. Checkpoint dtype also unverified.
|
||||
attn_backend = get_exec().kernel.attention_backend
|
||||
sinks_dtype = torch.float32 if attn_backend == "trtllm_mha" else torch.bfloat16
|
||||
sinks = nn.Parameter(torch.empty(num_heads, dtype=sinks_dtype), requires_grad=False)
|
||||
set_weight_attrs(sinks, {"weight_loader": sharded_weight_loader(0)})
|
||||
|
||||
Reference in New Issue
Block a user