From 0ee3749e07a857daca8313d3f7783aea9de03257 Mon Sep 17 00:00:00 2001 From: Khoa Pham Date: Fri, 21 Aug 2026 18:13:18 -0700 Subject: [PATCH] [Fix] Read the granite sinks dtype from the exec bag, not the legacy global shim (#35921) Co-authored-by: Claude Opus 5 (1M context) --- python/sglang/srt/models/granite.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/python/sglang/srt/models/granite.py b/python/sglang/srt/models/granite.py index 1fcb41acd..45a3f239e 100644 --- a/python/sglang/srt/models/granite.py +++ b/python/sglang/srt/models/granite.py @@ -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)})