Skip inkling sheared bias under batch invariance (#35161)

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(

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
+    Reaching the exact bar also needs the sheared relative-bias path out of the
-    ships; a third step has no weights and the draft head refuses to start.
+    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
-    `SGLANG_OPT_USE_INKLING_SHEARED_BIAS=0` is required for the exact bar. The
+    environment override; a regression there surfaces here as a nonzero KL.
-    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.
     """
 
     @classmethod
@@ -351,7 +344,6 @@ class TestUnifiedHybridMTPBitExact(CustomTestCase):
             env={
                 **os.environ,
                 "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1",
-                "SGLANG_OPT_USE_INKLING_SHEARED_BIAS": "0",
             },
         )