Fixes from PR comments, unit tests

Signed-off-by: Antoni Viros i Martin <[email protected]>

Author: ani300

fms_mo/aiu_addons/fp8/fp8_linear.py

@@ -31,7 +31,7 @@
 from fms_mo.prep import available_packages
 
 # pylint: disable=not-callable
-# torch.nn.functional.scaled_dot_product_attention not recognized as callable
+# torch.nn.functional.linear not recognized as callable
 # open issue in PyLint: https://github.com/pytorch/pytorch/issues/119482
 
 # Gated torchao imports for FP8 implementation

fms_mo/aiu_addons/fp8/fp8_spyre_op.py

@@ -99,6 +99,7 @@ def scaled_paged_attn_store(
     Scales key and value tensors, and stores them to the paged KV cache
     using the same schema as vLLM.
     """
+    print("Should never hit")
     result_key_cache = key_cache.clone()
     result_value_cache = value_cache.clone()
     for seq_i, slot_mapping_seq in enumerate(slot_mapping):
@@ -150,6 +151,7 @@ def scaled_paged_attn_compute(
     Implements a CPU fallback to run the kernel that has been confirmed
     to match the vLLM fused kernel.
     """
+    print("Should never hit")
     # torch.zeros(NUM_BLOCKS, BLOCK_SIZE, kvheads, head_size, dtype=model_dtype),
     output = torch.zeros_like(query)
     num_query_heads = query.shape[2]

pyproject.toml

@@ -30,8 +30,7 @@ dependencies = [
 "datasets>=3.0.0,<4.0",
 "pandas",
 "safetensors",
-"pkginfo>1.10",
-"torchao"
+"pkginfo>1.10"
 ]
 
 [project.optional-dependencies]

tox.ini

@@ -33,6 +33,8 @@ deps =
     pytest
     pylint>=2.16.2,<4.0
     pylint-pydantic
+    ibm-fms
+    torchao
 commands =
     {basepython} -m pylint --load-plugins pylint_pydantic fms_mo/ tests/