Additional k3 sync fixes

python/sglang/kernels/jit/csrc/elementwise/set_mla_kv_concat_q.cuh

@@ -342,10 +342,10 @@ struct SetMlaKVConcatQFp8Params {
   int32_t qo_stride_1;
 };
 
-// 2x bf16 -> 2x fp8 e4m3, float-mediated cvt.rn NOSAT (matches aten: overflow -> NaN).
+// 2x bf16 -> 2x fp8 e4m3, float-mediated cvt.rn.satfinite.
 SGL_DEVICE uint16_t bf16x2_to_fp8x2(const bf16x2_t v) {
   const float2 f = __bfloat1622float2(v);
-  return __nv_cvt_float2_to_fp8x2(f, __NV_NOSAT, __NV_E4M3);
+  return __nv_cvt_float2_to_fp8x2(f, __NV_SATFINITE, __NV_E4M3);
 }
 
 // Convert 8 bf16 (one int4 load) to 8 fp8 packed in a uint2.

python/sglang/kernels/ops/attention/utils.py

@@ -54,6 +54,7 @@
 from sglang.kernels.ops.kvcache.rope_cache import (
     fused_qk_rope_reshape_and_cache as fused_qk_rope_reshape_and_cache,
 )
+from sglang.kernels.ops.quantization.fp8_utils import to_fp8_satfinite
 from sglang.srt.utils import is_cuda
 
 _is_cuda = is_cuda()
@@ -212,8 +213,12 @@ def mla_quantize_without_rope_for_fp8(
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     """Quantize MLA components to FP8 without applying rotary embeddings."""
     attn_dtype = torch.float8_e4m3fn
-    q = concat_mla_absorb_q_general(q_nope, q_rope).to(attn_dtype)
-    return q, k_nope.to(attn_dtype), k_rope.to(attn_dtype)
+    q = to_fp8_satfinite(concat_mla_absorb_q_general(q_nope, q_rope), attn_dtype)
+    return (
+        q,
+        to_fp8_satfinite(k_nope, attn_dtype),
+        to_fp8_satfinite(k_rope, attn_dtype),
+    )
 
 
 def concat_mla_absorb_q_general(q_nope, q_rope):

python/sglang/kernels/ops/kimi_k3/kda_decode_mtp.py

@@ -277,7 +277,10 @@ def kda_decode_mtp_kernel(
     # CUDA-graph padding rows use slot == -1.
     if slot < 0:
         cute.arch.griddepcontrol_wait()
-        pad_bos = cu_seqlens[i_n]
+        # This wrapper requires a fixed dense verify width. Graph metadata
+        # collapses all padded sequence boundaries to the real-token endpoint,
+        # but each padded request still owns its own dense output interval.
+        pad_bos = i_n * T_LOOP
         for i_t in cutlass.range_constexpr(T_LOOP):
             if tidx < HEAD_DIM:
                 o[0, pad_bos + i_t, i_hv, tidx] = cutlass.BFloat16(0.0)

python/sglang/kernels/ops/quantization/fp8_utils.py

@@ -12,6 +12,16 @@
 )
 
 
+def to_fp8_satfinite(x: torch.Tensor, fp8_dtype: torch.dtype) -> torch.Tensor:
+    """Cast to an fp8 dtype with satfinite semantics: finite values beyond the
+    format's range clamp to +-finfo.max (448 for e4m3fn); NaN stays NaN.
+    """
+    if x.dtype == fp8_dtype:
+        return x
+    fp8_max = torch.finfo(fp8_dtype).max
+    return x.clamp(-fp8_max, fp8_max).to(fp8_dtype)
+
+
 def cuda_capability_uses_fp8_e4b15(cuda_capability: Tuple[int, int]) -> bool:
     """Triton names E4M3 as fp8e4b15 on CUDA architectures before SM89."""
     return cuda_capability < (8, 9)

python/sglang/kernels/ops/speculative/dflash.py

@@ -271,8 +271,11 @@ def _selector_walk_kernel(
             tl.float32
         )
         if greedy:
-            best = tl.max(scores, axis=0)
-            index = tl.min(tl.where(scores == best, offsets, top_k), axis=0)
+            # First candidate that no other candidate strictly beats. Total over
+            # any row (there is always at least one), so index < top_k and the
+            # row never resolves to a neighbour's candidates.
+            beaten = tl.sum((scores[None, :] > scores[:, None]).to(tl.int32), axis=1)
+            index = tl.min(tl.where(beaten == 0, offsets, top_k), axis=0)
             probabilities = tl.where(offsets == index, 1.0, 0.0)
         else:
             scaled = scores / temperature

python/sglang/srt/entrypoints/warmup.py

@@ -109,10 +109,13 @@ async def whisper_autodetect(
 async def voice_chat(disaggregation_mode: str, tokenizer_manager: TokenizerManager):
     # this warms up the fused_moe triton kernels and caches them
     # if we don't do this we break real time inference for voice chat
+    token_upper_bound = min(
+        2**16, tokenizer_manager.model_config.hf_text_config.vocab_size
+    )
     for i in tqdm.trange(1, 512):
         size = i * 4
         generate_req_input = GenerateReqInput(
-            input_ids=(np.random.randint(2**16, size=[size])).tolist(),
+            input_ids=(np.random.randint(token_upper_bound, size=[size])).tolist(),
             sampling_params={
                 "max_new_tokens": 30,
                 "temperature": 0.8,
@@ -134,6 +137,9 @@ async def prefill_shapes(disaggregation_mode: str, tokenizer_manager: TokenizerM
     Uses power-of-2 sizes plus intermediate points to cover the shape space
     that fused_moe, attention extend, and other Triton kernels may encounter.
     """
+    token_upper_bound = min(
+        2**16, tokenizer_manager.model_config.hf_text_config.vocab_size
+    )
     page_size = 64
     sizes = set()
     base = 64
@@ -148,7 +154,7 @@ async def prefill_shapes(disaggregation_mode: str, tokenizer_manager: TokenizerM
 
     for size in tqdm.tqdm(sizes, desc="Warmup prefill shapes (up to 32K)"):
         generate_req_input = GenerateReqInput(
-            input_ids=(np.random.randint(2**16, size=[size])).tolist(),
+            input_ids=(np.random.randint(token_upper_bound, size=[size])).tolist(),
             sampling_params={
                 "max_new_tokens": 1,
                 "temperature": 0.0,

python/sglang/srt/environ.py

@@ -1338,6 +1338,8 @@ class Envs:
     SGLANG_RESIZE_RESAMPLE = EnvStr("")
     SGLANG_MM_BUFFER_SIZE_MB = EnvInt(0)
     SGLANG_MM_PRECOMPUTE_HASH = EnvBool(False)
+    # Omit a redundant processor ID list after copying canonical request IDs.
+    SGLANG_MM_STRIP_PROCESSOR_INPUT_IDS = EnvBool(False)
     SGLANG_VIT_ENABLE_CUDA_GRAPH = EnvBool(False)
     # Use the fully-vectorized ViT position-embedding interpolation (no per-image
     # Python loop / CPU<->GPU sync). Bit-exact with the legacy implementation;
@@ -1349,6 +1351,9 @@ class Envs:
     # For pre-tokenized (list[int]) multimodal prompts,
     # preserve the user's original tokens to avoid retokenization drift.
     SGLANG_MM_AVOID_RETOKENIZE = EnvBool(True)
+    # Raise instead of cropping when final multimodal embedding rows do not
+    # match the placeholder tokens in the extend window.
+    SGLANG_ENABLE_STRICT_MM_EMBEDDING_LENGTH = EnvBool(False)
 
     # ===================================================================
     # Multimodal CUDA IPC transport

python/sglang/srt/layers/attention/trtllm_mla_backend.py

@@ -45,6 +45,7 @@
     get_num_page_per_block_flashmla,
 )
 from sglang.kernels.ops.quantization.fp8_kernel import scaled_fp8_quant
+from sglang.kernels.ops.quantization.fp8_utils import to_fp8_satfinite
 from sglang.srt.environ import envs
 from sglang.srt.layers.attention.flashinfer_mla_backend import (
     FlashInferMLAAttnBackend,
@@ -116,7 +117,7 @@ def grow_multi_ctas_kv_counter_buffer_if_needed(
 
 
 def _quantize_fp8_qkv(q, k, v, layer):
-    q = q.to(torch.float8_e4m3fn)
+    q = to_fp8_satfinite(q, torch.float8_e4m3fn)
 
     k_scale = getattr(layer, "k_scale_float", None)
     if k_scale is None:
@@ -130,7 +131,7 @@ def _quantize_fp8_qkv(q, k, v, layer):
         )
         k = k_2d.reshape(k.shape)
     else:
-        k = k.to(torch.float8_e4m3fn)
+        k = to_fp8_satfinite(k, torch.float8_e4m3fn)
 
     v_scale = getattr(layer, "v_scale_float", None)
     if v_scale is None:
@@ -144,7 +145,7 @@ def _quantize_fp8_qkv(q, k, v, layer):
         )
         v = v_2d.reshape(v.shape)
     else:
-        v = v.to(torch.float8_e4m3fn)
+        v = to_fp8_satfinite(v, torch.float8_e4m3fn)
 
     return q, k, v, k_scale, v_scale
 

python/sglang/srt/layers/quantization/kv_cache.py

@@ -49,6 +49,21 @@ def apply(self, layer: torch.nn.Module) -> torch.Tensor:
         raise RuntimeError(f"{self.__class__.__name__}.apply should not be called.")
 
     def process_weights_after_loading(self, layer) -> None:
+        if layer.k_scale.numel() != 1 or layer.v_scale.numel() != 1:
+            raise ValueError(
+                "Only support per-tensor scaling factor for fp8 KV cache, got "
+                f"k_scale shape={tuple(layer.k_scale.shape)} "
+                f"v_scale shape={tuple(layer.v_scale.shape)}"
+            )
+        k_scale_cpu = layer.k_scale.to("cpu")
+        v_scale_cpu = layer.v_scale.to("cpu")
+        if not (
+            torch.isfinite(k_scale_cpu).all() and torch.isfinite(v_scale_cpu).all()
+        ):
+            raise ValueError(
+                "fp8 KV cache scales must be finite, got "
+                f"k_scale={k_scale_cpu.item()} v_scale={v_scale_cpu.item()}"
+            )
         if layer.k_scale > 0.0 and layer.v_scale > 0.0:
             # We prefer to use separate k_scale and v_scale if present
             k_scale = layer.k_scale.to("cpu").tolist()

python/sglang/srt/managers/mm_schedule.py

@@ -703,11 +703,34 @@ def _count_mm_tokens_in_extend(
     return num_mm_tokens
 
 
+def _check_exact_embedding_length(
+    embedding: torch.Tensor, num_mm_tokens_in_input_ids: int
+) -> torch.Tensor:
+    # Chunk extraction has already selected the rows for this extend window, so
+    # any mismatch means rows would land on the wrong placeholder tokens.
+    num_mm_tokens_in_embedding = _embedding_token_count(embedding)
+    if num_mm_tokens_in_input_ids != num_mm_tokens_in_embedding:
+        hint = ""
+        if get_schedule().chunked_prefill_size != -1:
+            hint = (
+                " Chunked prefill is enabled; check that embedding chunks match "
+                "the requested token spans."
+            )
+        raise RuntimeError(
+            "Multimodal embedding length does not match the placeholder tokens in "
+            f"the input text: {num_mm_tokens_in_input_ids=} vs "
+            f"{num_mm_tokens_in_embedding=}.{hint}"
+        )
+    return embedding
+
+
 def _adjust_embedding_length(
     embedding: torch.Tensor,
     num_mm_tokens_in_input_ids: int,
     logger,
 ) -> torch.Tensor:
+    if envs.SGLANG_ENABLE_STRICT_MM_EMBEDDING_LENGTH.get():
+        return _check_exact_embedding_length(embedding, num_mm_tokens_in_input_ids)
     num_mm_tokens_in_embedding = embedding.shape[0]
     if num_mm_tokens_in_input_ids != num_mm_tokens_in_embedding:
         logger.warning(

python/sglang/srt/managers/schedule_batch.py

@@ -837,6 +837,9 @@ def merge(self, other: MultimodalInputs):
                 self.mrope_position_delta = torch.cat(
                     [self.mrope_position_delta, other.mrope_position_delta], dim=0
                 )
+        elif other.mrope_position_delta is not None:
+            self.mrope_position_delta = other.mrope_position_delta
+        self.mrope_position_delta_repeated_cache = None
 
         for key, val in other.__dict__.items():
             if "_id" in key:
@@ -1014,6 +1017,7 @@ def __init__(
         # Full-KV-derived boundary whose SWA window should be inserted after
         # the current prefill pass.
         self.swa_branching_seqlen: Optional[int] = None
+        self.penalty_cumulated_len = 0
 
         # for cross-encoder model
         self.token_type_ids = token_type_ids
@@ -1918,6 +1922,7 @@ def reset_for_retract(self):
         self.already_computed = 0
         assert not self.kv.holds_kv, "expect it is already released"
         self.kv.kv_committed_len = 0
+        self.penalty_cumulated_len = 0
         self.extend_batch_idx = 0
         self.decode_batch_idx = 0
 
@@ -3421,6 +3426,39 @@ def cumulate_penalty_output_tokens(self):
             latest_output_ids
         )
 
+    def cumulate_penalty_output_tokens_since_last(self):
+        """Feed every token committed since the previous call (all accepted
+        speculative tokens, not only the last) to the penalizers. The first
+        call also feeds the last prompt token, matching the non-spec path."""
+        new_tokens = []
+        for req in self.reqs:
+            seq_len = 1 + len(req.output_ids)
+            start = req.penalty_cumulated_len
+            if start == 0:
+                toks = [req.origin_input_ids[-1]] + list(req.output_ids)
+            else:
+                toks = list(req.output_ids[start - 1 :])
+            new_tokens.append(toks)
+            req.penalty_cumulated_len = seq_len
+
+        k = max((len(t) for t in new_tokens), default=0)
+        if k == 0:
+            return
+
+        pin_memory = is_pin_memory_available(self.device)
+        ids = torch.zeros((len(self.reqs), k), dtype=torch.int64, pin_memory=pin_memory)
+        num_valid = torch.tensor(
+            [len(t) for t in new_tokens], dtype=torch.int64, pin_memory=pin_memory
+        )
+        for i, t in enumerate(new_tokens):
+            if t:
+                ids[i, : len(t)] = torch.tensor(t, dtype=torch.int64)
+
+        self.sampling_info.penalizer_orchestrator.cumulate_output_tokens_multi(
+            ids.to(self.device, non_blocking=True),
+            num_valid.to(self.device, non_blocking=True),
+        )
+
     def prepare_for_decode(self):
         self.forward_mode = ForwardMode.DECODE
         # Decode embeds the last output token via embed_tokens; clear the stale