This document is relevant for: Trn2, Trn3

GPT-OSS Tail Chunked Megakernel Kernel API Reference#

GPT-OSS decode sampling tail (RMSNorm, lm_head, top-k, sampler) with a single collective.

Background#

The gpt_oss_tail_chunked_megakernel kernel implements the GPT-OSS decode sampling tail (final RMSNorm, lm_head, top-k, and sampler) with a single all-to-all collective. It brings the full vocabulary onto each rank and computes the global top-k locally with chunked_topk.

API Reference#

Source code for this kernel API can be found at: tail_chunked_megakernel.py

gpt_oss_tail_chunked_megakernel#

nkilib.experimental.gpt_oss_tail.gpt_oss_tail_chunked_megakernel(hidden_states: nl.NkiTensor, gamma: nl.NkiTensor, lm_head_weight_t: nl.NkiTensor, eps: float = 1e-05, hidden_actual: Optional[int] = None, zero_interleaved_pad: bool = True, k: int = 256, sampling_params: Optional[nl.NkiTensor] = None, replica_groups: Optional[tuple] = None, hop_meshes: Optional[tuple] = None, dummy_collectives: bool = False, emit_logits: bool = True, debug_taps: bool = False, deferred_rms: bool = True, skip_nan_fold: bool = True, a2a_batch_groups: int = 1) → Tuple[nl.NkiTensor, ...]#

GPT-OSS decode sampling tail (RMSNorm, lm_head, top-k, sampler) with a single collective.

Parameters:

  • hidden_states (nl.NkiTensor) – [gbs, H] post-MoE hidden states in HBM; fp32, fp16, or bf16. A 16-bit input is widened to fp32 through an explicit on-chip conversion copy immediately after its (half-sized) load, which is EXACT (both half formats are subsets of fp32) – so it produces bitwise the same outputs as passing a pre-widened fp32 copy, while halving this kernel’s largest input DMA and letting a 16-bit producer skip an HBM widen round trip. The copy is load-bearing: engines fed 16-bit data directly compute at input precision.

  • gamma (nl.NkiTensor) – [1, H] final RMSNorm weight, fp32.

  • lm_head_weight_t (nl.NkiTensor) – [H, V_s] bfloat16 lm_head shard for this rank, ALREADY TRANSPOSED and C-contiguous. V_s = V/VTP. bf16 because that is the checkpoint’s own precision.

  • eps (float) – RMSNorm epsilon (1e-5 for GPT-OSS).

  • hidden_actual (Optional[int]) – H_u, the unpadded hidden (2880). Defaults to H.

  • zero_interleaved_pad (bool) – True for the mxfp4 checkpoint (interleaved padding lanes).

  • k (int) – candidates kept (256 in production).

  • emit_logits (bool) – whether to return the fp32 [gbs, V_s] pre-collective logits. They are a DEBUG/GATING output – nothing downstream of this kernel reads them; the collective carries the bf16 mirror instead. Keeping them costs a full extra HBM write of the shard (805 KB at the production shape, roughly half this kernel’s hbm_write) plus its DMA and the fp32 SBUF tile. The tests set it True so the norm and the lm_head GEMM stay independently gated; a deployment that only wants tokens should set it False.

  • sampling_params (Optional[nl.NkiTensor]) – [gbs/VTP, 3] or [gbs/VTP, 4] fp32. Columns are [top_k, top_p, temperature]; the optional 4th column is the per-row inverse-transform draw u in [0, 1) – the RUNTIME form of the compiled-in _DETERMINISTIC_DRAW (0.5), so sweeping the draw is a data change rather than a NEFF rebuild. Rows are THIS RANK’S only – see owned_row_block for which rows those are when more than one hop is used. When None the sampler is skipped and no token is returned.

  • replica_groups (Optional[tuple]) – one group listing all VTP ranks IN RANK ORDER. Shorthand for a single-hop hop_meshes=(replica_groups,); supply exactly one of the two.

  • hop_meshes (Optional[tuple]) – one replica-group mesh per all-to-all hop, from compose_hops. Hop sizes must multiply to VTP. Needed because an all-to-all replica group must be device-local or one-core-per-device.

  • dummy_collectives (bool) – prepend a zero all-reduce fence per hop mesh so the fence, not the real hops, absorbs rank-arrival skew. A PROFILING instrument: it cannot change the answer but it adds a collective per mesh to every call, so it defaults off and profiling runs opt in. Without it a per-hop duration is mostly dispatch jitter.

  • debug_taps (bool) – forward chunked_topk’s seven diagnostic taps as extra outputs, APPENDED after everything else so the existing output positions do not move. A DIAGNOSTIC instrument for hardware-only index defects – see DEBUG TAPS in chunked_topk. It changes this kernel’s arity, so leave it off unless the consumer expects them.

  • deferred_rms (bool) – apply the per-row RMSNorm scale on the GEMM’s PSUM-read copies instead of on the GEMM’s input (exact algebra: rms is a per-row scalar and the GEMM is linear), which stops every matmul from waiting on the full row’s sum of squares – measured -3.8 us on the fenced gbs=128 head. Numerically the SAME one-bf16-rounding-per-element bound; only the last-ulp realization moves, so fixed-seed sampled tokens can shift at cdf crossings (same class as any resummation change). Default True; a SAFETY FALLBACK to the legacy in-chain form – (x*rms)*gamma feeding the GEMM, the pre-deferred realization bit-for-bit – is one trace-time flag away if a consumer needs it.

  • skip_nan_fold (bool) – skip chunked_topk’s stage-1 NaN fold. Sound here because the top-k input is this kernel’s OWN GEMM output: finite hidden/gamma/weight inputs give finite bf16 products summed in fp32, and even an overflow saturates to +/-inf – which nisa.topk orders fine (the fold only exists for NaN, which would take inf arithmetic INSIDE the GEMM, i.e. already-garbage inputs). The fold costs a full-width compare + predicated copy per stage-1 tile on the post-collective serial chain. Default True; the fold is one trace-time flag away for a caller that wants the belt-and-suspenders behavior back. Outputs are BITWISE identical either way on any NaN-free input – unlike deferred_rms this cannot move even an ulp, because it deletes ops whose effect on finite data is the identity.

  • a2a_batch_groups (int) – split the exchange into this many batch-group collectives, one per gbs/G whole lm_head tiles (single-hop only; G must divide rows_per_rank and the geometry must support the receive-layout fold – asserted; groups aligned to whole 8-row sort tiles overlap best but are not required). Group g’s exchange fires as soon as its own lm_head tiles’ mirror stores land – under the rest of the GEMM – pulling the exchange forward by up to (G-1)/G of its wire time minus (G-1) launch serializations. OWNERSHIP CONTRACT CHANGES with G > 1: rank r owns the STRIDED row set given by owned_rows (one extra mixed-radix digit), not owned_row_block’s contiguous block – sampling_params rows and every consumer of the per-rank outputs must use that map. Gate-validated, not bitwise: regrouping a collective changes the runtime’s internal transfer realization.

compose_hops#

nkilib.experimental.gpt_oss_tail.compose_hops(vtp: int, hops: Tuple[int, ...])#

Replica-group meshes for a sequence of all-to-all hops whose sizes multiply to vtp.

owned_row_block#

nkilib.experimental.gpt_oss_tail.owned_row_block(rank: int, vtp: int, gbs: int, hops=None)#

Batch rows that rank owns after the all-to-all. Returns (row_start, n_rows).

owned_rows#

nkilib.experimental.gpt_oss_tail.owned_rows(rank: int, vtp: int, gbs: int, hops=None, batch_groups: int = 1)#

The GLOBAL batch rows rank owns, IN RECEIVE ORDER, as a list.

This document is relevant for: Trn2, Trn3