@@ -1,27 +1,26 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
 # SPDX-License-Identifier: Apache-2.0
 
-"""Blackwell FMHA-prefill port from CUTLASS Python DSL to CUDA Lang.
+"""Blackwell FMHA-prefill CUDA Lang tutorial.
 
 This is a native fused implementation, with no SIMT, PyTorch-SDPA, library,
-source-kernel, or precompiled-PTX fallback.  It supports the source tutorial's
-fixed and variable sequence storage, D=32/64/128, FP16/BF16/E4M3 inputs and
-outputs, MHA/GQA, mask variants, LSE, sink logits, scaling controls, and
-correction/epilogue options.
+or precompiled-PTX fallback. It supports fixed and variable sequence storage,
+D=32/64/128, FP16/BF16/E4M3 inputs and outputs, MHA/GQA, mask variants, LSE,
+sink logits, scaling controls, and correction/epilogue options.
 
-The device kernel retains the source's structural contract:
+The device kernel uses:
 
 * one 512-thread CTA and one CTA per ``(256 Q rows, Q head, batch)`` tile;
 * softmax warps 0-7, correction warps 8-11, MMA warp 12, TMA warp 13,
   epilogue warp 14, and CLC scheduler warp 15;
 * CTA-group-1 tcgen05 QK (128x128x16) and PV (128x64x16);
 * two Q stages, three unified K/V stages, two output stages, and all 512
-  TMEM columns with the source S/P/O aliases;
+  TMEM columns with S/P/O aliases;
 * online softmax, packed-FP16 P in TMEM, FP32 O correction in TMEM, TMA
   global/shared transfers, and a one-slot 480-consumer CLC mailbox.
 
-The source-compatible host configuration, chunked reference, correctness
-matrix, cold-L2/CUDA-graph benchmark, and CLI share this public entry point.
+The host configuration, chunked reference, correctness matrix,
+cold-L2/CUDA-graph benchmark, and CLI share this public entry point.
 Invalid combinations raise clear diagnostics and are never silently redirected
 to another attention implementation.
 """
@@ -376,9 +375,8 @@ def _store_output_values(
     word_count = 16
     if output_kind == ELEMENT_E4M3:
         word_count = 8
-    # One PTX shared-memory vector store carries four 32-bit registers.  The
-    # swizzle is constant within each 16-byte group, so form the same v4
-    # transfer units as CUTLASS's store_swizzled lowering.
+    # One PTX shared-memory vector store carries four 32-bit registers. The
+    # swizzle is constant within each 16-byte group, so form v4 transfer units.
     for group in cl.static_iter(range(word_count // 4)):
         byte_offset = (
             slice_index * MMA_M * slice_bytes
@@ -421,8 +419,8 @@ def _store_output_global(
                     alignment=1,
                 )
     else:
-        # CUTLASS keeps the converted FP16/BF16 lanes scalar on the guarded
-        # tail path, which lowers to one b16 store per valid element.
+        # Keep converted FP16/BF16 lanes scalar on the guarded tail path, which
+        # lowers to one b16 store per valid element.
         dst_u16 = cl.bitcast(
             dst,
             cl.pointer_dtype(cl.uint16, cl.MemorySpace.GLOBAL),
@@ -489,12 +487,12 @@ def _key_trip_context(
     window_left,
     window_right,
 ):
-    """Return the source super-tile K interval and its Q0 invalid tail.
+    """Return the super-tile K interval and its Q0 invalid tail.
 
-    CUTLASS prunes whole 128-column K/V tiles using the union of the two
-    128-row Q halves.  Q0 can therefore have one fully-invalid trailing tile
-    that Q1 still needs.  That tile remains in the barrier protocol but does
-    not issue QK0/PV0 math.
+    Whole 128-column K/V tiles are pruned using the union of the two 128-row Q
+    halves. Q0 can therefore have one fully-invalid trailing tile that Q1 still
+    needs. That tile remains in the barrier protocol but does not issue QK0/PV0
+    math.
     """
     query_begin = seq_tile * CTA_M
     offset = cl.int32(0)
@@ -792,8 +790,8 @@ def _fmha_prefill_kernel(
     sched_valid = cl.shared_array(1, cl.int32, alignment=4)
 
     tid = cl.thread_index(0)
-    # Match CUTLASS's make_warp_uniform(warp_idx): broadcasting lane 0's
-    # value lets the backend treat every role predicate as warp-uniform.
+    # Broadcasting lane 0's warp index lets the backend treat every role
+    # predicate as warp-uniform.
     warp = cl.shfl_sync(tid // WARP_SIZE, 0)
     lane = tid % WARP_SIZE
     if variable_length:
@@ -855,9 +853,9 @@ def _fmha_prefill_kernel(
             l2_promotion=cl.TensorMapL2Promotion.NONE,
         )
 
-    # CUTLASS initializes Blackwell pipeline stages cooperatively with warp 0.
-    # Dynamic lane indexing keeps one static PTX site per pipeline half while
-    # still initializing every physical barrier exactly once.
+    # Initialize pipeline stages cooperatively with warp 0. Dynamic lane
+    # indexing keeps one static PTX site per pipeline half while initializing
+    # every physical barrier exactly once.
     if warp == SOFTMAX0_WARPS[0]:
         if lane < Q_STAGES:
             cl.mbarrier_initialize(q_full.get_element_pointer(lane), 1)
@@ -1313,8 +1311,8 @@ def _fmha_prefill_kernel(
         kv_index = 0
         kv_full_phase = 0
         while valid:
-            # CUTLASS reuses one elected MMA-warp lane for all UMMA math in a
-            # work tile. Pipeline commits elect independently in their helper.
+            # Reuse one elected MMA-warp lane for all UMMA math in a work tile.
+            # Pipeline commits elect independently in their helper.
             mma_elect_one = cl.elect_sync()
             query_base, key_base, task_seq_q, task_seq_k = _sequence_context(
                 cumulative_q,
@@ -1364,7 +1362,7 @@ def _fmha_prefill_kernel(
 
             # QK0(0), QK1(0): both halves share K0 before its ring slot is
             # released. The acquire mbarrier wait orders the TMA async writes
-            # before the tcgen MMA, matching CUTLASS's TMA-to-UMMA pipeline.
+            # before the tcgen MMA in the TMA-to-UMMA pipeline.
             k_stage = kv_index
             _wait_mbarrier(kv_full.get_element_pointer(k_stage), kv_full_phase)
             k_desc = _qk_descriptor(
@@ -1842,9 +1840,9 @@ def _fmha_prefill_kernel(
                             for chunk in cl.static_iter(range(4))
                         )
 
-                    # Match CUTLASS's four 32-value vector reductions:
-                    # MAXIMUMF propagates NaNs within each chunk, while the
-                    # outer maxnum combines retain the source's semantics.
+                    # Use four 32-value vector reductions. MAXIMUMF propagates
+                    # NaNs within each chunk, while the outer maxnum combines
+                    # retain the required semantics.
                     new_row_max = cl.float32(-float("inf"))
                     for chunk in cl.static_iter(score_chunks):
                         chunk_max = chunk[:32].reduce(
@@ -1878,18 +1876,16 @@ def _fmha_prefill_kernel(
                         stats_vec,
                     )
                     # Acquire the first P-ready stage while the statistics
-                    # store is in flight.  The second acquire and STTM wait
-                    # are split into the initial FMA lookahead below, matching
-                    # CUTLASS's machine schedule.
+                    # store is in flight. The second acquire and STTM wait are
+                    # split into the initial FMA lookahead below.
                     _wait_mbarrier(
                         p_empty.get_element_pointer((qid, 0)),
                         1 ^ (pv_event & 1),
                     )
 
                     block_sums = cl.Vector(0.0, 0.0, dtype=cl.float32)
                     minus_safe_max_scale = -safe_max * scale_softmax_log2
-                    # Keep the DKG constexpr-loop topology while using flattened
-                    # local arrays in place of its mutable Python lists.
+                    # Use constexpr loops with flattened local arrays.
                     if input_kind == ELEMENT_E4M3:
                         p_window_elems = 4
                         fma_pipe_windows = 3
@@ -3211,10 +3207,10 @@ def benchmark(
         return 0.0
 
     if use_cuda_graphs:
-        # Match the CUTLASS testing helper's semantics: capture the complete
-        # warmup and profiling loops on a non-default stream, replay each graph
-        # once, and divide the profiled graph time by its launch count.  Each
-        # workspace has static graph arguments, so cold-L2 rotation is retained.
+        # Capture the complete warmup and profiling loops on a non-default
+        # stream, replay each graph once, and divide the profiled graph time by
+        # its launch count. Each workspace has static graph arguments, so
+        # cold-L2 rotation is retained.
         workspace_index = 0
         if warmup_iterations:
             warmup_graph = torch.cuda.CUDAGraph()
 
@@ -3,10 +3,9 @@
 
 """Shared helpers for the CUDA Lang Blackwell FMHA-prefill tutorial.
 
-This module mirrors DKG's single fmha_prefill_helpers.py dependency. It
-contains the pure-Python mask/work-tile formulas, specialization traits,
-CUDA Lang leaf device helpers, and host-side tensor/reference utilities used
-by fmha_prefill.py and its focused tests.
+This module contains the pure-Python mask/work-tile formulas, specialization
+traits, CUDA Lang leaf device helpers, and host-side tensor/reference utilities
+used by fmha_prefill.py and its focused tests.
 """
 
 from __future__ import annotations
@@ -1218,9 +1217,9 @@ def _validate_output_traits(traits: OutputTraits) -> None:
 # Leaf CUDA Lang device helpers.
 # ---------------------------------------------------------------------------
 
-# CUTLASS Swizzle<B, 4, 3> uses B={1,2,3} for 32/64/128-byte TMA
-# swizzles.  The XOR source bits start at byte-address bit 7 and the
-# destination bits start at bit 4.
+# Swizzle<B, 4, 3> uses B={1,2,3} for 32/64/128-byte TMA swizzles. The XOR
+# source bits start at byte-address bit 7 and the destination bits start at
+# bit 4.
 SWIZZLE_32B_BITS = 1
 SWIZZLE_64B_BITS = 2
 SWIZZLE_128B_BITS = 3
@@ -1312,7 +1311,7 @@ def unpack_e4m3_lane(packed, lane):
 
 @cl.function
 def swizzle_byte_offset(byte_offset, swizzle_bits):
-    """Apply CUTLASS ``Swizzle<B,4,3>`` to a slice-local byte offset.
+    """Apply ``Swizzle<B,4,3>`` to a slice-local byte offset.
 
     The operation is its own inverse because it XORs bits 4..6 using bits
     7..9, which are not modified.  ``swizzle_bits`` must be 1, 2, or 3.
@@ -1906,8 +1905,8 @@ def prepare_tensors(
     """Allocate fixed or flattened-varlen tensors on ``device``.
 
     Custom Q/K/V inputs are interpreted in storage order and must be supplied
-    together.  Reference inputs retain the original FP32 source values, matching
-    the CUTLASS example's ``create_and_pad_tensor`` contract.
+    together. Reference inputs retain the original FP32 values used for the
+    correctness calculation.
     """
 
     torch_module = _require_torch()
 
@@ -1,57 +1,15 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
 # SPDX-License-Identifier: Apache-2.0
 
-"""CUDA Lang port of the CuTe DSL ``fp16_gemm_0.py`` tutorial.
-
-Computes ``C = A @ B.T`` with an optional FP16 row bias. The kernel retains
-the source tutorial's 128x128x64 single-CTA baseline: one shared-memory stage,
-explicit TMA and mbarrier operations, tcgen05 MMA into TMEM, FP32 accumulation,
-and an FP16 epilogue.
-
-CuTe DSL -> CUDA Lang API mapping used here:
-
-* ``@cute.kernel`` maps to ``@cl.kernel``. CuTe's ``warp_idx``, thread/block
-  indices, ``ceil_div``, and ``range_constexpr`` map to the corresponding
-  ``cl`` index operations, ``cl.cdiv``, and ``cl.static_iter``.
-* ``cutlass.Array(..., space=cutlass.AddressSpace.smem)`` maps to
-  ``cl.shared_array``. The A/B arrays are flattened in CUDA Lang, but retain
-  the source sizes, 128-byte alignment, and physical shared-memory layout.
-* ``cuda.create_tensor_map_tiled_from_view`` maps to ``cl.tensor_map_tiled``.
-  A and B remain row-major, K-contiguous 2-D tensors. CUDA Lang expresses the
-  TMA coordinate order explicitly as ``(K, M)`` / ``(K, N)`` using tile shapes
-  ``(BLOCK_K, BLOCK_M)`` / ``(BLOCK_K, BLOCK_N)`` and ``order="F"``; this is
-  equivalent to CuTe DSL automatically selecting K as the leading mode.
-* ``prims.cp_async_bulk_tensor_shared_cta_global`` maps to
-  ``cl.copy_async_bulk_tensor_global_to_shared``. The expected transaction
-  byte count and the TMA-completion mbarrier are unchanged.
-* ``prims.mbarrier_init``, ``mbarrier_arrive_expect_tx``, and
-  ``mbarrier_try_wait_parity`` map to ``cl.mbarrier_initialize``,
-  ``cl.mbarrier_arrive_expect_transaction``, and
-  ``cl.mbarrier_wait_parity``. The full/empty/accumulator barrier topology
-  and phase progression are preserved.
-* ``prims.Tcgen05InstrDesc.build`` and ``Tcgen05SmemDesc.build`` map to
-  ``cl.Tcgen05InstructionDescriptor`` and
-  ``cl.Tcgen05SharedMemoryDescriptor``, each followed by ``encode()``. Both
-  builders take byte-valued leading/stride offsets. Once encoded, the per-MMA
-  K-step is added in descriptor units of 16 bytes, hence ``2 * kk`` for each
-  32-byte F16 MMA K-step.
-* ``prims.tcgen05_alloc``, ``tcgen05_mma``, ``tcgen05_commit``,
-  ``tcgen05_ld``, and ``tcgen05_dealloc`` map directly to the corresponding
-  ``cl.tcgen05_*`` operations. Allocation size, MMA accumulation predicate,
-  commit barriers, 32-column TMEM loads, and deallocation size match the
-  source.
-* The CuTe ``@cute.jit`` host function builds tensor maps and launches the
-  device kernel. CUDA Lang instead declares tensor maps in ``@cl.kernel``;
-  compilation hoists their construction into the generated host program, so
-  both public launch interfaces still accept ordinary 2-D tensors.
-* ``cute.compile`` has no explicit CUDA Lang counterpart here. ``run()`` calls
-  ``cl.launch`` directly; the first launch compiles the kernel and later
-  launches reuse CUDA Lang's internal cache. K and ``has_bias`` are
-  ``cl.Constant`` specialization keys, while M and N remain runtime dimensions
-  of C.
-
-The CLI retains the source tutorial's correctness-checking interface. CUDA
-Lang-specific restrictions are validated explicitly.
+"""Single-CTA CUDA Lang FP16 GEMM tutorial.
+
+Computes ``C = A @ B.T`` with an optional FP16 row bias using a 128x128x64
+tile, one shared-memory stage, explicit TMA and mbarrier operations, tcgen05
+MMA into TMEM, FP32 accumulation, and an FP16 epilogue. A and B are row-major,
+K-contiguous tensors; ``order="F"`` expresses K as the leading TMA coordinate.
+Tensor maps declared in ``@cl.kernel`` are hoisted into the generated host
+program. K and ``has_bias`` are specialization keys, while M and N remain
+runtime dimensions of C. The CLI validates inputs and checks correctness.
 """
 
 from __future__ import annotations
@@ -100,8 +58,7 @@ def _kernel(
 
     # A/B are ordinary row-major (rows, K) arrays. ``order="F"`` reverses the
     # tensor-map modes to (K, rows), making K the contiguous TMA coordinate.
-    # This is the explicit CUDA Lang equivalent of CuTe DSL's leading-mode
-    # detection in create_tensor_map_tiled_from_view.
+    # The explicit mode order makes K the leading tensor-map dimension.
     a_tmap = cl.tensor_map_tiled(
         a,
         (BLOCK_K, BLOCK_M),
 
@@ -1,29 +1,17 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
 # SPDX-License-Identifier: Apache-2.0
 
-"""CUDA Lang port of the CuTe DSL ``fp16_gemm_1.py`` tutorial.
+"""Two-CTA clustered CUDA Lang FP16 GEMM tutorial.
 
 Computes ``C = A @ B.T`` with an optional FP16 row bias. Two CTAs collaborate
 on each 256x256 output tile: each CTA loads one 128-row A slice and one
 128-column B slice, the leader issues CTA_2 tcgen05 MMA instructions, and both
 CTAs store their 128 output rows.
 
-Additional CuTe DSL -> CUDA Lang notes beyond ``fp16_gemm_0.py``:
-
-* CuTe DSL differentiates CTA-local and cluster TMA with different primitive
-  names: ``cp_async_bulk_tensor_shared_cta_global`` in ``fp16_gemm_0.py`` and
-  ``cp_async_bulk_tensor_shared_cluster_global`` here. CUDA Lang uses the same
-  ``cl.copy_async_bulk_tensor_global_to_shared`` API for both; cluster TMA is
-  selected by passing a destination pointer in
-  ``cl.MemorySpace.SHARED_CLUSTER`` plus the CTA_2 multicast arguments.
-  Completion still has to arrive on the leader CTA's mbarrier, so this sample
-  maps the local ``ab_full`` pointer with ``cl.map_shared_to_leader_block``.
-* The source kernel takes K as a runtime argument. This sample keeps M, N, and K
-  runtime dimensions; only the optional-bias path is specialized.
-
-The implementation otherwise retains the source's cluster shape, work
-decomposition, data movement, barrier phases, CTA_2 MMA sequence, TMEM
-epilogue, and vector stores.
+Cluster TMA uses a destination pointer in ``cl.MemorySpace.SHARED_CLUSTER``
+with CTA_2 multicast arguments. Completion arrives on the leader CTA's
+mbarrier, mapped with ``cl.map_shared_to_leader_block``. M, N, and K remain
+runtime dimensions; only the optional-bias path is specialized.
 """
 
 from __future__ import annotations
 
@@ -1,23 +1,16 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES.
 # SPDX-License-Identifier: Apache-2.0
 
-"""CUDA Lang port of the CuTe DSL ``fp16_gemm_2.py`` tutorial.
+"""Two-CTA clustered CUDA Lang FP16 GEMM tutorial with FP32 output.
 
 Computes ``C = A @ B.T`` with an optional FP32 row bias. Two CTAs collaborate
 on each 256x128 output tile: each CTA loads one 128-row A slice and one
 64-column B slice, the leader issues CTA_2 tcgen05 MMA instructions, and both
 CTAs store their 128 output rows.
 
-Additional CuTe DSL -> CUDA Lang notes beyond ``fp16_gemm_1.py``:
-
-* The collective N tile is 128 instead of 256. CTA_2 still splits the tile
-  across the two CTA ranks, so each CTA loads a 64-column B slice.
-* C and the optional row bias are FP32. The epilogue stores FP32 accumulator
-  vectors directly instead of converting them to FP16.
-
-The implementation otherwise retains the source's cluster shape, work
-decomposition, data movement, barrier phases, CTA_2 MMA sequence, TMEM
-epilogue, and vector stores.
+The collective N tile is 128 and is split across the two CTA ranks, so each CTA
+loads a 64-column B slice. C and the optional row bias are FP32; the epilogue
+stores accumulator vectors directly without an FP16 conversion.
 """
 
 from __future__ import annotations