Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@examples/gemm_fp4/example_fusedmoe_a8w4_sm100.py`:
- Around line 40-46: Add an explicit guard that TL_MOE_HIDDEN is even before
treating expert weight tensors as packed FP4 (two 4-bit values per byte);
validate this in the config/parsing path and before any use of
unpack_fp4_to_float and places where expert tensors are created/reshaped
(references: function unpack_fp4_to_float and the code regions handling expert
weight shapes around the other occurrences). If TL_MOE_HIDDEN is odd, raise a
clear ValueError (or argparse/config error) explaining that packed FP4 storage
requires an even hidden dimension, so the failure happens early instead of
during unpacking.

In `@examples/gemm_fp4/example_fusedmoe_a8w4_sm120.py`:
- Around line 136-141: The zero-input smoke test currently only prints PASS/FAIL
and can allow the script to exit successfully on failure; change it to a hard
assertion so failures stop execution: after calling jit_kernel(z_input, z_gate,
z_up) and computing c_zero, replace the print line with an assertion that
c_zero.abs().max().item() == 0.0 (or torch.equal(c_zero,
torch.zeros_like(c_zero))) and include a descriptive message (e.g., "zeros in ->
zeros out failed") so the test fails closed; refer to the variables z_input,
z_gate, z_up and the function jit_kernel to locate where to apply this change.

In `@examples/gemm_fp4/example_gemm_a8w4_sm100.py`:
- Around line 35-41: The code assumes K is even when packing/unpacking FP4 (see
unpack_fp4_to_float) which silently truncates when TL_A8W4_K is odd; add an
explicit check (raise ValueError or assert) that TL_A8W4_K % 2 == 0 before any
packing or buffer allocation/reshape, and apply the same guard near the FP4
packing logic and any other unpacking uses (the unpack_fp4_to_float function and
the FP4 weight-packing sites) so the script fails fast with a clear error
instead of allocating a truncated buffer.

In `@examples/gemm_fp4/example_gemm_fp4_sm100.py`:
- Around line 131-137: The unpack_fp4_to_float routine (and corresponding calls
like make_random_fp4) assume the packed axis length K//2 (or cols//2) is an
integer; add an upfront validation in unpack_fp4_to_float (and the code that
calls it, e.g., make_random_fp4 and any uses conditioned on
TL_FP4_TRANSPOSE_B/TL_FP4_K/TL_FP4_N) to check that K (and any cols passed as
packed length*2) is even and raise a clear ValueError if not (or alternatively
document/pad explicitly). Locate the unpack_fp4_to_float function and the
make_random_fp4 call sites referenced in the diff and enforce this guard so odd
TL_FP4_K or TL_FP4_N cannot silently truncate the buffer.

In `@examples/gemm_fp4/example_gemm_nvfp4_sm120.py`:
- Around line 71-109: The kernel assumes full 16x8x64 tiles but the launch dims
use ceildiv, so for any M, N, or K not divisible by block_M=16, block_N=8,
block_K=64 the kernel will read/write out-of-bounds; fix by rejecting
unsupported sizes early in matmul_nvfp4_sm120 (before building/returning main) —
check that M % block_M == 0, N % block_N == 0, and K % block_K == 0 and raise a
clear exception (e.g., ValueError) if not, so callers must provide padded inputs
or a different kernel; reference block_M/block_N/block_K and the prim func main
where copies (T.copy to/from A_shared/B_shared and final T.copy to C) currently
assume full tiles.
- Around line 112-124: The example lacks a runtime SM120 capability check before
calling tilelang.compile for the SM120-only frontend (T.nvfp4_gemm) in the
__main__ path; add a preflight that queries the current device capability (via
tilelang.utils.target.target_is_sm120 or equivalent helper) and fail fast with a
clear "SM120 required" message if the device is not SM120, returning/non-zero
exit instead of proceeding to tilelang.compile and kernel creation; update both
example_gemm_nvfp4_sm120 entrypoint logic and any similar SM120-only example
functions to guard before compilation and import/use of SM120-specific
primitives.

In `@src/cuda/codegen/codegen_cuda.cc`:
- Around line 4746-4754: The scope check using GetPtrStorageScope(buffer_var) is
incomplete—GetBufferRef() first consults alloc_storage_scope_ so allocated
shared/local FP4 buffers can be misclassified as global/packed; update the FP4
packed-load/store branches to resolve scope the same way GetBufferRef() does by
checking alloc_storage_scope_ for buffer_var.get() first and falling back to
GetPtrStorageScope(buffer_var), then use that resolved scope when computing
is_packed_fp4_scope; apply this change in both the load (tl_fp4_packed_load) and
store branches.

In `@src/cuda/op/copy.cc`:
- Around line 1721-1728: The code path handling is_align16b_subbyte_layout
currently calls TMABytesFromElements and so ignores transaction-byte accounting
(and the float4_e2m1_unpacked FP4 special case), causing mbarrier_expect_tx to
be overstated; modify the branch where total_bytes is set (the block that uses
inner_box_dim, instruction_dim, total_elements, and shared_tensor->dtype) to
call TMATransactionBytesFromElements(total_elements * loop_extent,
shared_tensor->dtype) instead of TMABytesFromElements so transaction-byte
accounting (and FP4 handling) is used for ALIGN16B loads.

In `@src/layout/gemm_layouts.cc`:
- Around line 1027-1035: DetectSwizzleMode() is collapsing ALIGN16B into
SwizzleMode::kFull causing MergeSwizzleLayouts() to reconstruct the wrong layout
with makeFullBankSwizzleLayout(buffer); add a distinct swizzle mode (e.g.,
SwizzleMode::kAlign16B) or, when merging, branch on info.element_size < 8 to
detect ALIGN16B and reconstruct it with makeAlign16BSwizzleLayout(buffer)
instead of makeFullBankSwizzleLayout(buffer); update the SwizzleMode enum and
the switch/merge logic in MergeSwizzleLayouts() to return/use the new kAlign16B
(or the element_size-based branch) so ALIGN16B mappings for sub-byte types
remain preserved.

In `@src/tl_templates/cuda/instruction/mma.h`:
- Around line 361-381: The FP4 values loaded by ptx_ldmatrix_b4x16_* end up in
the low nibble and must be repacked the same way as in mma_sync; update
mma_sync_blockscaled to apply the FP4 repack to the A and B operands before
calling Dispatcher::exec (i.e., perform the same left-shift/repack used in
mma_sync on the arrays of Dispatcher::ARegType and Dispatcher::BRegType prior to
dispatch) so BlockScaledMmaDispatcher::exec receives the correctly encoded FP4
data and produces correct accumulations.

In `@tilelang/cuda/intrinsics/macro/mma_macro_generator.py`:
- Around line 678-709: The scale operands are never advanced per k_inner, so
SFA_local_buf and SFB_local_buf always use base 0; compute SFA_offset and
SFB_offset analogous to A_offset/B_offset using k_inner (e.g. advancing by
k_inner * element_stride_or_vector_size used for scale storage) and pass those
offsets into T.ptx_mma_blockscaled instead of always 0; update the
_atom_mma_blockscaled macro to use the new SFA_offset and SFB_offset
(referencing k_inner, SFA_local_buf, SFB_local_buf, and the
_atom_mma_blockscaled/T.ptx_mma_blockscaled call) so each ki iteration uses the
correct per-k_inner scale vectors.
- Around line 126-136: The check for FP4 chunk sizes must reject non-multiples
rather than only enforcing a minimum: in the block-scaled path (when
self.is_blockscaled and str(a_dtype) == "float4_e2m1fn") validate that
self.chunk is a multiple of 64 (e.g. self.chunk % 64 == 0) and raise a
ValueError if not; in the plain FP4 path (when str(a_dtype) == "float4_e2m1fn")
validate that self.chunk is a multiple of 32 (self.chunk % 32 == 0) instead of
only checking self.chunk < 32, and raise a clear ValueError if the modulus check
fails; keep setting self.k_dim = 64 (block-scaled) or self.k_dim = min(32,
self.chunk) after the validation.

In `@tilelang/language/gemm_op.py`:
- Around line 269-320: The nvfp4_gemm path dropped the _gemm_impl base-offset
and rank guards so row-sliced packed tiles can lose their outer-dimension base
offset and malformed inputs raise IndexError; restore the same validation: after
computing A_offset = retrieve_offset(A_region) and B_offset =
retrieve_offset(B_region) (and their shapes via retrieve_shape if needed),
assert that A_offset and B_offset have the expected rank (e.g., len(A_offset) >=
2 and len(B_offset) >= 2) and assert A_offset[-2] == 0 and B_offset[-2] == 0 (or
alternatively call _gemm_impl to reuse its checks) before forwarding
A_offset[-1] and B_offset[-1] into tirx.call_intrin in nvfp4_gemm so
outer-dimension base offsets cannot be silently lost.

In `@tilelang/layout/swizzle.py`:
- Around line 91-120: This code currently flattens all leading dims into a
single stride and calls _ffi_api.make_tcgen05mma_swizzled_layout which lets the
swizzle cross outer-dimension boundaries; instead, compute the 2D layout from
the trailing two dims only (use shape[-2:] and the corresponding last-two
stride/continuity values from _get_stride_continuous/stride), call
_ffi_api.make_tcgen05mma_swizzled_layout to build the base 2D layout for those
last two dims, then expand that 2D base across the leading dimensions
(shape[:-2]) using the same ExpandLayout2D-style expansion used by the native
C++ factories so the Python path matches the buffer-based path for rank>2
buffers.

---

Outside diff comments:
In `@tilelang/cuda/intrinsics/macro/tcgen05_macro_generator.py`:
- Around line 399-418: runtime_instr_desc is computed once from sf_a_id/sf_b_id
before the ki loop, so all atoms use the same scale IDs even though
descriptors/offsets advance; move the computation of runtime_instr_desc (or
otherwise recompute sf_a_id/sf_b_id) inside the ki loop so each atom uses the
current scale IDs for its descriptor, i.e. update how runtime_instr_desc is
built before calling tcgen05_blockscaled_atom (referencing runtime_instr_desc,
sf_a_id, sf_b_id, tcgen05_blockscaled_atom, sfa_data, sfb_data, a_params,
b_params) so multi-atom iterations use the correct per-atom scale vectors.

---

Nitpick comments:
In `@examples/gemm_fp4/example_fusedmoe_nvfp4_sm120.py`:
- Around line 118-132: The X activation tile is copied twice per ko; change the
loops so X_bytes is loaded into X_shared once and reused for both GEMMs: for
each ko, do a single T.copy(X_bytes[by * block_tokens, ko *
packed_block_hidden], X_shared) then call T.nvfp4_gemm for gate (using
gate_shared, SFA_local/SFB_local, gate_local, clear_accum=(ko == 0)) and then
call T.nvfp4_gemm for up (using up_shared, SFA_local/SFB_local, up_local,
clear_accum=(ko == 0)), removing the duplicate T.copy from the second loop
(references: X_bytes, X_shared, W_gate_bytes, W_up_bytes, T.nvfp4_gemm,
gate_local, up_local).

In `@tilelang/cuda/intrinsics/layout/utils.py`:
- Around line 22-29: The function get_ldmatrix_offset currently advertises
dtype: Literal["float16", "int8", "int4"] but the implementation contains a
runtime path for FP4; update the public type contract to include "fp4" (or "FP4"
matching project convention) so type-checkers accept the code, and adjust any
surrounding docstrings/error messages to explicitly list and distinguish "fp4"
and "int4" (and/or their canonical casing) where dtype choices are described or
validated (also update the equivalent signature/annotations and messages in the
similar helper at lines ~52-63 that handle small-int/fp4 cases). Ensure you
reference and change the dtype literal in get_ldmatrix_offset and the
corresponding validation/error text so FP4 and INT4 remain clearly separated.

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@src/cuda/codegen/codegen_cuda.cc`:
- Around line 3411-3433: When use_mxf4nvf4 is true, validate that
kind_dtype/dtype_enum is the FP4 variant expected by the mxf4nvf4 intrinsic
before emitting tcgen05mma_mxf4nvf4_blockscaled_ss: call
tl::codegen::ptx::DTypeFromString(kind_dtype) (already present as dtype_enum)
and assert (ICHECK) that it matches the FP4 dtype enum the intrinsic requires,
otherwise fail with a clear error message (or avoid emitting the FP4 intrinsic);
update the branch that constructs tcgen05_call (and/or use_mxf4nvf4 decision) to
perform this dtype check so a mismatched caller errors at codegen rather than
silently lowering to the wrong intrinsic.

In `@src/op/tcgen5_meta.h`:
- Around line 358-359: The helper GetTCGEN5MXF4NVF4BlockScaledInstrDesc
currently writes a_sf_id into desc via set_bits(static_cast<uint32_t>(a_sf_id),
29, 2) even though a_sf_id is not a parameter and the field is fixed to zero for
scale_vec::4X; remove that stale write (or alternatively add an a_sf_id
parameter to the function signature and propagate it through if intended) so
that desc no longer attempts to set bits 29-30 from an undefined a_sf_id; update
any related comments to reflect the removal.

In `@tilelang/cuda/intrinsics/macro/tcgen05_macro_generator.py`:
- Around line 468-529: The guard and descriptor for B are incorrect: replace the
assert that assumes B is non-transposed and ensure b_transposed is used
consistently when building b_params; specifically, update the top assertion to
require the correct B layout (use self.b_transposed as done by
tcgen05mma_blockscaled() and compute_tcgen05_b_desc_params()), compute
b_is_k_major = self.b_transposed and pass that into the TCGEN05DescriptorParams
for b_params (instead of hardcoding is_k_major=True), and adjust any related
leading/stride byte offset logic to follow the K-major vs N-major branch
consistent with the _determinate_swizzle_mode(B_buf, ...) handling.
- Around line 531-556: The code reads SFA_tmem.dtype before normalizing
BufferRegion/Buffer, which fails for sliced regions; update the function to
first normalize SFA_tmem and SFB_tmem into their underlying Buffer or
Buffer.data (using the existing BufferRegion/Buffer checks for SFA_tmem and
SFB_tmem) and only then read dtype to compute sf_dtype and is_mxfp4 and call
_ffi_api.get_tcgen5_mxf4nvf4_blockscaled_instr_desc; reference symbols:
SFA_tmem, SFB_tmem, BufferRegion, Buffer, sf_dtype, is_mxfp4, and
_ffi_api.get_tcgen5_mxf4nvf4_blockscaled_instr_desc to locate where to move the
dtype access and normalization logic.

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@examples/gemm_fp4/example_gemm_nvfp4_sm100.py`:
- Around line 84-86: SFB_tmem is allocated with the wrong tile dimension causing
incorrect sizing for non-square tiles; change the alloc_tmem call that creates
SFB_tmem (currently T.alloc_tmem([block_M, 4], "uint32")) to use block_N instead
of block_M so it becomes T.alloc_tmem([block_N, 4], "uint32"), ensuring SFB_tmem
matches per-row scales of B; update the allocation site where SFB_tmem is
defined and verify any uses of SFB_tmem assume the N-dimension.

---

Nitpick comments:
In `@tilelang/cuda/op/gemm/gemm_tcgen05.py`:
- Around line 286-341: Extract the duplicated is_nvfp4 dispatch into a small
helper function or lambda and call it from both _gemm_blockscaled_cond and
_gemm_blockscaled to avoid repeating the MMA call pairs; specifically create a
helper (e.g., _call_mma_blockscaled or a local lambda) that takes (A_shared,
B_shared, C_local, SFA_tmem, SFB_tmem, mbarptr, clear_accum, sf_a_id=None,
sf_b_id=None) and inside performs the is_nvfp4 check to call
mma_emitter.tcgen05mma_mxf4nvf4_blockscaled or
mma_emitter.tcgen05mma_blockscaled accordingly, then replace the duplicated
blocks in _gemm_blockscaled_cond and _gemm_blockscaled with a single call to
that helper.

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@src/layout/gemm_layouts.cc`:
- Around line 543-556: The current ICHECK at line 543 only validates that the
FP4 count is even, but this is insufficient because the downstream code uses
16-byte packed vectors even when bank equals 1. Add an additional validation
that byte_continuous is divisible by vector_size (16 bytes) to reject dense FP4
shapes that do not cover a full 16-byte vector. This should be enforced in the
same ICHECK or a new one immediately following the existing check, ensuring that
continuous is at least 32 FP4 elements (since byte_continuous = continuous / 2,
requiring byte_continuous >= 16 means continuous >= 32).

In `@tilelang/layout/swizzle.py`:
- Around line 177-189: The function `make_dense_fp4_swizzled_layout` calls
`_ffi_api.make_dense_fp4_swizzled_layout` directly without checking if the FFI
symbol exists, unlike the similar `make_align16b_swizzled_layout` function which
guards the call with hasattr. Add a hasattr check on _ffi_api for the
make_dense_fp4_swizzled_layout symbol before attempting to call it, and either
provide a graceful fallback or raise a descriptive error message if the symbol
is not available. This ensures users receive a clear, helpful error instead of
an AttributeError when the native library lacks this symbol.

---

Outside diff comments:
In `@src/layout/gemm_layouts.cc`:
- Around line 998-1011: The TCGEN05 path in makeGemmABLayoutSm100 at lines
998-1011 needs to distinguish between dense FP4 and ALIGN16B layouts for
sub-byte operands, but currently it treats them uniformly. Additionally, the
logic at lines 1076-1084 and 1093-1100 returns kNone for dense FP4 layouts
without proper differentiation, and MergeSwizzleLayouts at lines 1149-1156 lacks
a merge case for the dense FP4 SwizzleMode. To fix this, introduce an explicit
layout kind parameter that distinguishes dense FP4 from ALIGN16B layouts and
thread it through the SM100 selection logic at the anchor location, update the
layout kind determination at lines 1093-1100 to return the appropriate dense FP4
kind instead of kNone, update the TCGEN05 routing at lines 1076-1084 to handle
the dense case, and add a corresponding dense FP4 SwizzleMode merge case in
MergeSwizzleLayouts at lines 1149-1156 to preserve and reconstruct the dense
layout.

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@examples/gemm_fp4/example_gemm_fp4_sm100.py`:
- Around line 296-305: Add validation to ensure the K dimension is even in both
FP4 example files. In examples/gemm_fp4/example_gemm_fp4_sm100.py (lines
296-305), after the lines that resolve args.m, args.n, and args.k from
command-line arguments and environment variables, add a check that raises
ValueError with message "K must be even for packed FP4 inputs" if args.k is odd.
Apply the identical validation check in
examples/gemm_fp4/example_gemm_fp4_sm120.py (lines 232-240) after the args
resolution section to prevent silent truncation of packed FP4 buffers when K
values are not divisible by 2.

---

Duplicate comments:
In `@examples/gemm_fp4/example_gemm_fp4_sm100.py`:
- Around line 174-208: The K dimension must be validated to be even before use
in FP4 mode, since make_random_fp4 and unpack_fp4_to_float both use K // 2 for
packed buffers. Add a validation check after extracting K from args in the
run_fp4 function to ensure K is even; if odd, raise an error or print a clear
message explaining the requirement. Apply the same validation to any other entry
points that accept K as a parameter for FP4 operations.

---

Nitpick comments:
In `@examples/gemm_fp4/example_gemm_fp4_sm100.py`:
- Around line 165-171: The _arch_and_flags function uses a hardcoded default of
aarch64-linux for cuda_target, which fails on x86_64 systems that need
x86_64-linux instead. Auto-detect the host platform to determine the correct
default cuda_target based on the system's architecture (x86_64 vs aarch64). This
detection should be applied at the location where the cuda_target argument
default is defined (around line 292 where the argument parser is likely
configured) and ensure the detected value flows through to where it is used in
the _arch_and_flags function at line 165-171 when constructing the cuda_include
path. The auto-detection should happen early so that args.cuda_target receives
the platform-appropriate default unless explicitly overridden by the user.

In `@examples/gemm_fp4/example_gemm_fp4_sm120.py`:
- Around line 105-112: The accumulator C_local is being cleared twice
redundantly: once explicitly with T.clear(C_local) before the loop and once
implicitly during the first iteration of the loop via the clear_accum=(ko == 0)
parameter in the T.nvfp4_gemm call. Remove the explicit T.clear(C_local) call
since the T.nvfp4_gemm operation already handles accumulator initialization on
the first K-tile iteration through its clear_accum parameter.