flashinfer.cake_rmsnorm_det_train.cake_rmsnorm_det_train_backward

flashinfer.cake_rmsnorm_det_train.cake_rmsnorm_det_train_backward(g: Tensor, x: Tensor, w: Tensor, rstd: Tensor, *, deterministic: bool = True, g_residual: Tensor | None = None, workspace: Tensor | None = None, dx: Tensor | None = None, dw: Tensor | None = None) → tuple[Tensor, Tensor]

Fused BF16 RMSNorm backward: dx and the FP32 weight gradient dw.

Parameters:

  • g (torch.Tensor) – BF16 [T, H] gradient of y of any layout. Views with a contiguous feature dimension, 16-byte aligned storage and a row stride that is a multiple of 16 bytes (expanded or sliced autograd gradients included) are read in place; other layouts are materialized once. An N-D gradient, [..., H], takes x’s shape and collapses its leading dimensions into rows under the forward’s rule.

  • x (torch.Tensor) – The BF16 [T, H] (or [..., H]) tensor that was normalized (x of the plain forward, h_new of the residual forward); same layout rules as the forward’s x.

  • w (torch.Tensor) – BF16 or FP32 [H] weight (the forward’s). BF16 is read in place when contiguous and 16-byte aligned, FP32 when contiguous and 32-byte aligned; any other layout is copied once. An FP32 weight needs the FP32-weight route of the width (weight_dtypes()), else ValueError.

  • rstd (torch.Tensor) – FP32 [T] reciprocal RMS returned by the forward (of the leading shape [...] for an N-D call, or flat [prod(...)]).

  • deterministic (bool) – Only the deterministic fixed-order dw reduction is exported; False is rejected.

  • g_residual (torch.Tensor, optional) – BF16 gradient of g’s shape flowing directly through h_new (residual forward), any layout. It is added once so that the returned dx is the gradient of both residual-add inputs.

  • workspace (torch.Tensor, optional) – Caller-owned uint8 buffer from cake_rmsnorm_det_train_backward_workspace() sized for at least this many rows; allocated when omitted.

  • dx (torch.Tensor, optional) – Preallocated outputs: BF16 dx of x’s shape and FP32 [H]. Contiguous aligned buffers are written directly; any other layout receives a copy after the launch.

  • dw (torch.Tensor, optional) – Preallocated outputs: BF16 dx of x’s shape and FP32 [H]. Contiguous aligned buffers are written directly; any other layout receives a copy after the launch.

Returns:

  • (dx, dw) with dw in FP32; the caller decides whether to cast it.

  • For a fixed token count dw is bitwise reproducible on devices with the

  • same SM count (the rows are split over ctas_per_sm x SM count CTAs, so)

  • devices of one architecture with different SM counts can give different

  • dw bits.

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 (bwd_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 bwd_h1536: (0, 109] a772b70f ? (109, 1024] 7dc29629 = (1024, 4096] c1a9efdf = (4096, 32768] 21c8bb55 = (32768, 36352] b457c4d1 = (36352, inf) b5bdabba

  • sm_100a bwd_h2048: (0, 288] 24d642a5 ? (288, 3072] caf6978e = (3072, 4096] 9ac8d21e ~ (4096, 16384] c73074cf = (16384, 58880] ceeef043 = (58880, inf) 0cb817d1

  • sm_100a bwd_h256: (0, 87] 083ea823 ? (87, 2048] a385ae16 = (2048, 5120] af005462 = (5120, 12288] a385ae16 = (12288, 40960] af005462 = (40960, 80896] a385ae16 = (80896, 1572864] 083ea823 = (1572864, inf) c18297d7

  • sm_100a bwd_h4096: (0, 203] bbc80180 ? (203, 256] 3e644ea1 = (256, 1024] 59270fb2 = (1024, 4096] 9a830956 = (4096, 44032] 35a0b0ca = (44032, inf) 7122dbcd

  • sm_100a bwd_h512: (0, 227] 5f3e3fcc ? (227, 1024] 09601d42 ~ (1024, 4096] 123db0b2 = (4096, 32768] ad1a5a3c = (32768, inf) 8877322a

  • sm_100a bwd_h5120: (0, 121] 6a1cfc92 ? (121, 256] db01cec1 = (256, 6144] b482c84d = (6144, inf) c145ef19

  • sm_100a bwd_h6144: (0, 20] e60cfafd ? (20, 256] b93bf060 = (256, 4096] 5c804aed = (4096, 42496] b93bf060 = (42496, inf) c440e57b

  • sm_100a bwd_residual_h6144: (0, 32] d2749326 ? (32, inf) 0dcdc563

  • sm_103a bwd_h1536: (0, 31] ad547a3c ? (31, 214] 53285b95 ? (214, 1024] 7dc29629 = (1024, 4096] c1a9efdf = (4096, 14928] 1df7db17 = (14928, 15566] 4df08ce5 = (15566, 40448] f522f5b8 = (40448, inf) cbeae168

  • sm_103a bwd_w32_h1536: (0, 31] 775d9130 ? (31, 214] 5c5be9a6 ? (214, 1024] baff9ecb = (1024, 4096] 17d9cdfd = (4096, 14928] a99e2b38 = (14928, 15566] 7a72e97b = (15566, 40448] 109f19bb = (40448, inf) ca6211f3

  • sm_103a bwd_h2048: (0, 2] ce658895 ? (2, 171] 2b4c0f5b ? (171, 3072] caf6978e = (3072, 4096] 9ac8d21e ~ (4096, 16384] caf6978e = (16384, 32768] caf6978e = (32768, 55296] 302f64db = (55296, inf) 018f0cd7

  • sm_103a bwd_h256: (0, 1680] 083ea823 ? (1680, 2048] a385ae16 = (2048, 5120] af005462 = (5120, 12288] a385ae16 = (12288, 40960] af005462 = (40960, 99840] a385ae16 = (99840, 720896] 083ea823 = (720896, 1572864] 0291949b = (1572864, inf) c18297d7

  • sm_103a bwd_h4096: (0, 152] ff0156d5 ? (152, 256] 50044ee3 = (256, 1024] ad2bb26e = (1024, 4096] 24643edc = (4096, 48640] 35a0b0ca = (48640, inf) 850eec3a

  • sm_103a bwd_h512: (0, 4] 117a66e4 ? (4, 227] 877751ba ? (227, 1024] 09601d42 ~ (1024, 4096] 8e7f4d81 = (4096, 32768] 274b0055 = (32768, inf) cfaae3ce

  • sm_103a bwd_h5120: (0, 25] 505283a5 ? (25, 5120] 2fc4932f = (5120, inf) 7ffb443a

  • sm_103a bwd_h6144: (0, 20] 38b1c6bc ? (20, 4096] 5c804aed = (4096, 48640] 322574af = (48640, inf) 289dc9ab

  • sm_103a bwd_residual_h6144: (0, 32] 07135550 ? (32, inf) 0dcdc563

  • sm_107a bwd_h1536: (0, 576] 9e71ed0d ? (576, 1024] 7dc29629 = (1024, 16384] 6a428da2 = (16384, 45056] 56d26996 = (45056, inf) e694b0a2

  • sm_107a bwd_h2048: (0, 256] 32eb12e2 ? (256, 1024] 845cdcd8 ~ (1024, 16384] 81c3d8ab = (16384, 50688] 845cdcd8 = (50688, inf) 694decfb

  • sm_107a bwd_h256: (0, 1572864] 083ea823 = (1572864, inf) c18297d7

  • sm_107a bwd_h4096: (0, 256] 44511fa0 ? (256, 1024] a647fccc = (1024, 4096] d41a514f = (4096, 24064] ee4a9053 = (24064, inf) 99a01f2a

  • sm_107a bwd_h512: (0, 768] bcd8d6a9 ? (768, 1024] 3feed502 = (1024, 4096] 3feed502 = (4096, 16384] dfaf378a ? (16384, inf) eb141cc9

  • sm_107a bwd_h5120: (0, 1280] 505283a5 ? (1280, 4096] 838c077d ? (4096, 19968] 6816fb80 ? (19968, inf) 47da668e

  • sm_107a bwd_h6144: (0, 1664] 57d45549 ? (1664, 19456] 5d292172 = (19456, inf) 19e18371

  • sm_107a bwd_residual_h6144: (0, 960] c15868b8 ? (960, inf) 0dcdc563