tests/unittest/_torch/thop/parallel/test_arcquant_fp4.py (1)

9-40: Add return type annotations to the new helpers.

These utilities are missing explicit return types, which is out of step with the repo's Python typing rules.

✍️ Suggested cleanup

-def _benchmark(fn, warmup=10, iters=100):
+def _benchmark(fn, warmup: int = 10, iters: int = 100) -> float:
@@
-def _effective_bytes(M, K, KE, input_dtype=torch.bfloat16):
+def _effective_bytes(M: int, K: int, KE: int, input_dtype: torch.dtype = torch.bfloat16) -> int:
@@
-def _gb_per_sec(total_bytes, time_ms):
+def _gb_per_sec(total_bytes: int, time_ms: float) -> float:

As per coding guidelines, "Always annotate Python function return types; use None if the function does not return anything".

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tests/unittest/_torch/thop/parallel/test_arcquant_fp4.py` around lines 9 -
40, Add explicit return type annotations to the three helper functions: annotate
_benchmark(...) -> float, _effective_bytes(M, K, KE, input_dtype: torch.dtype =
torch.bfloat16) -> int, and _gb_per_sec(total_bytes, time_ms) -> float; update
the function signatures (not bodies) to include these types so they comply with
the project's typing rules.

cpp/tensorrt_llm/kernels/arcquantFP4.cu (1)

463-503: Label the get_sf_out_offset_128x4 arguments inline.

These call sites are hard to audit with five positional arguments and two std::optional values. Inline parameter names would make the swizzled-layout indexing much easier to review.

🧭 Suggested cleanup

-                auto sfOffset
-                    = get_sf_out_offset_128x4(std::nullopt, rowIdx, sfIdx, std::optional<int>(numRows), numSFCols);
+                auto sfOffset = get_sf_out_offset_128x4(
+                    /*batchIdx=*/std::nullopt, /*mIdx=*/rowIdx, /*kIdx=*/sfIdx,
+                    /*numRows=*/std::optional<int>(numRows), /*numColVecs=*/numSFCols);
@@
-                auto sfOffset
-                    = get_sf_out_offset_128x4(std::nullopt, rowIdx, pos, std::optional<int>(numRows), numSFCols);
+                auto sfOffset = get_sf_out_offset_128x4(
+                    /*batchIdx=*/std::nullopt, /*mIdx=*/rowIdx, /*kIdx=*/pos,
+                    /*numRows=*/std::optional<int>(numRows), /*numColVecs=*/numSFCols);
@@
-                    auto sfOffsetRes = get_sf_out_offset_128x4(
-                        std::nullopt, rowIdx, pos + 1, std::optional<int>(numRows), numSFCols);
+                    auto sfOffsetRes = get_sf_out_offset_128x4(
+                        /*batchIdx=*/std::nullopt, /*mIdx=*/rowIdx, /*kIdx=*/pos + 1,
+                        /*numRows=*/std::optional<int>(numRows), /*numColVecs=*/numSFCols);

As per coding guidelines, "In C++ function calls with non-obvious parameters, use inline C comments with the format /*paramName=*/ to document parameters".

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@cpp/tensorrt_llm/kernels/arcquantFP4.cu` around lines 463 - 503, Call sites
of get_sf_out_offset_128x4 are hard to read because they pass multiple
positional args and std::optional values; update each call (e.g., the ones using
std::nullopt, rowIdx, pos, std::optional<int>(numRows), numSFCols and the pos+1
variant) to use inline parameter-name comments like /*argName=*/ before each
argument (for example /*swizzle=*/std::nullopt, /*row=*/rowIdx, /*col=*/pos,
/*numRows=*/std::optional<int>(numRows), /*numSFCols=*/numSFCols) so the meaning
of each parameter is explicit and easier to audit.

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes