# Copyright (c) 2026 LightSeek Foundation # # Permission is hereby granted, free of charge, to any person obtaining a copy # of this software and associated documentation files (the "Software"), to deal # in the Software without restriction, including without limitation the rights # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell # copies of the Software, and to permit persons to whom the Software is # furnished to do so, subject to the following conditions: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. """Logits processing.""" import dataclasses import torch import triton import triton.language as tl from tokenspeed_kernel.ops.communication.triton import all_gather_inner, create_state from tokenspeed_kernel.ops.sampling import argmax as sampling_argmax from tokenspeed_kernel.ops.sampling.cute_dsl import ( create_dist_argmax_state, distributed_argmax, ) from tokenspeed_kernel.platform import current_platform from torch import nn from tokenspeed.runtime.distributed.comm_ops import all_gather_into_tensor from tokenspeed.runtime.distributed.process_group_manager import ( process_group_manager as pg_manager, ) from tokenspeed.runtime.execution.context import ForwardContext from tokenspeed.runtime.execution.forward_batch_info import ( CaptureHiddenMode, ForwardMode, ) from tokenspeed.runtime.layers.vocab_parallel_embedding import ( VocabParallelEmbedding, ) from tokenspeed.runtime.sampling.dp_sampling_config import ( DpSamplingRuntimeConfig, ) from tokenspeed.runtime.sampling.logits_layout import ( LogitsLayoutExecutor, LogitsLayoutPlan, ) from tokenspeed.runtime.utils import get_colorful_logger logger = get_colorful_logger(__name__) @dataclasses.dataclass class LogitsProcessorOutput: ## Part 1: This part will be assigned in python/tokenspeed/runtime/layers/logits_processor.py::LogitsProcessor # The logits of the next tokens. shape: [#seq, vocab_size] next_token_logits: torch.Tensor # Used when ``do_argmax=True``. shape: [#seq] next_token_ids: torch.Tensor | None = None # Used by speculative decoding. # The last hidden layers hidden_states: torch.Tensor | None = None logits_layout_plan: LogitsLayoutPlan | None = None ## Part 2: Populated by the active SamplingBackend during sample()/verify(). # The logprobs of the next tokens. shape: [#seq] next_token_logprobs: torch.Tensor | None = None # The logprobs and ids of the top-k tokens in output positions. shape: [#seq, k] next_token_top_logprobs_val: list | None = None next_token_top_logprobs_idx: list | None = None # The logprobs and ids of the requested token ids in output positions. shape: [#seq, n] (n is the number of requested token ids) next_token_token_ids_logprobs_val: list | None = None next_token_token_ids_logprobs_idx: list | None = None ## Part 3: Prefill-only. This part will be assigned in python/tokenspeed/runtime/layers/logits_processor.py::LogitsProcessor # The logprobs of input tokens. shape: [#token] input_token_logprobs: torch.Tensor | None = None # The logprobs and ids of the top-k tokens in input positions. shape: [#seq, #token, k] input_top_logprobs_val: list = None input_top_logprobs_idx: list = None # The logprobs and ids of the requested token ids in input positions. shape: [#seq, n] (n is the number of requested token ids) input_token_ids_logprobs_val: list | None = None input_token_ids_logprobs_idx: list | None = None @dataclasses.dataclass class LogitsMetadata: forward_mode: ForwardMode capture_hidden_mode: CaptureHiddenMode = CaptureHiddenMode.NULL gather_ids: torch.Tensor | None = None extend_return_logprob: bool = False extend_return_top_logprob: bool = False extend_token_ids_logprob: bool = False extend_seq_lens_cpu: list[int] | None = None extend_logprob_start_lens_cpu: list[int] | None = None extend_logprob_pruned_lens_cpu: list[int] | None = None top_logprobs_nums: list[int] | None = None extend_input_logprob_token_ids_gpu: torch.Tensor | None = None token_ids_logprobs: list[list[int]] | None = None # logits and logprobs post processing temp_scaled_logprobs: bool = False temperature: torch.Tensor = None top_p_normalized_logprobs: bool = False top_p: torch.Tensor = None # DP attention metadata. Not needed when DP attention is not used. # Number of tokens in the request. global_num_tokens_gpu: torch.Tensor | None = None # The start position of local hidden states. dp_local_start_pos: torch.Tensor | None = None dp_local_num_tokens: torch.Tensor | None = None gathered_buffer: torch.Tensor | None = None # Buffer to gather logits from all ranks. forward_batch_gathered_buffer: torch.Tensor | None = None @classmethod def from_forward_context(cls, ctx: ForwardContext): return cls( forward_mode=ctx.forward_mode, capture_hidden_mode=ctx.capture_hidden_mode, gather_ids=ctx.gather_ids, ) _FUSED_LM_HEAD_GEMM = None def _get_fused_lm_head_gemm(): """Lazily import the fused lm_head GEMM kernel. The kernel is only present when tokenspeed-kernel was built with a compatible nvcc. Cache a sentinel when unavailable so we fall back to ``torch.matmul`` silently on subsequent calls. """ global _FUSED_LM_HEAD_GEMM if _FUSED_LM_HEAD_GEMM is not None: return _FUSED_LM_HEAD_GEMM if not current_platform().is_nvidia: _FUSED_LM_HEAD_GEMM = (None, None) return _FUSED_LM_HEAD_GEMM try: from tokenspeed_kernel.thirdparty.cuda.lm_head_gemm import ( lm_head_gemm, should_use_fused, ) _FUSED_LM_HEAD_GEMM = (should_use_fused, lm_head_gemm) except Exception: _FUSED_LM_HEAD_GEMM = (None, None) return _FUSED_LM_HEAD_GEMM def _lm_head_matmul(hidden_states: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: """Compute ``hidden_states @ weight.T``. Routes to the fused ``lm_head_gemm`` when the shape matches a compiled template and the bench-driven perf gate accepts (``should_use_fused``). Otherwise falls back to ``torch.matmul``. Only enabled for Kimi (``model_type == "kimi_k2"``) at the call site — on DSv3 the fused kernel's PDL launch surface caused a downstream EAGLE3 spec decode AR regression that we have not characterised end-to-end; on Kimi the perf win is the largest and the regression has not been reproduced, so we gate the fused path to Kimi only. """ cast_hidden = hidden_states.to(weight.dtype) should_use_fused, lm_head_gemm = _get_fused_lm_head_gemm() if should_use_fused is not None and should_use_fused(cast_hidden, weight): return lm_head_gemm(cast_hidden, weight, enable_pdl=True) return torch.matmul(cast_hidden, weight.T) class LogitsProcessor(nn.Module): _LOGITS_AG_MAX_TOKENS = 128 _LOGITS_AG_STATE_UNINITIALIZED = object() _LOGITS_AG_STATES = {} _LOGITS_DIST_ARGMAX_MAX_TOKENS = 8192 _LOGITS_DIST_ARGMAX_UNINITIALIZED = object() _LOGITS_DIST_ARGMAX_STATES = {} def __init__( self, config, skip_all_gather: bool = False, do_argmax: bool = False, logit_scale: float | None = None, tp_rank: int | None = None, tp_size: int | None = None, tp_group: tuple[int, ...] | None = None, ): super().__init__() self.config = config self.skip_all_gather = skip_all_gather self.do_argmax = do_argmax self.dp_sampling_enabled = False self.dp_num_tokens_per_req = 1 self.dp_sampling_min_bs = 0 self.logit_scale = logit_scale self._logits_layout_executor: LogitsLayoutExecutor | None = None if tp_rank is None: if tp_size is not None or tp_group is not None: raise ValueError("tp_size and tp_group require tp_rank.") tp_rank, tp_size = 0, 1 elif tp_size is None: raise ValueError("tp_size is required when tp_rank is provided.") if not 0 <= tp_rank < tp_size: raise ValueError(f"Invalid tensor-parallel rank: {tp_rank}/{tp_size}.") if tp_size != 1 and tp_group is None: raise ValueError("tp_group is required when tp_size > 1.") self.tp_rank, self.tp_size, self.tp_group = tp_rank, tp_size, tp_group self._all_gather_state = self._LOGITS_AG_STATE_UNINITIALIZED self._dist_argmax_state = self._LOGITS_DIST_ARGMAX_UNINITIALIZED self.final_logit_softcapping = getattr( self.config, "final_logit_softcapping", None ) if ( self.final_logit_softcapping is not None and self.final_logit_softcapping < 0 ): self.final_logit_softcapping = None # Gate the fused lm_head GEMM to Kimi only. See ``_lm_head_matmul``. self._use_fused_lm_head = getattr(self.config, "model_type", None) == "kimi_k2" def configure_dp_logits_layout(self, runtime: DpSamplingRuntimeConfig) -> None: if ( not runtime.enabled or runtime.topology is None or runtime.min_bs is None or runtime.max_bucket_bs is None or runtime.vocab_size is None or runtime.device is None ): raise RuntimeError("enabled DP sampling runtime is incomplete") topology = runtime.topology self.dp_sampling_enabled = True self.dp_num_tokens_per_req = runtime.num_tokens_per_req self.dp_sampling_min_bs = runtime.min_bs self._logits_layout_executor = LogitsLayoutExecutor( tp_rank=topology.tp_rank, tp_size=topology.tp_size, tp_group=topology.tp_group, max_bucket_bs=runtime.max_bucket_bs, num_tokens_per_req=runtime.num_tokens_per_req, vocab_size=runtime.vocab_size, device=runtime.device, ) def _resolve_logits_layout_plan( self, hidden_states: torch.Tensor, logits_metadata: LogitsMetadata, ) -> LogitsLayoutPlan | None: if not self.dp_sampling_enabled: return None if not logits_metadata.forward_mode.is_decode(): return None n = self.dp_num_tokens_per_req rows = hidden_states.shape[0] if rows % n != 0: raise ValueError(f"hidden_states have {rows} rows, not divisible by N={n}") effective_bs = rows // n bucket_bs = ((effective_bs + self.tp_size - 1) // self.tp_size) * self.tp_size if effective_bs < self.dp_sampling_min_bs: return None return LogitsLayoutPlan( effective_bs=effective_bs, bucket_bs=bucket_bs, tp_size=self.tp_size, num_tokens_per_req=n, ) def _init_all_gather_state(self, lm_head: VocabParallelEmbedding): if not current_platform().is_nvidia: return None if self.tp_size == 1 or self.skip_all_gather: return None vocab_padded = lm_head.weight.size(0) * self.tp_size if vocab_padded % (self.tp_size * 8) != 0: return None key = (self.tp_group, vocab_padded) if key not in self._LOGITS_AG_STATES: self._LOGITS_AG_STATES[key] = create_state( group=pg_manager.get_process_group("nccl", self.tp_group), rank_in_group=self.tp_rank, max_tokens=self._LOGITS_AG_MAX_TOKENS, hidden_size=vocab_padded, ) return self._LOGITS_AG_STATES[key] def _init_dist_argmax_state(self, lm_head: VocabParallelEmbedding): if not current_platform().is_nvidia: return None if self.tp_size == 1 or self.skip_all_gather or self.dp_sampling_enabled: return None vocab_per_rank = lm_head.weight.size(0) if vocab_per_rank * self.tp_size != self.config.vocab_size: return None # padded vocab: sharded argmax could pick a pad column if vocab_per_rank < 4096 or vocab_per_rank % 32 != 0: return None # below the kernel's vocab floor / alignment key = (self.tp_group, vocab_per_rank) if key not in self._LOGITS_DIST_ARGMAX_STATES: self._LOGITS_DIST_ARGMAX_STATES[key] = create_dist_argmax_state( group=pg_manager.get_process_group("nccl", self.tp_group), rank_in_group=self.tp_rank, max_M=self._LOGITS_DIST_ARGMAX_MAX_TOKENS, dtype=lm_head.weight.dtype, device=lm_head.weight.device, skip_ping_pong=True, ) return self._LOGITS_DIST_ARGMAX_STATES[key] def forward( self, input_ids, hidden_states, lm_head: VocabParallelEmbedding, logits_metadata: LogitsMetadata, aux_hidden_states: torch.Tensor | None = None, ) -> LogitsProcessorOutput: # Get the last hidden states and last logits for the next token prediction if not logits_metadata.extend_return_logprob: gather_ids = logits_metadata.gather_ids if gather_ids is not None: # Shapes align iff midlayer already pruned to one row per request # (draft first-step reduce). Other paths emit [N, H] with N > bs. if gather_ids.shape[0] == hidden_states.shape[0]: pruned_states = hidden_states if aux_hidden_states is not None: aux_pruned_states = list(aux_hidden_states) else: pruned_states = hidden_states[gather_ids] if aux_hidden_states is not None: aux_pruned_states = [h[gather_ids] for h in aux_hidden_states] else: if logits_metadata.forward_mode.is_extend_or_mixed(): raise RuntimeError( "EXTEND/MIXED forward must set gather_ids on ForwardContext" ) pruned_states = hidden_states if aux_hidden_states is not None: aux_pruned_states = list(aux_hidden_states) sample_indices = None input_logprob_indices = None else: # Input logprobs are required. # Find 3 different indices. # 1. pruned_states: hidden states that we want logprobs from. # 2. sample_indices: Indices that have sampled tokens. # 3. input_logprob_indices: Indices that have input logprob tokens. sample_index_pt = -1 sample_indices = [] input_logprob_indices_pt = 0 input_logprob_indices = [] pt, pruned_states = 0, [] for extend_logprob_start_len, extend_len in zip( logits_metadata.extend_logprob_start_lens_cpu, logits_metadata.extend_seq_lens_cpu, ): # It can happen in chunked prefill. We still need to sample 1 token, # But we don't want to include it in input logprob. if extend_len == extend_logprob_start_len: start_len = extend_logprob_start_len - 1 else: start_len = extend_logprob_start_len # We always need at least 1 token to sample because that's required # by a caller. if extend_len <= start_len: raise RuntimeError("extend_len must be greater than start_len.") pruned_states.append(hidden_states[pt + start_len : pt + extend_len]) pt += extend_len sample_index_pt += extend_len - start_len sample_indices.append(sample_index_pt) input_logprob_indices.extend( [ input_logprob_indices_pt + i for i in range(extend_len - extend_logprob_start_len) ] ) input_logprob_indices_pt += extend_len - start_len pruned_states = torch.cat(pruned_states) sample_indices = torch.tensor( sample_indices, device=pruned_states.device, dtype=torch.int64 ) input_logprob_indices = torch.tensor( input_logprob_indices, device=pruned_states.device, dtype=torch.int64 ) # Compute logits for both input and sampled tokens. logits_layout_plan = self._resolve_logits_layout_plan( pruned_states, logits_metadata ) logits = self._get_logits( pruned_states, lm_head, logits_metadata, plan=logits_layout_plan ) sampled_logits = ( logits[sample_indices] if sample_indices is not None else logits ) hidden_states_to_store: torch.Tensor | None = None if logits_metadata.capture_hidden_mode.need_capture(): if logits_metadata.capture_hidden_mode.is_full(): if aux_hidden_states is not None: aux_hidden_states = ( aux_hidden_states[0] if len(aux_hidden_states) == 1 else torch.cat(aux_hidden_states, dim=-1) ) hidden_states_to_store = aux_hidden_states else: hidden_states_to_store = hidden_states elif logits_metadata.capture_hidden_mode.is_last(): # Get the last token hidden states. If sample_indices is None, # pruned states only contain the last tokens already. if aux_hidden_states is not None: aux_pruned_states = ( aux_pruned_states[0] if len(aux_pruned_states) == 1 else torch.cat(aux_pruned_states, dim=-1) ) hidden_states_to_store = ( aux_pruned_states[sample_indices] if sample_indices is not None else aux_pruned_states ) else: hidden_states_to_store = ( pruned_states[sample_indices] if sample_indices is not None else pruned_states ) else: raise RuntimeError("Should never reach") if not logits_metadata.extend_return_logprob: # Decode mode or extend mode without return_logprob. # Greedy draft path: emit token ids here, fusing the cross-rank # vocab reduction into the argmax when gated on. next_token_ids = self._argmax(sampled_logits) if self.do_argmax else None return LogitsProcessorOutput( next_token_logits=sampled_logits, next_token_ids=next_token_ids, hidden_states=hidden_states_to_store, logits_layout_plan=logits_layout_plan, ) else: input_logprobs = logits[input_logprob_indices] del hidden_states, logits # Normalize the logprob w/o temperature, top-p pruned_lens = torch.tensor( logits_metadata.extend_logprob_pruned_lens_cpu, device=input_logprobs.device, ) if logits_metadata.temp_scaled_logprobs: logits_metadata.temperature = torch.repeat_interleave( logits_metadata.temperature.view(-1), pruned_lens, ).view(-1, 1) if logits_metadata.top_p_normalized_logprobs: logits_metadata.top_p = torch.repeat_interleave( logits_metadata.top_p, pruned_lens, ) input_logprobs = self.compute_temp_top_p_normalized_logprobs( input_logprobs, logits_metadata ) # Get the logprob of top-k tokens if logits_metadata.extend_return_top_logprob: ( input_top_logprobs_val, input_top_logprobs_idx, ) = self.get_top_logprobs(input_logprobs, logits_metadata) else: input_top_logprobs_val = input_top_logprobs_idx = None # Get the logprob of given token id if logits_metadata.extend_token_ids_logprob: ( input_token_ids_logprobs_val, input_token_ids_logprobs_idx, ) = self.get_token_ids_logprobs(input_logprobs, logits_metadata) else: input_token_ids_logprobs_val = input_token_ids_logprobs_idx = None input_token_logprobs = input_logprobs[ torch.arange(input_logprobs.shape[0], device=input_logprobs.device), logits_metadata.extend_input_logprob_token_ids_gpu, ] return LogitsProcessorOutput( next_token_logits=sampled_logits, input_token_logprobs=input_token_logprobs, input_top_logprobs_val=input_top_logprobs_val, input_top_logprobs_idx=input_top_logprobs_idx, hidden_states=hidden_states_to_store, input_token_ids_logprobs_val=input_token_ids_logprobs_val, input_token_ids_logprobs_idx=input_token_ids_logprobs_idx, ) def _get_logits( self, hidden_states: torch.Tensor, lm_head: VocabParallelEmbedding, logits_metadata: LogitsMetadata, embedding_bias: torch.Tensor | None = None, plan: LogitsLayoutPlan | None = None, ) -> torch.Tensor: """Get logits from hidden_states. If sampled_logits_only is True, it means hidden_states only contain the last position (e.g., extend without input logprobs). The caller should guarantee the given hidden_states follow this constraint. """ dp_sampling = plan is not None if dp_sampling and not self.dp_sampling_enabled: raise RuntimeError( "DP logits layout plan was provided but LogitsProcessor was not " "configured with dp_sampling" ) if dp_sampling and self.skip_all_gather: if self._logits_layout_executor is None: raise RuntimeError( "dp_sampling logits layout executor is not configured" ) hidden_states = self._logits_layout_executor.slice_hidden_states( hidden_states, plan ) if hasattr(lm_head, "weight"): if self._use_fused_lm_head: logits = _lm_head_matmul(hidden_states, lm_head.weight) else: logits = torch.matmul( hidden_states.to(lm_head.weight.dtype), lm_head.weight.T ) else: # GGUF models logits = lm_head.linear_method.apply(lm_head, hidden_states, embedding_bias) if self.logit_scale is not None: logits.mul_(self.logit_scale) if dp_sampling and not self.skip_all_gather: if self._logits_layout_executor is None: raise RuntimeError( "dp_sampling logits layout executor is not configured" ) logits = self._logits_layout_executor.swap_batch_vocab(logits, plan) elif not dp_sampling and self.tp_size > 1 and not self.skip_all_gather: if self.do_argmax: if self._dist_argmax_state is self._LOGITS_DIST_ARGMAX_UNINITIALIZED: self._dist_argmax_state = self._init_dist_argmax_state(lm_head) if ( self._dist_argmax_state is not None and not self.final_logit_softcapping and logits.size(0) <= self._LOGITS_DIST_ARGMAX_MAX_TOKENS ): return logits if self._all_gather_state is self._LOGITS_AG_STATE_UNINITIALIZED: self._all_gather_state = self._init_all_gather_state(lm_head) if ( self._all_gather_state is not None and logits.size(0) <= self._LOGITS_AG_MAX_TOKENS ): # skip_entry_sync=True assumes other sync points existing between two all_gather_inner calls. logits = all_gather_inner( self._all_gather_state, logits, tp_hidden_dim=logits.size(-1) * self.tp_size, skip_entry_sync=True, safe=False, ) else: num_rows = logits.size(0) local_vocab_size = logits.size(1) gathered_logits = torch.empty( self.tp_size * num_rows, local_vocab_size, dtype=logits.dtype, device=logits.device, ) all_gather_into_tensor(gathered_logits, logits, self.tp_group) logits = ( gathered_logits.view(self.tp_size, num_rows, local_vocab_size) .transpose(0, 1) .contiguous() .view(num_rows, local_vocab_size * self.tp_size) ) logits = logits[:, : self.config.vocab_size].contiguous() if self.final_logit_softcapping: fused_softcap_generic(logits, self.final_logit_softcapping) return logits def _argmax(self, logits: torch.Tensor) -> torch.Tensor: if ( self._dist_argmax_state not in (self._LOGITS_DIST_ARGMAX_UNINITIALIZED, None) and not self.final_logit_softcapping and logits.size(0) <= self._LOGITS_DIST_ARGMAX_MAX_TOKENS ): _, idx = distributed_argmax(self._dist_argmax_state, logits) return idx else: return sampling_argmax(logits) @staticmethod def get_top_logprobs(all_logprobs: torch.Tensor, logits_metadata: LogitsMetadata): max_k = max(logits_metadata.top_logprobs_nums) ret = all_logprobs.topk(max_k, dim=1) values = ret.values.tolist() indices = ret.indices.tolist() input_top_logprobs_val, input_top_logprobs_idx = [], [] pt = 0 for k, pruned_len in zip( logits_metadata.top_logprobs_nums, logits_metadata.extend_logprob_pruned_lens_cpu, ): if pruned_len <= 0: input_top_logprobs_val.append([]) input_top_logprobs_idx.append([]) continue input_top_logprobs_val.append( [values[pt + j][:k] for j in range(pruned_len)] ) input_top_logprobs_idx.append( [indices[pt + j][:k] for j in range(pruned_len)] ) pt += pruned_len return input_top_logprobs_val, input_top_logprobs_idx @staticmethod def get_token_ids_logprobs( all_logprobs: torch.Tensor, logits_metadata: LogitsMetadata ): input_token_ids_logprobs_val, input_token_ids_logprobs_idx = [], [] pt = 0 for token_ids, pruned_len in zip( logits_metadata.token_ids_logprobs, logits_metadata.extend_logprob_pruned_lens_cpu, ): if pruned_len <= 0: input_token_ids_logprobs_val.append([]) input_token_ids_logprobs_idx.append([]) continue input_token_ids_logprobs_val.append( [all_logprobs[pt + j, token_ids].tolist() for j in range(pruned_len)] ) input_token_ids_logprobs_idx.append([token_ids for _ in range(pruned_len)]) pt += pruned_len return input_token_ids_logprobs_val, input_token_ids_logprobs_idx @staticmethod def compute_temp_top_p_normalized_logprobs( last_logits: torch.Tensor, logits_metadata: LogitsMetadata ) -> torch.Tensor: """ compute logprobs for the output token from the given logits. Returns: torch.Tensor: logprobs from logits """ last_logits = last_logits.float() # Scale logits if temperature scaling is enabled if logits_metadata.temp_scaled_logprobs: last_logits = last_logits / logits_metadata.temperature # Normalize logprobs if top_p normalization is enabled # only normalize logprobs when top_p is set and not equal to 1.0 if ( logits_metadata.top_p_normalized_logprobs and (logits_metadata.top_p != 1.0).any() ): from tokenspeed.runtime.sampling.utils import top_p_normalize_probs_torch probs = torch.softmax(last_logits, dim=-1) del last_logits probs = top_p_normalize_probs_torch(probs, logits_metadata.top_p) return torch.log(probs) else: return torch.nn.functional.log_softmax(last_logits, dim=-1) @triton.jit def fused_softcap_kernel( full_logits_ptr, softcapping_value, n_elements, BLOCK_SIZE: tl.constexpr, ): pid = tl.program_id(0).to(tl.int64) block_start = pid * BLOCK_SIZE offsets = block_start + tl.arange(0, BLOCK_SIZE) mask = offsets < n_elements # Load values x = tl.load(full_logits_ptr + offsets, mask=mask).to(tl.float32) # Perform operations in-place x = x / softcapping_value # Stable tanh form; the exp ratio overflows to inf/inf for large logits. x = 2 * tl.sigmoid(2 * x) - 1 x = x * softcapping_value # Store result tl.store(full_logits_ptr + offsets, x, mask=mask) def fused_softcap(full_logits, final_logit_softcapping): n_elements = full_logits.numel() BLOCK_SIZE = 1024 grid = ((n_elements + BLOCK_SIZE - 1) // BLOCK_SIZE, 1, 1) fused_softcap_kernel[grid]( full_logits_ptr=full_logits, softcapping_value=final_logit_softcapping, n_elements=n_elements, BLOCK_SIZE=BLOCK_SIZE, ) return full_logits def fused_softcap_generic(full_logits, final_logit_softcapping): return fused_softcap(full_logits, final_logit_softcapping)