flashinfer.cake_rmsnorm_det_train.cake_rmsnorm_det_train_forward

flashinfer.cake_rmsnorm_det_train.cake_rmsnorm_det_train_forward(x: Tensor, w: Tensor, eps: float, *, residual: Tensor | None = None, out: Tensor | None = None, rstd: Tensor | None = None, residual_out: Tensor | None = None) → tuple[Tensor, ...]

Fused BF16 RMSNorm forward with FP32 statistics.

Parameters:

  • x (torch.Tensor) – BF16 [T, H] input of any layout. Row-strided views with a contiguous feature dimension, 32-byte aligned storage and a row stride that is a multiple of 32 bytes are read in place; other layouts (odd column offset, odd row stride, transposed storage) are materialized once. A tensor of more dimensions, [..., H], is normalized row by row when its feature dimension is contiguous and its leading dimensions are contiguous with one another in row-major order (stride(i) == stride(i + 1) * size(i + 1); a [T, heads, H] slice of a wider [T, heads, W] buffer qualifies): the leading dimensions collapse into prod(...) rows, the outputs take the input’s leading shape (rstd is [...]), w is shared by every row and the backward’s dw sums them all; any other N-D layout raises ValueError naming the rule. With residual this is the residual stream h and the normalized value is h_new = BF16(h + residual).

  • w (torch.Tensor) – BF16 or FP32 [H] weight. BF16 is read in place when contiguous and 32-byte aligned (the forward kernels load it with 32-byte vectors), FP32 – the effective weight 1 + w a zero-centred model computes in FP32 – when contiguous and 32-byte aligned; any other layout is copied once. An FP32 weight needs the FP32-weight route of the width on the device (weight_dtypes()), else ValueError.

  • eps (float) – Variance epsilon (FP32 inside the kernel).

  • residual (torch.Tensor, optional) – BF16 update of x’s shape added to x before normalization (same layout rules as x). Enables the fused residual-add route.

  • out (torch.Tensor, optional) – Preallocated outputs: BF16 y of x’s shape, FP32 reciprocal RMS of its leading shape ([T]) and (residual only) BF16 h_new of x’s shape. Contiguous aligned buffers are written directly; any other layout receives a copy after the launch.

  • rstd (torch.Tensor, optional) – Preallocated outputs: BF16 y of x’s shape, FP32 reciprocal RMS of its leading shape ([T]) and (residual only) BF16 h_new of x’s shape. Contiguous aligned buffers are written directly; any other layout receives a copy after the launch.

  • residual_out (torch.Tensor, optional) – Preallocated outputs: BF16 y of x’s shape, FP32 reciprocal RMS of its leading shape ([T]) and (residual only) BF16 h_new of x’s shape. Contiguous aligned buffers are written directly; any other layout receives a copy after the launch.

Returns:

  • (y, rstd) or (y, rstd, h_new) when residual is given.

  • rstd is the FP32 per-row 1 / sqrt(mean(h_new^2) + eps) the backward

  • needs; no other [T, H] temporary is produced.

Notes

Program selection by token count

Each route – one architecture, hidden size and direction – ships several programs of one kernel and selects one per call from the token count T: the program whose window above < T <= at_most admits it (cake_jit.program_windows returns the windows; the table below lists them). Within one window a row’s y, rstd and dx do not depend on T: two calls whose token counts fall in the same forward window (and, for the backward, the same backward window) return bitwise-identical results on their common rows. Across a forward edge between programs of different per-row thread groups the sum of squares is reduced in another order. On the rows both programs compute, rstd then moves by at most two FP32 ulps (three on SM107 at the residual route’s edge 1024) on 10.8-18.2 % of the rows (24.0-31.6 % on the residual route); y moves by one BF16 ulp, only on rows whose rstd moved, on at most 2.34 % of the rows (4.30 % on the residual route); dx follows on those rows and the residual output h_new does not change. Across every other forward edge, every edge of H = 256 and H = 5120 among them, the outputs are bitwise identical. Measured over three random inputs per edge, with BF16 and FP32 weights: on B200 (sm_100a) at most two ulps on up to 18.0 % of the rows (27.3 % on the residual route), y on up to 2.03 % (2.83 %); on GB300 (sm_103a) at most two ulps on up to 17.3 % (27.8 %), y on up to 2.03 % (2.34 %); on R200 (sm_107a, 212- and 216-SM boards) at most three ulps on up to 18.2 % (31.6 %), y on up to 2.34 % (4.30 %). Across a backward-only edge dx is bitwise identical: the H = 6144 route, which switched between the base backward kernel and the bulk-copy one with dx differing on about 1-2 % of the rows, now runs the bulk-copy kernel at every T on every architecture. dw sums every row and depends on T by definition; for a fixed T it is bitwise reproducible on devices with the same SM count. The backward splits the rows over ctas_per_sm x SM count CTAs, so devices of one architecture with different SM counts (B200 with 148 and GB200 with 152 SMs, R200 boards with 212 and 216 SMs) cut different chunk tables and can give different dw bits. In the window of its smallest token counts a backward route other than H = 256 runs a column-task program instead: one CTA sums each column of dw over every row in a fixed order (exact products and TwoSum, one FP32 rounding), so there dw does not depend on the SM count. Every backward route but H = 512 and the residual H = 6144 switches from plain to compensated dw sums at an accuracy edge (ACCURACY_EDGES) measured on B200 (148 SMs), GB300 (152 SMs) and R200 (212 SMs; the H = 5120 and H = 256 edges also on 216 SMs). A device with fewer SMs than ACCURACY_MIN_SMS gives each CTA more rows to sum at the same T; there these routes select their programs as if T were the edge plus one, so their compensated programs run at every T. Measured on B200 (sm_100a), GB300 (sm_103a) and R200 (sm_107a).

The windows per architecture and route, (above, at_most] and the program (the first 8 hex digits of its name). The mark after a window names the measured outcome at its upper edge, which the package tests hold: = bitwise-identical outputs on the common rows across it, ~ the forward’s measured differences (y by at most one BF16 ulp, rstd by at most two FP32 ulps, three on R200 (216 SMs) at the residual route’s edge 1024), dx the backward kernel switch, ? not yet measured (the tests bound the differences there).

The FP32-weight routes (fwd_w32_...) select among the FP32-weight twins of the same programs over the same windows, with the same outcome at every edge, and are not listed; a route re-cut for the FP32 weight (another twin in a window) is listed with its own windows.

  • sm_100a fwd_h1536: (0, 1536] 1ea88523 ~ (1536, 3584] 2fdbab2a = (3584, 6144] 7de1e510 = (6144, inf) 2fdbab2a

  • sm_100a fwd_h2048: (0, 1024] b44107e0 ~ (1024, 2048] 780b9833 ~ (2048, 4096] b44107e0 ~ (4096, inf) 780b9833

  • sm_100a fwd_h256: (0, 8192] 0489c784 = (8192, 40960] 1b284bae = (40960, 81920] 86d51eff = (81920, 114688] 1388ab77 = (114688, 196608] 466331ff = (196608, 327680] ed9483d8 = (327680, inf) b8f32dd4

  • sm_100a fwd_h4096: (0, 640] 8392d1a0 ~ (640, 1024] fd913ecd = (1024, 6144] d447b7f6 = (6144, inf) fd913ecd

  • sm_100a fwd_h512: (0, 1024] 8757f6da ~ (1024, 4096] 7673978d = (4096, 24576] 3d20ab4a = (24576, inf) 3430e5d9

  • sm_100a fwd_h5120: (0, 640] 5eef28cd = (640, 6144] b5e173f5 = (6144, inf) 5eef28cd

  • sm_100a fwd_w32_h5120: (0, 640] ba28d2c0 = (640, 6144] 9ecb62c2 = (6144, 16384] b49de9b4 = (16384, inf) 9ecb62c2

  • sm_100a fwd_h6144: (0, 768] bad4db91 = (768, 1024] bad4db91 = (1024, inf) 288d7b1d

  • sm_100a fwd_w32_h6144: (0, 768] 2763cd5e = (768, 1024] 2763cd5e = (1024, 6144] 3d619cd1 = (6144, inf) 5f89e5fd

  • sm_100a fwd_residual_h6144: (0, 256] 308da182 = (256, 1024] 03bdf566 ~ (1024, inf) e646bbdb

  • sm_103a fwd_h1536: (0, 1536] 1ea88523 ~ (1536, 3584] 2fdbab2a = (3584, 6144] 7de1e510 = (6144, inf) 2fdbab2a

  • sm_103a fwd_h2048: (0, 1024] c5a8099e = (1024, 4096] b44107e0 ~ (4096, inf) 780b9833

  • sm_103a fwd_h256: (0, 8192] 0489c784 = (8192, 20480] bbb05314 = (20480, 40960] 1b284bae = (40960, 81920] 86d51eff = (81920, 114688] 466331ff = (114688, 196608] 86d51eff = (196608, 327680] ed9483d8 = (327680, 368640] 37f66318 = (368640, inf) b8f32dd4

  • sm_103a fwd_h4096: (0, 640] 8392d1a0 ~ (640, 1024] fd913ecd = (1024, 3072] d447b7f6 = (3072, 6144] d447b7f6 = (6144, inf) fd913ecd

  • sm_103a fwd_h512: (0, 1024] 8757f6da ~ (1024, 4096] 7673978d = (4096, 16384] 3d20ab4a = (16384, inf) 3430e5d9

  • sm_103a fwd_h5120: (0, 640] 5eef28cd = (640, 6144] b5e173f5 = (6144, inf) 5eef28cd

  • sm_103a fwd_w32_h5120: (0, 640] 1defffdf = (640, 6144] 9ecb62c2 = (6144, inf) 9ecb62c2

  • sm_103a fwd_h6144: (0, 1024] bad4db91 = (1024, inf) 288d7b1d

  • sm_103a fwd_w32_h6144: (0, 1024] 2763cd5e = (1024, 6144] 2763cd5e = (6144, inf) 5f89e5fd

  • sm_103a fwd_residual_h6144: (0, 256] 308da182 = (256, 1024] 03bdf566 ~ (1024, inf) e646bbdb

  • sm_107a fwd_h1536: (0, 2048] 1ea88523 ~ (2048, 12288] 4cdc4c49 = (12288, inf) 682675bd

  • sm_107a fwd_h2048: (0, 1024] b44107e0 ~ (1024, 8192] a4da8513 = (8192, 12288] 0ad4d53c = (12288, 16384] 0ad4d53c = (16384, inf) bde56bd5

  • sm_107a fwd_w32_h2048: (0, 1024] 0da8e18d ~ (1024, 8192] a72c6ece = (8192, 12288] 595f5060 = (12288, 16384] 40210950 = (16384, inf) 40210950

  • sm_107a fwd_h256: (0, 6144] c5ca409f = (6144, 32768] bbb05314 = (32768, 98304] 466331ff = (98304, inf) 37f66318

  • sm_107a fwd_h4096: (0, 512] 8392d1a0 ~ (512, 1024] fd913ecd = (1024, 4096] d447b7f6 = (4096, 20480] bf933e16 = (20480, 49152] bf933e16 = (49152, inf) fd913ecd

  • sm_107a fwd_w32_h4096: (0, 512] 4371270e ~ (512, 1024] f0db3068 = (1024, 4096] f9f3423b = (4096, 20480] f9f3423b = (20480, 49152] f9f3423b = (49152, inf) f9f3423b

  • sm_107a fwd_h512: (0, 256] 7673978d = (256, 6144] 2a42e5c0 = (6144, inf) 60f47c5d

  • sm_107a fwd_h5120: (0, 512] fd1097c4 = (512, 1152] 5eef28cd = (1152, 6144] b5e173f5 = (6144, 20480] 044d4a2a = (20480, 49152] 044d4a2a = (49152, inf) 5eef28cd

  • sm_107a fwd_w32_h5120: (0, 512] ba28d2c0 = (512, 1152] 1defffdf = (1152, 6144] 9ecb62c2 = (6144, 20480] b49de9b4 = (20480, 49152] b49de9b4 = (49152, inf) b49de9b4

  • sm_107a fwd_h6144: (0, 1024] bad4db91 = (1024, 4096] cf1d0554 = (4096, inf) e7f4221a

  • sm_107a fwd_residual_h6144: (0, 256] e646bbdb ~ (256, 1024] 03bdf566 ~ (1024, inf) e646bbdb