docs: fix API indexes and RMSNorm RST rendering

docs/api/comm.rst

@@ -773,6 +773,40 @@ Prepared QKV workspaces and ordinary staging workspaces are not interchangeable.
 .. autoclass:: UlyssesQKVWorkspace
     :members:
 
+Low-Precision Implementation Primitives
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+The following internal helpers support numerical tests, trace definitions,
+and kernel benchmarks. They are not exported from ``flashinfer.comm``;
+applications should use the prepared ``UlyssesCommunicator`` QKV API above.
+
+.. currentmodule:: flashinfer.comm._ulysses_lowp
+
+.. autosummary::
+    :toctree: ../generated
+
+    capability
+    payload_spec
+    k_sum_v_amax
+    q_grouped_amax
+    k_grouped_amax
+    boundary_descriptors
+    merge_boundary_amax
+    k_boundary_minmax
+    derive_k_boundary_amax
+    zero_scale_and_padding
+    quant_q_into_payload
+    quant_kv_into_payload
+    quant_qkv_pack
+    unpack_for_sage
+    local_stats
+    finalize_stats
+    quant_and_pack
+    verify_duplicate_scale_slots
+    quant_v_fp8_with_scale
+
+.. currentmodule:: flashinfer.comm
+
 Head-Chunk Layout and Transport Primitives
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 

docs/api/gemm.rst

@@ -41,7 +41,7 @@ Quantized-base GEMM with runtime BF16 LoRA (SM90, SM100, SM103)
 BF16 LoRA factors consumed as given (Cake generated programs). ``lora_scale=None`` disables the
 adapter even if ``d`` / ``lora_b`` are passed; pass a device scalar to apply it.  The NVFP4
 operator also takes a prepared adapter (``lora_adapter=prepare_nvfp4_lora_adapter(lora_b,
-lora_scale, alpha)``: the adapter folded once at the switch, SVDQuant-equivalent numerics).
+lora_scale, alpha)``: the adapter is folded once at the switch, with SVDQuant-equivalent numerics).
 
 .. autosummary::
     :toctree: ../generated

docs/api/kda.rst

@@ -26,7 +26,7 @@ backend of its own. These prefill routes leave decode and speculative-decode
 dispatch unchanged.
 
 Automatic prefill policy on GB200 and GB300
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
 For one packed sequence, BF16 activations, head dimension 128 and an existing
 BF16 or FP32 state, ``auto`` uses these coarse performance regions:
@@ -114,7 +114,7 @@ and decode retains its existing layout contract.
     recurrent_kda
 
 Static PTX prefill on B300
-~~~~~~~~~~~~~~~~~~~~~~~~~
+~~~~~~~~~~~~~~~~~~~~~~~~~~
 
 ``flashinfer.recurrent_kda(..., backend="ptx")`` explicitly selects four
 retained PTX programs from `NVlabs/kda (260927-kda-for-kda)
@@ -715,7 +715,7 @@ should rotate fewer buffer sets, or call
 which releases all of it.
 
 Persistent CuTe DSL prefill with FP32 state
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
 ``recurrent_kda(..., backend="cute-dsl-persistent")`` selects the persistent
 M128/BT32 kernel adapted from humanfia/kda-for-kda-release, now

docs/api/mhc.rst

@@ -175,3 +175,19 @@ Names: ``flashinfer.cake_mhc_det_train.cake_mhc_det_train_<name>`` ->
 non-deterministic package's nodes). The legacy module
 ``flashinfer.cake_mhc_train`` keeps the same ``cake_mhc_train_*`` names as
 aliases of the deterministic objects, without a ``deterministic`` keyword.
+
+.. currentmodule:: flashinfer.cake_mhc_train
+
+.. autosummary::
+    :toctree: ../generated
+
+    cake_mhc_train_forward_prepare
+    cake_mhc_train_forward_stats
+    cake_mhc_train_forward_finish
+    cake_mhc_train_backward_prepare
+    cake_mhc_train_backward_finish
+    cake_mhc_train_prepare
+    cake_mhc_train_finish
+    cake_mhc_train_head_forward
+    cake_mhc_train_head_backward
+    cake_mhc_train_head

docs/api/norm.rst

@@ -78,7 +78,7 @@ The names the training RMSNorm had before the deterministic package took its
   call computes ``dx`` row by row and adds the ``dw`` sums of its CTAs into
   FP64 sums in the workspace, in the order the CTAs finish: ``dw`` is not
   guaranteed bitwise reproducible; ``dx`` is, and it is within one BF16 ulp of
-  max(|dx|, rms(dx)) of the deterministic backward's ``dx`` (not always bitwise
+  ``max(|dx|, rms(dx))`` of the deterministic backward's ``dx`` (not always bitwise
   equal to it). In the windows the backward's notes name, a twosum single kernel
   keeps each thread's ``dw`` column sums exact (error-free FP32 products, TwoSum
   compensation) before the FP64 additions, so its ``dw`` error does not grow

flashinfer/cake_rmsnorm_det_train_backend/cake_jit.py

@@ -15311,7 +15311,7 @@ def supported_hidden_sizes(arch: str) -> tuple[int, ...]:
 
 
 def weight_dtypes(arch: str, hidden: int, kind: str, residual: bool) -> tuple[str, ...]:
-    """The weight dtypes (tokens of ``WEIGHT_DTYPES``) the route ships for on ``arch``."""
+    """The weight dtypes (tokens of ``WEIGHT_DTYPES``) the route ships on ``arch``."""
 
     routes = ROUTES.get(arch, {})
     return tuple(

flashinfer/cake_rmsnorm_train.py

@@ -40,7 +40,7 @@
   CTAs claim row chunks, write ``dx`` and add their per-column ``dw`` sums into
   FP64 sums in the workspace in the order they finish; the CTA that finishes
   last rounds the sums into ``dw``.  ``dx`` is bitwise reproducible and within
-  one BF16 ulp of max(|dx|, rms(dx)) of the deterministic backward's ``dx``
+  one BF16 ulp of ``max(|dx|, rms(dx))`` of the deterministic backward's ``dx``
   (not always bitwise equal to it); ``dw`` is not guaranteed bitwise
   reproducible (the order of the FP64 additions depends on scheduling).  In
   the windows the backward docstring's notes name, a twosum single kernel keeps

flashinfer/cake_rmsnorm_train_backend/__init__.py

@@ -1190,7 +1190,7 @@ def cake_rmsnorm_train_backward(
         for it: in the windows that ship it, one kernel that writes ``dx`` and
         adds the ``dw`` sums of its CTAs into FP64 sums in the workspace, which
         the CTA that finishes last rounds into ``dw`` (``dx`` is bitwise
-        reproducible and within one BF16 ulp of max(|dx|, rms(dx)) of the
+        reproducible and within one BF16 ulp of ``max(|dx|, rms(dx))`` of the
         deterministic kernel's ``dx``, not always bitwise equal to it; the
         order of the FP64 additions
         depends on scheduling, so ``dw`` is not guaranteed bitwise

flashinfer/cake_rmsnorm_train_backend/cake_jit.py

@@ -15585,7 +15585,7 @@ def supported_hidden_sizes(arch: str) -> tuple[int, ...]:
 
 
 def weight_dtypes(arch: str, hidden: int, kind: str, residual: bool) -> tuple[str, ...]:
-    """The weight dtypes (tokens of ``WEIGHT_DTYPES``) the route ships for on ``arch``."""
+    """The weight dtypes (tokens of ``WEIGHT_DTYPES``) the route ships on ``arch``."""
 
     routes = ROUTES.get(arch, {})
     return tuple(

flashinfer/cake_rmsnorm_train_backend/cake_windows.py

@@ -1435,7 +1435,7 @@
 ``y`` by at most one BF16 ulp, on up to about 32 % of the rows (measured on B200, GB300
 and R200).  A backward window (``deterministic=False``) runs one of two forms.  The
 single kernel computes ``dx`` row by row (bitwise reproducible, within one BF16 ulp of
-max(|dx|, rms(dx)) of the deterministic backward's ``dx``, not always bitwise equal to
+``max(|dx|, rms(dx))`` of the deterministic backward's ``dx``, not always bitwise equal to
 it) and adds the ``dw`` sums of its CTAs into FP64 column sums in a
 scheduling-dependent order: its ``dw`` is not bitwise reproducible, even for a fixed
 ``T``.  The top window of ``bwd_h256`` on sm_100a and sm_107a, on sm_103a of
@@ -1501,7 +1501,7 @@
 ``y`` by at most one BF16 ulp, on up to about 32 % of the rows (measured on B200, GB300
 and R200).  A backward window (``deterministic=False``) runs one of two forms.  The
 single kernel computes ``dx`` row by row (bitwise reproducible, within one BF16 ulp of
-max(|dx|, rms(dx)) of the deterministic backward's ``dx``, not always bitwise equal to
+``max(|dx|, rms(dx))`` of the deterministic backward's ``dx``, not always bitwise equal to
 it) and adds the ``dw`` sums of its CTAs into FP64 column sums in a
 scheduling-dependent order: its ``dw`` is not bitwise reproducible, even for a fixed
 ``T``.  The top window of ``bwd_h256`` on sm_100a and sm_107a, on sm_103a of
@@ -1653,7 +1653,7 @@
 ``y`` by at most one BF16 ulp, on up to about 32 % of the rows (measured on B200, GB300
 and R200).  A backward window (``deterministic=False``) runs one of two forms.  The
 single kernel computes ``dx`` row by row (bitwise reproducible, within one BF16 ulp of
-max(|dx|, rms(dx)) of the deterministic backward's ``dx``, not always bitwise equal to
+``max(|dx|, rms(dx))`` of the deterministic backward's ``dx``, not always bitwise equal to
 it) and adds the ``dw`` sums of its CTAs into FP64 column sums in a
 scheduling-dependent order: its ``dw`` is not bitwise reproducible, even for a fixed
 ``T``.  The top window of ``bwd_h256`` on sm_100a and sm_107a, on sm_103a of

flashinfer/gemm/cake_quant_lora_gemm.py

@@ -743,7 +743,7 @@ def _lora_carrier(record: dict, device: torch.device) -> torch.Tensor:
 
 
 class StaleAdapterError(ValueError):
-    """A prepared adapter's source tensor (``lora_b`` / ``lora_scale`` / ``alpha``) moved since its prepare (snapshot
+    """A prepared adapter's source tensor (``lora_b`` / ``lora_scale`` / ``alpha``) moved since its preparation (snapshot
     semantics): re-prepare (``prepare_nvfp4_lora_adapter(..., out=adapter)``) or capture the refold inside the graph."""
 
 

flashinfer/quantization/kernels/nvfp4_kv_materialize.py

@@ -722,7 +722,23 @@ def __init__(
         device: Union[str, int, torch.device] = "cuda",
         uses_shared_paged_kv_idx: bool = True,
     ) -> None:
-        """Check the layout and device, and compile or load the kernel."""
+        """Check the layout and device, and compile or load the kernel.
+
+        Parameters
+        ----------
+        num_kv_heads : int
+            Number of KV heads per page.
+        page_size : int
+            Tokens per page. Must be divisible by 4.
+        head_dim : int
+            Head dimension. Must be divisible by 64.
+        device : Union[str, int, torch.device]
+            CUDA device that holds the tensors passed to :meth:`run`.
+        uses_shared_paged_kv_idx : bool
+            Whether K and V share page indices. ``block_tables`` is
+            ``[batch_size, max_num_pages]`` if so, else
+            ``[batch_size, 2, max_num_pages]``.
+        """
         device = torch.device(device)
         if device.type != "cuda":
             raise ValueError(f"device must be a CUDA device, got {device}")