diff --git a/python/sglang/srt/models/inkling_common/attn.py b/python/sglang/srt/models/inkling_common/attn.py index 895258811..75e6aa701 100644 --- a/python/sglang/srt/models/inkling_common/attn.py +++ b/python/sglang/srt/models/inkling_common/attn.py @@ -6,6 +6,9 @@ from functools import cache import torch from torch import nn +from sglang.kernels.ops.attention.flash_attn.cute.batch_invariance import ( + is_batch_invariant, +) from sglang.kernels.ops.attention.inkling_rel_proj import rel_proj_small_t from sglang.kernels.ops.attention.inkling_row_scale import row_compact_bf16 from sglang.kernels.ops.attention.log_scaling_tau import ( @@ -943,7 +946,15 @@ class InklingAttention(nn.Module): # stay bf16 here and the MXFP8 pool's set_kv_buffer quantizes and # stores them in one fused kernel (absent descales signal it). - if envs.SGLANG_OPT_USE_INKLING_SHEARED_BIAS.get() and fa4: + # The sheared-bias kernel is not batch invariant: its bias tile geometry + # follows the query count, so the same absolute (q, k) pair accumulates in + # a different order for a few-query decode step than for a many-query + # prefill. Deterministic mode takes the score_mod path instead. + if ( + envs.SGLANG_OPT_USE_INKLING_SHEARED_BIAS.get() + and fa4 + and not is_batch_invariant() + ): # FA4 sheared-bias kernel: pass rel_logits directly; the kernel shears # it into a column-aligned pre-softmax bias. attn_output = self.attn( diff --git a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_hybrid_bitexact.py b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_hybrid_bitexact.py index 7d2919926..ce8a88380 100644 --- a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_hybrid_bitexact.py +++ b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_hybrid_bitexact.py @@ -307,21 +307,14 @@ class TestUnifiedHybridMTPBitExact(CustomTestCase): from the verify path. #29792 was a wrong-slot pick in the non-spec save; nothing exercises the spec-side save today. - Two settings are load-bearing rather than incidental: + `--speculative-num-steps 2` is load-bearing rather than incidental: it + matches the two MTP heads this checkpoint ships, and a third step has no + weights and the draft head refuses to start. - `--speculative-num-steps 2` matches the two MTP heads this checkpoint - ships; a third step has no weights and the draft head refuses to start. - - `SGLANG_OPT_USE_INKLING_SHEARED_BIAS=0` is required for the exact bar. The - sheared relative-bias path shears on `max_seqlen_q`, so a verify pass - (several queries) lands the same absolute (q, k) pair on a different tile - than a prefill pass and one output element differs. Measured on this - checkpoint: with the sheared path on, 12 of 16 tokens diverge from the - first token past a tile boundary, up to 8.1e-03; with it off, every token - reads exactly 0. Tracked in - https://github.com/sgl-project/sglang/issues/34899 -- when the shear is - made query-count invariant this override should be dropped, and this class - is what will prove it. + Reaching the exact bar also needs the sheared relative-bias path out of the + way, since its bias tile geometry follows the query count. Deterministic + mode now selects the invariant path on its own, so this class carries no + environment override; a regression there surfaces here as a nonzero KL. """ @classmethod @@ -351,7 +344,6 @@ class TestUnifiedHybridMTPBitExact(CustomTestCase): env={ **os.environ, "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1", - "SGLANG_OPT_USE_INKLING_SHEARED_BIAS": "0", }, )