1+ # Copyright 2026 FlagOS Contributors
2+ #
3+ # Licensed under the Apache License, Version 2.0 (the "License");
4+ # you may not use this file except in compliance with the License.
5+ # You may obtain a copy of the License at
6+ #
7+ # http://www.apache.org/licenses/LICENSE-2.0
8+ #
9+ # Unless required by applicable law or agreed to in writing, software
10+ # distributed under the License is distributed on an "AS IS" BASIS,
11+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+ # See the License for the specific language governing permissions and
13+ # limitations under the License.
14+
15+ import logging
16+
17+ import torch
18+ import triton
19+ import triton .language as tl
20+
21+ logger = logging .getLogger (__name__ )
22+
23+
24+ @triton .jit
25+ def _beam_search_score_kernel (
26+ log_probs ,
27+ beam_scores ,
28+ output ,
29+ N ,
30+ V : tl .constexpr ,
31+ BLOCK : tl .constexpr ,
32+ NEED_MASK : tl .constexpr ,
33+ NEED_RNE : tl .constexpr ,
34+ ):
35+ """Flat 1D beam search score kernel: out[i] = log_probs[i] + beam_scores[i // V].
36+
37+ Continuous flat index space [0, N) with N = batch * vocab. `V` is a
38+ constexpr so the row division `offs // V` lowers to a shift (V is a power
39+ of two in every exercised shape); each lane then adds the scalar beam
40+ score of its row. NEED_MASK covers the tail when N % BLOCK != 0.
41+
42+ NEED_RNE enables a manual round-to-nearest-even emulation of the
43+ fp32->bf16 conversion before the store: the Kunlunxin backend lowers
44+ fp32->bf16 casts with round-toward-zero, which differs from torch's RNE
45+ semantics on ~10% of elements (1 ULP). fp16/fp32 store conversions on this
46+ backend are already RNE-correct / exact, so the emulation is only applied
47+ for bf16 outputs.
48+ """
49+ pid = tl .program_id (0 )
50+ offs = pid * BLOCK + tl .arange (0 , BLOCK )
51+ if NEED_MASK :
52+ mask = offs < N
53+ row = offs // V
54+ v = tl .load (log_probs + offs , mask = mask , other = 0.0 ).to (tl .float32 )
55+ b = tl .load (beam_scores + row , mask = mask , other = 0.0 ).to (tl .float32 )
56+ else :
57+ row = offs // V
58+ v = tl .load (log_probs + offs ).to (tl .float32 )
59+ b = tl .load (beam_scores + row ).to (tl .float32 )
60+ acc = v + b
61+ if NEED_RNE :
62+ bits = acc .to (tl .int32 , bitcast = True )
63+ lsb = (bits >> 16 ) & 1
64+ rnd = (bits + 0x7FFF + lsb ) & - 65536 # RNE round to bf16 precision
65+ out_val = rnd .to (tl .float32 , bitcast = True )
66+ else :
67+ out_val = acc
68+ if NEED_MASK :
69+ tl .store (output + offs , out_val , mask = mask )
70+ else :
71+ tl .store (output + offs , out_val )
72+
73+
74+ def _block_and_warps (numel , dtype ):
75+ """Empirically tuned per-size dispatch (XPU7 sweep, 2026-08-17).
76+
77+ Flat BLOCK values: larger tiles reduce program count for launch-bound
78+ big shapes; 8192-class tiles win for small shapes. bf16 keeps 16384 at
79+ the largest size because the RNE emulation path degrades on 64K-lane
80+ tiles.
81+ """
82+ if dtype == torch .float32 :
83+ if numel <= 131072 :
84+ return 8192 , 4
85+ return 65536 , 4
86+ if dtype == torch .float16 :
87+ if numel <= 32768 :
88+ return 8192 , 8
89+ if numel <= 131072 :
90+ return 16384 , 8
91+ if numel <= 524288 :
92+ return 16384 , 4
93+ return 65536 , 8
94+ # bfloat16
95+ if numel <= 32768 :
96+ return 8192 , 8
97+ if numel <= 131072 :
98+ return 16384 , 8
99+ if numel <= 524288 :
100+ return 16384 , 2
101+ return 16384 , 8
102+
103+
104+ def _launch_beam_search_score (log_probs , beam_scores , outputs ):
105+ if log_probs .dim () != 2 :
106+ raise ValueError ("beam_search_score expects 2D log_probs on Kunlunxin" )
107+ batch_size , vocab_size = log_probs .shape
108+ if beam_scores .numel () != batch_size :
109+ raise ValueError (
110+ "beam_scores must contain one score per batch entry on Kunlunxin"
111+ )
112+ numel = log_probs .numel ()
113+ if numel == 0 or batch_size == 0 :
114+ return outputs
115+ if not log_probs .is_contiguous ():
116+ log_probs = log_probs .contiguous ()
117+ beam_flat = beam_scores
118+ if not beam_flat .is_contiguous ():
119+ beam_flat = beam_flat .contiguous ()
120+ beam_flat = beam_flat .reshape (- 1 )
121+ block , num_warps = _block_and_warps (numel , log_probs .dtype )
122+ need_mask = 1 if numel % block else 0
123+ grid = (triton .cdiv (numel , block ),)
124+ _beam_search_score_kernel [grid ](
125+ log_probs ,
126+ beam_flat ,
127+ outputs ,
128+ numel ,
129+ V = vocab_size ,
130+ BLOCK = block ,
131+ NEED_MASK = need_mask ,
132+ NEED_RNE = log_probs .dtype == torch .bfloat16 ,
133+ num_warps = num_warps ,
134+ )
135+ return outputs
136+
137+
138+ def _flat_beam_scores (beam_scores , batch_size ):
139+ """Normalize beam_scores to [B] flat. Accepts 1D [B] or 2D [B, 1]."""
140+ if beam_scores .dim () > 2 or (
141+ beam_scores .dim () == 2 and beam_scores .shape [- 1 ] != 1
142+ ):
143+ raise ValueError (
144+ "beam_scores must have shape [batch] or [batch, 1] on Kunlunxin"
145+ )
146+ return beam_scores .reshape (batch_size )
147+
148+
149+ def beam_search_score (log_probs , beam_scores ):
150+ """Out-of-place beam search score: log_probs [B, V] + beam_scores [B]."""
151+ logger .debug ("GEMS_KUNLUNXIN BEAM_SEARCH_SCORE" )
152+ batch_size = log_probs .shape [0 ]
153+ beam_flat = _flat_beam_scores (beam_scores , batch_size )
154+ outputs = torch .empty_like (log_probs )
155+ return _launch_beam_search_score (log_probs , beam_flat , outputs )
156+
157+
158+ def beam_search_score_ (log_probs , beam_scores ):
159+ """In-place variant writing back into log_probs."""
160+ logger .debug ("GEMS_KUNLUNXIN BEAM_SEARCH_SCORE_" )
161+ batch_size = log_probs .shape [0 ]
162+ beam_flat = _flat_beam_scores (beam_scores , batch_size )
163+ return _launch_beam_search_score (log_probs , beam_flat , log_probs )
0 commit comments